diff --git a/arro3-compute/python/arro3/compute/_compute.pyi b/arro3-compute/python/arro3/compute/_compute.pyi index 6b2c39c..9085bde 100644 --- a/arro3-compute/python/arro3/compute/_compute.pyi +++ b/arro3-compute/python/arro3/compute/_compute.pyi @@ -1,4 +1,4 @@ -from typing import Sequence, overload +from typing import overload # Note: importing with # `from arro3.core import Array` @@ -54,109 +54,6 @@ def dictionary_encode( The dictionary-encoded array. """ -@overload -def dictionary_dictionary(array: types.ArrowArrayExportable) -> core.Array: ... -@overload -def dictionary_dictionary(array: types.ArrowStreamExportable) -> core.ArrayReader: ... -def dictionary_dictionary( - array: types.ArrowArrayExportable | types.ArrowStreamExportable, -) -> core.Array | core.ArrayReader: - """ - Access the `dictionary` of a dictionary array. - - This is equivalent to the [`.dictionary`][pyarrow.DictionaryArray.dictionary] - attribute on a PyArrow [DictionaryArray][pyarrow.DictionaryArray]. - - Args: - array: Argument to compute function. - - Returns: - The keys of a dictionary-encoded array. - """ - -@overload -def dictionary_indices(array: types.ArrowArrayExportable) -> core.Array: ... -@overload -def dictionary_indices(array: types.ArrowStreamExportable) -> core.ArrayReader: ... -def dictionary_indices( - array: types.ArrowArrayExportable | types.ArrowStreamExportable, -) -> core.Array | core.ArrayReader: - """ - Access the indices of a dictionary array. - - This is equivalent to the [`.indices`][pyarrow.DictionaryArray.indices] - attribute on a PyArrow [DictionaryArray][pyarrow.DictionaryArray]. - - Args: - array: Argument to compute function. - - Returns: - The indices of a dictionary-encoded array. - """ - -@overload -def list_flatten(input: types.ArrowArrayExportable) -> core.Array: ... -@overload -def list_flatten(input: types.ArrowStreamExportable) -> core.ArrayReader: ... -def list_flatten( - input: types.ArrowArrayExportable | types.ArrowStreamExportable, -) -> core.Array | core.ArrayReader: - """Unnest this ListArray, LargeListArray or FixedSizeListArray. - - Args: - input: Input data. - - Raises: - Exception if not a list-typed array. - - Returns: - The flattened Arrow data. - """ - -@overload -def list_offsets( - input: types.ArrowArrayExportable, *, logical: bool = True -) -> core.Array: ... -@overload -def list_offsets( - input: types.ArrowStreamExportable, *, logical: bool = True -) -> core.ArrayReader: ... -def list_offsets( - input: types.ArrowArrayExportable | types.ArrowStreamExportable, - *, - logical: bool = True, -) -> core.Array | core.ArrayReader: - """Access the offsets of this ListArray or LargeListArray - - Args: - input: _description_ - physical: If False, return the physical (unsliced) offsets of the provided list array. If True, adjust the list offsets for the current array slicing. Defaults to `True`. - - Raises: - Exception if not a list-typed array. - - Returns: - _description_ - """ - -def struct_field( - values: types.ArrowArrayExportable, - /, - indices: int | Sequence[int], -) -> core.Array: - """Access a column within a StructArray by index - - Args: - values: Argument to compute function. - indices: List of indices for chained field lookup, for example [4, 1] will look up the second nested field in the fifth outer field. - - Raises: - Exception if not a struct-typed array. - - Returns: - _description_ - """ - def take( values: types.ArrowArrayExportable, indices: types.ArrowArrayExportable ) -> core.Array: diff --git a/arro3-compute/src/lib.rs b/arro3-compute/src/lib.rs index 0e447ef..938630c 100644 --- a/arro3-compute/src/lib.rs +++ b/arro3-compute/src/lib.rs @@ -3,9 +3,6 @@ use pyo3::prelude::*; mod cast; mod concat; mod dictionary; -mod list_flatten; -mod list_offsets; -mod struct_field; mod take; const VERSION: &str = env!("CARGO_PKG_VERSION"); @@ -24,9 +21,6 @@ fn _compute(_py: Python, m: &Bound) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(dictionary::dictionary_dictionary))?; m.add_wrapped(wrap_pyfunction!(dictionary::dictionary_encode))?; m.add_wrapped(wrap_pyfunction!(dictionary::dictionary_indices))?; - m.add_wrapped(wrap_pyfunction!(list_flatten::list_flatten))?; - m.add_wrapped(wrap_pyfunction!(list_offsets::list_offsets))?; - m.add_wrapped(wrap_pyfunction!(struct_field::struct_field))?; m.add_wrapped(wrap_pyfunction!(take::take))?; Ok(()) diff --git a/arro3-core/python/arro3/core/__init__.py b/arro3-core/python/arro3/core/__init__.py index 15f81e2..b83cbfd 100644 --- a/arro3-core/python/arro3/core/__init__.py +++ b/arro3-core/python/arro3/core/__init__.py @@ -1,32 +1,4 @@ -from ._core import ( - Array, - ArrayReader, - ChunkedArray, - DataType, - Field, - RecordBatch, - RecordBatchReader, - Schema, - Table, - fixed_size_list_array, - list_array, - struct_array, - ___version, # noqa, -) +from ._core import * +from ._core import ___version __version__: str = ___version() - -__all__ = ( - "Array", - "ArrayReader", - "ChunkedArray", - "DataType", - "Field", - "RecordBatch", - "RecordBatchReader", - "Schema", - "Table", - "fixed_size_list_array", - "list_array", - "struct_array", -) diff --git a/arro3-core/python/arro3/core/_core.pyi b/arro3-core/python/arro3/core/_core.pyi index 447152f..fa676e2 100644 --- a/arro3-core/python/arro3/core/_core.pyi +++ b/arro3-core/python/arro3/core/_core.pyi @@ -20,7 +20,7 @@ class Array: /, type: ArrowSchemaExportable | None = None, ) -> None: - """Create arro3.core.Array instance from a sequence of Python objects. + """Create arro3.Array instance from a sequence of Python objects. Args: obj: A sequence of input objects. @@ -119,7 +119,7 @@ class ArrayReader: """A stream of Arrow `Array`s. This is similar to the [`RecordBatchReader`][arro3.core.RecordBatchReader] but each - item yielded from the stream is an [`Array`][arro3.core.Array], not a + item yielded from the stream is an [`Array`][arro3.Array], not a [`RecordBatch`][arro3.core.RecordBatch]. """ def __arrow_c_schema__(self) -> object: @@ -167,7 +167,7 @@ class ArrayReader: """Construct this from an existing Arrow object. This is an alias of and has the same behavior as - [`from_arrow`][arro3.core.ArrayReader.from_arrow], but is included for parity + [`from_arrow`][arro3.ArrayReader.from_arrow], but is included for parity with [`pyarrow.RecordBatchReader`][pyarrow.RecordBatchReader]. """ @property @@ -1719,21 +1719,124 @@ class Table: _description_ """ +@overload +def dictionary_dictionary(array: ArrowArrayExportable) -> Array: ... +@overload +def dictionary_dictionary(array: ArrowStreamExportable) -> ArrayReader: ... +def dictionary_dictionary( + array: ArrowArrayExportable | ArrowStreamExportable, +) -> Array | ArrayReader: + """ + Access the `dictionary` of a dictionary array. + + This is equivalent to the [`.dictionary`][pyarrow.DictionaryArray.dictionary] + attribute on a PyArrow [DictionaryArray][pyarrow.DictionaryArray]. + + Args: + array: Argument to compute function. + + Returns: + The keys of a dictionary-encoded array. + """ + +@overload +def dictionary_indices(array: ArrowArrayExportable) -> Array: ... +@overload +def dictionary_indices(array: ArrowStreamExportable) -> ArrayReader: ... +def dictionary_indices( + array: ArrowArrayExportable | ArrowStreamExportable, +) -> Array | ArrayReader: + """ + Access the indices of a dictionary array. + + This is equivalent to the [`.indices`][pyarrow.DictionaryArray.indices] + attribute on a PyArrow [DictionaryArray][pyarrow.DictionaryArray]. + + Args: + array: Argument to compute function. + + Returns: + The indices of a dictionary-encoded array. + """ + +@overload +def list_flatten(input: ArrowArrayExportable) -> Array: ... +@overload +def list_flatten(input: ArrowStreamExportable) -> ArrayReader: ... +def list_flatten( + input: ArrowArrayExportable | ArrowStreamExportable, +) -> Array | ArrayReader: + """Unnest this ListArray, LargeListArray or FixedSizeListArray. + + Args: + input: Input data. + + Raises: + Exception if not a list-typed array. + + Returns: + The flattened Arrow data. + """ + +@overload +def list_offsets(input: ArrowArrayExportable, *, logical: bool = True) -> Array: ... +@overload +def list_offsets( + input: ArrowStreamExportable, *, logical: bool = True +) -> ArrayReader: ... +def list_offsets( + input: ArrowArrayExportable | ArrowStreamExportable, + *, + logical: bool = True, +) -> Array | ArrayReader: + """Access the offsets of this ListArray or LargeListArray + + Args: + input: _description_ + physical: If False, return the physical (unsliced) offsets of the provided list array. If True, adjust the list offsets for the current array slicing. Defaults to `True`. + + Raises: + Exception if not a list-typed array. + + Returns: + _description_ + """ + +def struct_field( + values: ArrowArrayExportable, + /, + indices: int | Sequence[int], +) -> Array: + """Access a column within a StructArray by index + + Args: + values: Argument to compute function. + indices: List of indices for chained field lookup, for example [4, 1] will look up the second nested field in the fifth outer field. + + Raises: + Exception if not a struct-typed array. + + Returns: + _description_ + """ + def fixed_size_list_array( values: ArrowArrayExportable, list_size: int, *, type: ArrowSchemaExportable | None = None, ) -> Array: - """_summary_ + """Construct a new fixed size list array Args: - values: _description_ - list_size: _description_ - type: _description_. Defaults to None. + values: the values of the new fixed size list array + list_size: the number of elements in each item of the list. + + Keyword Args: + type: the type of output array. This must have fixed size list type. You may pass a `Field` into this parameter to associate extension metadata with the created array. Defaults to None, in which case it is inferred. Returns: - _description_ + a new Array with fixed size list type """ def list_array( @@ -1742,28 +1845,35 @@ def list_array( *, type: ArrowSchemaExportable | None = None, ) -> Array: - """_summary_ + """Construct a new list array Args: - offsets: _description_ - values: _description_ - type: _description_. Defaults to None. + offsets: the offsets for the output list array. This array must have type int32 or int64, depending on whether you wish to create a list array or large list array. + values: the values for the output list array. + + Keyword Args: + type: the type of output array. This must have list or large list type. You may pass a `Field` into this parameter to associate extension metadata with the created array. Defaults to None, in which case it is inferred. Returns: - _description_ + a new Array with list or large list type """ def struct_array( arrays: Sequence[ArrowArrayExportable], *, fields: Sequence[ArrowSchemaExportable], + type: ArrowSchemaExportable | None = None, ) -> Array: - """_summary_ + """Construct a new struct array Args: - arrays: _description_ - fields: _description_ + arrays: a sequence of arrays for the struct children + + Keyword Args: + fields: a sequence of fields that represent each of the struct children + type: the type of output array. This must have struct type. You may pass a `Field` into this parameter to associate extension metadata with the created array. Defaults to None, in which case it is inferred . + Returns: - _description_ + a new Array with struct type """ diff --git a/arro3-core/src/accessors/dictionary.rs b/arro3-core/src/accessors/dictionary.rs new file mode 100644 index 0000000..ca22e6b --- /dev/null +++ b/arro3-core/src/accessors/dictionary.rs @@ -0,0 +1,104 @@ +use arrow_array::cast::AsArray; +use arrow_array::ArrayRef; +use arrow_schema::{ArrowError, DataType, Field}; +use pyo3::prelude::*; +use pyo3_arrow::error::PyArrowResult; +use pyo3_arrow::ffi::ArrayIterator; +use pyo3_arrow::input::AnyArray; +use pyo3_arrow::{PyArray, PyArrayReader}; + +#[pyfunction] +pub(crate) fn dictionary_indices(py: Python, array: AnyArray) -> PyArrowResult { + match array { + AnyArray::Array(array) => { + let (array, _field) = array.into_inner(); + let output_array = _dictionary_indices(array)?; + Ok(PyArray::from_array_ref(output_array).to_arro3(py)?) + } + AnyArray::Stream(stream) => { + let reader = stream.into_reader()?; + let existing_field = reader.field(); + let out_field = match existing_field.data_type() { + DataType::Dictionary(key_type, _value_type) => { + Field::new("", *key_type.clone(), true) + } + _ => { + return Err(ArrowError::ComputeError( + "Expected dictionary-typed Array".to_string(), + ) + .into()) + } + }; + let iter = reader + .into_iter() + .map(move |array| _dictionary_indices(array?)); + Ok( + PyArrayReader::new(Box::new(ArrayIterator::new(iter, out_field.into()))) + .to_arro3(py)?, + ) + } + } +} + +/// Access the dictionary of the dictionary array +/// +/// This is equivalent to the `.dictionary` attribute on a PyArrow DictionaryArray. +#[pyfunction] +pub(crate) fn dictionary_dictionary(py: Python, array: AnyArray) -> PyArrowResult { + match array { + AnyArray::Array(array) => { + let (array, _field) = array.into_inner(); + let output_array = _dictionary_dictionary(array)?; + Ok(PyArray::from_array_ref(output_array).to_arro3(py)?) + } + AnyArray::Stream(stream) => { + let reader = stream.into_reader()?; + let existing_field = reader.field(); + let out_field = match existing_field.data_type() { + DataType::Dictionary(_key_type, value_type) => { + Field::new("", *value_type.clone(), true) + } + _ => { + return Err(ArrowError::ComputeError( + "Expected dictionary-typed Array".to_string(), + ) + .into()) + } + }; + let iter = reader + .into_iter() + .map(move |array| _dictionary_dictionary(array?)); + Ok( + PyArrayReader::new(Box::new(ArrayIterator::new(iter, out_field.into()))) + .to_arro3(py)?, + ) + } + } +} + +fn _dictionary_indices(array: ArrayRef) -> Result { + match array.data_type() { + DataType::Dictionary(_, _) => { + let dict_arr = array.as_any_dictionary(); + let keys_arr = dict_arr.keys(); + let keys_arr_ref = keys_arr.slice(0, keys_arr.len()); + Ok(keys_arr_ref) + } + _ => Err(ArrowError::ComputeError( + "Expected dictionary-typed Array".to_string(), + )), + } +} + +fn _dictionary_dictionary(array: ArrayRef) -> Result { + match array.data_type() { + DataType::Dictionary(_, _) => { + let dict_arr = array.as_any_dictionary(); + let values_arr = dict_arr.values().clone(); + Ok(values_arr) + } + _ => Err(ArrowError::ComputeError( + "Expected dictionary-typed Array".to_string(), + )), + } +} diff --git a/arro3-compute/src/list_flatten.rs b/arro3-core/src/accessors/list_flatten.rs similarity index 98% rename from arro3-compute/src/list_flatten.rs rename to arro3-core/src/accessors/list_flatten.rs index c2e86aa..0ebd24c 100644 --- a/arro3-compute/src/list_flatten.rs +++ b/arro3-core/src/accessors/list_flatten.rs @@ -1,4 +1,4 @@ -use arrow::array::AsArray; +use arrow_array::cast::AsArray; use arrow_array::ArrayRef; use arrow_schema::{ArrowError, DataType, FieldRef}; use pyo3::prelude::*; diff --git a/arro3-compute/src/list_offsets.rs b/arro3-core/src/accessors/list_offsets.rs similarity index 98% rename from arro3-compute/src/list_offsets.rs rename to arro3-core/src/accessors/list_offsets.rs index 45ae074..08e81f9 100644 --- a/arro3-compute/src/list_offsets.rs +++ b/arro3-core/src/accessors/list_offsets.rs @@ -1,6 +1,6 @@ use std::sync::Arc; -use arrow::array::AsArray; +use arrow_array::cast::AsArray; use arrow_array::{ArrayRef, Int32Array, Int64Array, OffsetSizeTrait}; use arrow_buffer::{OffsetBuffer, ScalarBuffer}; use arrow_schema::{ArrowError, DataType, Field}; diff --git a/arro3-core/src/accessors/mod.rs b/arro3-core/src/accessors/mod.rs new file mode 100644 index 0000000..2a44dfd --- /dev/null +++ b/arro3-core/src/accessors/mod.rs @@ -0,0 +1,4 @@ +pub(crate) mod dictionary; +pub(crate) mod list_flatten; +pub(crate) mod list_offsets; +pub(crate) mod struct_field; diff --git a/arro3-compute/src/struct_field.rs b/arro3-core/src/accessors/struct_field.rs similarity index 97% rename from arro3-compute/src/struct_field.rs rename to arro3-core/src/accessors/struct_field.rs index d948f84..f09f5eb 100644 --- a/arro3-compute/src/struct_field.rs +++ b/arro3-core/src/accessors/struct_field.rs @@ -1,4 +1,4 @@ -use arrow::array::AsArray; +use arrow_array::cast::AsArray; use arrow_array::ArrayRef; use arrow_schema::{ArrowError, DataType, FieldRef}; use pyo3::prelude::*; diff --git a/arro3-core/src/constructors.rs b/arro3-core/src/constructors.rs index cd72843..c156a85 100644 --- a/arro3-core/src/constructors.rs +++ b/arro3-core/src/constructors.rs @@ -9,7 +9,7 @@ use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use pyo3_arrow::error::PyArrowResult; -use pyo3_arrow::{PyArray, PyDataType, PyField}; +use pyo3_arrow::{PyArray, PyField}; #[pyfunction] #[pyo3(signature=(values, list_size, *, r#type=None))] @@ -17,13 +17,17 @@ pub(crate) fn fixed_size_list_array( py: Python, values: PyArray, list_size: i32, - r#type: Option, + r#type: Option, ) -> PyArrowResult { let (values_array, values_field) = values.into_inner(); - let list_data_type = r#type - .map(|t| t.into_inner()) - .unwrap_or_else(|| DataType::FixedSizeList(values_field.clone(), list_size)); - let inner_field = match &list_data_type { + let output_field = r#type.map(|t| t.into_inner()).unwrap_or_else(|| { + Arc::new(Field::new( + "", + DataType::FixedSizeList(values_field.clone(), list_size), + true, + )) + }); + let inner_field = match output_field.data_type() { DataType::FixedSizeList(inner_field, _) => inner_field, _ => { return Err( @@ -32,7 +36,7 @@ pub(crate) fn fixed_size_list_array( } }; let array = FixedSizeListArray::try_new(inner_field.clone(), list_size, values_array, None)?; - Ok(PyArray::new(Arc::new(array), Field::new("", list_data_type, true).into()).to_arro3(py)?) + Ok(PyArray::new(Arc::new(array), output_field).to_arro3(py)?) } #[pyfunction] @@ -41,7 +45,7 @@ pub(crate) fn list_array( py: Python, offsets: PyArray, values: PyArray, - r#type: Option, + r#type: Option, ) -> PyArrowResult { let (values_array, values_field) = values.into_inner(); let (offsets_array, _) = offsets.into_inner(); @@ -54,14 +58,18 @@ pub(crate) fn list_array( ) } }; - let list_data_type = r#type.map(|t| t.into_inner()).unwrap_or_else(|| { + let output_field = r#type.map(|t| t.into_inner()).unwrap_or_else(|| { if large_offsets { - DataType::LargeList(values_field.clone()) + Arc::new(Field::new( + "", + DataType::LargeList(values_field.clone()), + true, + )) } else { - DataType::List(values_field.clone()) + Arc::new(Field::new("", DataType::List(values_field.clone()), true)) } }); - let inner_field = match &list_data_type { + let inner_field = match output_field.data_type() { DataType::List(inner_field) | DataType::LargeList(inner_field) => inner_field, _ => { return Err( @@ -85,20 +93,29 @@ pub(crate) fn list_array( None, )?) }; - Ok(PyArray::new( - Arc::new(list_array), - Field::new("", list_data_type, true).into(), - ) - .to_arro3(py)?) + Ok(PyArray::new(Arc::new(list_array), output_field).to_arro3(py)?) } #[pyfunction] -#[pyo3(signature=(arrays, *, fields))] +#[pyo3(signature=(arrays, *, fields, r#type=None))] pub(crate) fn struct_array( py: Python, arrays: Vec, fields: Vec, + r#type: Option, ) -> PyArrowResult { + let output_field = r#type.map(|t| t.into_inner()).unwrap_or_else(|| { + let fields = fields + .into_iter() + .map(|field| field.into_inner()) + .collect::>(); + Arc::new(Field::new_struct("", fields, true)) + }); + let inner_fields = match output_field.data_type() { + DataType::Struct(inner_fields) => inner_fields.clone(), + _ => return Err(PyValueError::new_err("Expected struct as the outer data type").into()), + }; + let arrays = arrays .into_iter() .map(|arr| { @@ -106,12 +123,7 @@ pub(crate) fn struct_array( arr }) .collect::>(); - let fields = fields - .into_iter() - .map(|field| field.into_inner()) - .collect::>(); - let array = StructArray::try_new(fields.clone().into(), arrays, None)?; - let field = Field::new_struct("", fields, true); - Ok(PyArray::new(Arc::new(array), field.into()).to_arro3(py)?) + let array = StructArray::try_new(inner_fields, arrays, None)?; + Ok(PyArray::new(Arc::new(array), output_field).to_arro3(py)?) } diff --git a/arro3-core/src/lib.rs b/arro3-core/src/lib.rs index 07b1330..a958d56 100644 --- a/arro3-core/src/lib.rs +++ b/arro3-core/src/lib.rs @@ -1,5 +1,6 @@ use pyo3::prelude::*; +mod accessors; mod constructors; const VERSION: &str = env!("CARGO_PKG_VERSION"); @@ -24,6 +25,14 @@ fn _core(_py: Python, m: &Bound) -> PyResult<()> { m.add_class::()?; m.add_class::()?; + m.add_wrapped(wrap_pyfunction!( + accessors::dictionary::dictionary_dictionary + ))?; + m.add_wrapped(wrap_pyfunction!(accessors::dictionary::dictionary_indices))?; + m.add_wrapped(wrap_pyfunction!(accessors::list_flatten::list_flatten))?; + m.add_wrapped(wrap_pyfunction!(accessors::list_offsets::list_offsets))?; + m.add_wrapped(wrap_pyfunction!(accessors::struct_field::struct_field))?; + m.add_wrapped(wrap_pyfunction!(constructors::fixed_size_list_array))?; m.add_wrapped(wrap_pyfunction!(constructors::list_array))?; m.add_wrapped(wrap_pyfunction!(constructors::struct_array))?; diff --git a/docs/api/core/accessors.md b/docs/api/core/accessors.md new file mode 100644 index 0000000..e3213c8 --- /dev/null +++ b/docs/api/core/accessors.md @@ -0,0 +1,10 @@ +# Accessors + +::: arro3.core + options: + members: + - dictionary_dictionary + - dictionary_indices + - list_flatten + - list_offsets + - struct_field diff --git a/docs/api/core/constructors.md b/docs/api/core/constructors.md new file mode 100644 index 0000000..f08acff --- /dev/null +++ b/docs/api/core/constructors.md @@ -0,0 +1,8 @@ +# Constructors + +::: arro3.core + options: + members: + - fixed_size_list_array + - list_array + - struct_array diff --git a/mkdocs.yml b/mkdocs.yml index 4d45052..18dcfa0 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -30,6 +30,8 @@ nav: - api/core/record-batch-reader.md - api/core/schema.md - api/core/table.md + - api/core/constructors.md + - api/core/accessors.md - api/core/types.md - api/compute.md - api/io.md diff --git a/tests/core/test_constructors.py b/tests/core/test_constructors.py index 39df4e7..f8e1f7a 100644 --- a/tests/core/test_constructors.py +++ b/tests/core/test_constructors.py @@ -8,7 +8,7 @@ list_array, struct_array, ) -from arro3.compute import list_offsets +from arro3.core import list_offsets def test_fixed_size_list_array(): diff --git a/tests/compute/test_list_flatten.py b/tests/core/test_list_flatten.py similarity index 93% rename from tests/compute/test_list_flatten.py rename to tests/core/test_list_flatten.py index 9fb0a35..3d7920e 100644 --- a/tests/compute/test_list_flatten.py +++ b/tests/core/test_list_flatten.py @@ -1,5 +1,5 @@ import pyarrow as pa -from arro3.compute import list_flatten +from arro3.core import list_flatten def test_list_flatten(): diff --git a/tests/compute/test_list_offsets.py b/tests/core/test_list_offsets.py similarity index 95% rename from tests/compute/test_list_offsets.py rename to tests/core/test_list_offsets.py index c9452db..bd6afd7 100644 --- a/tests/compute/test_list_offsets.py +++ b/tests/core/test_list_offsets.py @@ -1,5 +1,5 @@ import pyarrow as pa -from arro3.compute import list_offsets +from arro3.core import list_offsets def test_list_flatten(): diff --git a/tests/compute/test_struct_field.py b/tests/core/test_struct_field.py similarity index 97% rename from tests/compute/test_struct_field.py rename to tests/core/test_struct_field.py index 2b5aa5f..f270bac 100644 --- a/tests/compute/test_struct_field.py +++ b/tests/core/test_struct_field.py @@ -1,6 +1,6 @@ import pyarrow as pa import pytest -from arro3.compute import struct_field +from arro3.core import struct_field def test_struct_field(): diff --git a/tests/compute/test_dictionary.py b/tests/test_dictionary.py similarity index 95% rename from tests/compute/test_dictionary.py rename to tests/test_dictionary.py index fcce444..25c6053 100644 --- a/tests/compute/test_dictionary.py +++ b/tests/test_dictionary.py @@ -1,8 +1,9 @@ -import pyarrow as pa from datetime import datetime + +import pyarrow as pa import pyarrow.compute as pc -from arro3.core import ChunkedArray -from arro3.compute import dictionary_encode, dictionary_dictionary, dictionary_indices +from arro3.compute import dictionary_encode +from arro3.core import ChunkedArray, dictionary_dictionary, dictionary_indices def test_dictionary_encode():