feat: chunked array handling

ptars-python/src/lib.rs

@@ -191,7 +191,7 @@ impl MessageHandler {
     /// Convert a binary array of serialized protobuf messages to a record batch.
     ///
     /// Each element in the binary array is expected to be a serialized protobuf message.
-    /// The resulting record batch will have one column per field in the message descriptor.
+    /// Returns a `pyarrow.RecordBatch`.
     fn array_to_record_batch(&self, array: &Bound<PyAny>, py: Python<'_>) -> PyResult<Py<PyAny>> {
         let array_data = arrow::array::ArrayData::from_pyarrow_bound(array).map_err(|e| {
             pyo3::exceptions::PyTypeError::new_err(format!("Failed to convert array: {}", e))
@@ -206,6 +206,41 @@ impl MessageHandler {
         Ok(record_batch.to_pyarrow(py)?.unbind())
     }
 
+    /// Convert a chunked binary array of serialized protobuf messages to a table.
+    ///
+    /// Each chunk produces a record batch. Empty chunks are skipped.
+    /// Returns a `pyarrow.Table`.
+    fn chunked_array_to_table(
+        &self,
+        chunked_array: &Bound<PyAny>,
+        py: Python<'_>,
+    ) -> PyResult<Py<PyAny>> {
+        let chunks: Vec<Bound<PyAny>> = chunked_array.getattr("chunks")?.extract()?;
+
+        let mut batches: Vec<Py<PyAny>> = Vec::new();
+        for chunk in &chunks {
+            let array_data = arrow::array::ArrayData::from_pyarrow_bound(chunk).map_err(|e| {
+                pyo3::exceptions::PyTypeError::new_err(format!("Failed to convert chunk: {}", e))
+            })?;
+            if !array_data.is_empty() {
+                let arrow_array = BinaryArray::from(array_data);
+                let record_batch = ptars::binary_array_to_record_batch_direct(
+                    &arrow_array,
+                    &self.message_descriptor,
+                    &self.config,
+                )
+                .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
+                batches.push(record_batch.to_pyarrow(py)?.unbind());
+            }
+        }
+
+        let pa = py.import("pyarrow")?;
+        let table_class = pa.getattr("Table")?;
+        let batch_list = pyo3::types::PyList::new(py, &batches)?;
+        let table = table_class.call_method1("from_batches", (batch_list,))?;
+        Ok(table.unbind())
+    }
+
     /// Read size-delimited protobuf messages from a file and convert to a record batch.
     ///
     /// Each message in the file should be preceded by its size encoded as a varint.

ptars/src/converter.rs

@@ -2,10 +2,7 @@
 mod tests {
     use crate::arrow_to_proto::record_batch_to_array;
     use crate::config::PtarsConfig;
-    use crate::proto_to_arrow::{
-        binary_array_to_messages, binary_array_to_record_batch_direct, messages_to_record_batch,
-        messages_to_record_batch_with_config,
-    };
+    use crate::proto_to_arrow::{messages_to_record_batch, messages_to_record_batch_with_config};
     use arrow::array::Array;
     use prost_reflect::prost_types::{
         field_descriptor_proto::{Label, Type},

python/test/unit/test_ptars.py

@@ -502,6 +502,55 @@ def test_array_to_record_batch_complex_message(pool):
     assert record_batch["bool_value"].to_pylist() == [True, False]
 
 
+def test_chunked_array_to_table(pool):
+    """Test chunked_array_to_table returns a Table with one batch per chunk."""
+    handler = pool.get_for_message(SearchRequest.DESCRIPTOR)
+
+    messages = [
+        SearchRequest(query="hello", page_number=1, result_per_page=10),
+        SearchRequest(query="world", page_number=2, result_per_page=20),
+        SearchRequest(query="foo", page_number=3, result_per_page=30),
+    ]
+    payloads = [message.SerializeToString() for message in messages]
+
+    # Create a chunked binary array from two separate chunks
+    chunk1 = pa.array(payloads[:2], type=pa.binary())
+    chunk2 = pa.array(payloads[2:], type=pa.binary())
+    chunked_array = pa.chunked_array([chunk1, chunk2])
+
+    table = handler.chunked_array_to_table(chunked_array)
+
+    assert isinstance(table, pa.Table)
+    assert table.num_rows == 3
+    assert table.column("query").to_pylist() == ["hello", "world", "foo"]
+    assert table.column("page_number").to_pylist() == [1, 2, 3]
+    assert table.column("result_per_page").to_pylist() == [10, 20, 30]
+    # Each chunk becomes a separate batch
+    assert len(table.to_batches()) == 2
+
+
+def test_chunked_array_to_table_skips_empty_chunks(pool):
+    """Test that empty chunks are skipped."""
+    handler = pool.get_for_message(SearchRequest.DESCRIPTOR)
+
+    messages = [
+        SearchRequest(query="hello", page_number=1, result_per_page=10),
+    ]
+    payloads = [message.SerializeToString() for message in messages]
+
+    chunk1 = pa.array([], type=pa.binary())
+    chunk2 = pa.array(payloads, type=pa.binary())
+    chunk3 = pa.array([], type=pa.binary())
+    chunked_array = pa.chunked_array([chunk1, chunk2, chunk3])
+
+    table = handler.chunked_array_to_table(chunked_array)
+
+    assert isinstance(table, pa.Table)
+    assert table.num_rows == 1
+    assert len(table.to_batches()) == 1
+    assert table.column("query").to_pylist() == ["hello"]
+
+
 def test_nested_primitive():
     data = [
         NestedExampleMessage(example_message=ExampleMessage()),