diff --git a/ptars-python/src/lib.rs b/ptars-python/src/lib.rs index a22909a..ee34b95 100644 --- a/ptars-python/src/lib.rs +++ b/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, py: Python<'_>) -> PyResult> { 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, + py: Python<'_>, + ) -> PyResult> { + let chunks: Vec> = chunked_array.getattr("chunks")?.extract()?; + + let mut batches: Vec> = 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. diff --git a/ptars/src/converter.rs b/ptars/src/converter.rs index dc2f084..a66f35f 100644 --- a/ptars/src/converter.rs +++ b/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}, diff --git a/python/test/unit/test_ptars.py b/python/test/unit/test_ptars.py index 1e13717..604d1d8 100644 --- a/python/test/unit/test_ptars.py +++ b/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()),