Export list and struct array constructors (#71)

kylebarron · Jul 26, 2024 · e584d0d · e584d0d
1 parent 264f20f
commit e584d0d
Show file tree

Hide file tree

Showing 4 changed files with 66 additions and 6 deletions.
diff --git a/arro3-core/python/arro3/core/__init__.py b/arro3-core/python/arro3/core/__init__.py
@@ -8,6 +8,9 @@
     RecordBatchReader,
     Schema,
     Table,
+    fixed_size_list_array,
+    list_array,
+    struct_array,
     ___version,  # noqa,
 )
 
@@ -23,4 +26,7 @@
     "RecordBatchReader",
     "Schema",
     "Table",
+    "fixed_size_list_array",
+    "list_array",
+    "struct_array",
 )
diff --git a/arro3-core/python/arro3/core/_rust.pyi b/arro3-core/python/arro3/core/_rust.pyi
@@ -591,3 +591,52 @@ class Table:
         Returns:
             _description_
         """
+
+def fixed_size_list_array(
+    values: ArrowArrayExportable,
+    list_size: int,
+    *,
+    type: ArrowSchemaExportable | None = None,
+) -> Array:
+    """_summary_
+
+    Args:
+        values: _description_
+        list_size: _description_
+        type: _description_. Defaults to None.
+
+    Returns:
+        _description_
+    """
+
+def list_array(
+    offsets: ArrowArrayExportable,
+    values: ArrowArrayExportable,
+    *,
+    type: ArrowSchemaExportable | None = None,
+) -> Array:
+    """_summary_
+
+    Args:
+        offsets: _description_
+        values: _description_
+        type: _description_. Defaults to None.
+
+    Returns:
+        _description_
+    """
+
+def struct_array(
+    arrays: Sequence[ArrowArrayExportable],
+    *,
+    fields: Sequence[ArrowSchemaExportable],
+) -> Array:
+    """_summary_
+
+    Args:
+        arrays: _description_
+        fields: _description_
+
+    Returns:
+        _description_
+    """
diff --git a/arro3-core/src/constructors.rs b/arro3-core/src/constructors.rs
@@ -13,8 +13,7 @@ use pyo3_arrow::{PyArray, PyField};
 
 #[pyfunction]
 #[pyo3(signature=(values, list_size, *, r#type=None))]
-#[allow(dead_code)]
-fn fixed_size_list_array(
+pub(crate) fn fixed_size_list_array(
     py: Python,
     values: PyArray,
     list_size: i32,
@@ -36,8 +35,7 @@ fn fixed_size_list_array(
 
 #[pyfunction]
 #[pyo3(signature=(offsets, values, *, r#type=None))]
-#[allow(dead_code)]
-fn list_array(
+pub(crate) fn list_array(
     py: Python,
     offsets: PyArray,
     values: PyArray,
@@ -82,8 +80,11 @@ fn list_array(
 
 #[pyfunction]
 #[pyo3(signature=(arrays, *, fields))]
-#[allow(dead_code)]
-fn struct_array(py: Python, arrays: Vec<PyArray>, fields: Vec<PyField>) -> PyArrowResult<PyObject> {
+pub(crate) fn struct_array(
+    py: Python,
+    arrays: Vec<PyArray>,
+    fields: Vec<PyField>,
+) -> PyArrowResult<PyObject> {
     let arrays = arrays
         .into_iter()
         .map(|arr| {

diff --git a/arro3-core/src/lib.rs b/arro3-core/src/lib.rs
@@ -24,5 +24,9 @@ fn _rust(_py: Python, m: &Bound<PyModule>) -> PyResult<()> {
     m.add_class::<pyo3_arrow::PySchema>()?;
     m.add_class::<pyo3_arrow::PyTable>()?;
 
+    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))?;
+
     Ok(())
 }