feat(python): Add column-wise buffer builder

python/src/nanoarrow/_lib.pyx

@@ -1009,6 +1009,11 @@ cdef class CSchemaView:
         if self.extension_name or self._schema_view.type != self._schema_view.storage_type:
             return None
 
+        # String/binary types do not have format strings as far as the Python
+        # buffer protocol is concerned
+        if self.layout.n_buffers != 2:
+            return None
+
         cdef char out[128]
         cdef int element_size_bits = 0
         if self._schema_view.type == NANOARROW_TYPE_FIXED_SIZE_BINARY:
@@ -1632,6 +1637,22 @@ cdef class CArrayView:
 
     @property
     def null_count(self):
+        if self._ptr.null_count != -1:
+            return self._ptr.null_count
+
+        cdef ArrowBufferType buffer_type = self._ptr.layout.buffer_type[0]
+        cdef uint8_t* validity_bits = self._ptr.buffer_views[0].data.as_uint8
+
+        if buffer_type != NANOARROW_BUFFER_TYPE_VALIDITY:
+            self._ptr.null_count = 0
+        elif validity_bits == NULL:
+            self._ptr.null_count = 0
+        elif self._device is DEVICE_CPU:
+            self._ptr.null_count = (
+                self._ptr.length -
+                ArrowBitCountSet(validity_bits, self.offset, self.length)
+            )
+
         return self._ptr.null_count
 
     @property
@@ -1869,7 +1890,7 @@ cdef class CBufferView:
         return self._format.decode("UTF-8")
 
     @property
-    def item_size(self):
+    def itemsize(self):
         return self._strides
 
     def __len__(self):
@@ -1957,7 +1978,7 @@ cdef class CBufferView:
 
         cdef int64_t c_offset = offset
         cdef int64_t c_length = length
-        cdef int64_t c_item_size = self.item_size
+        cdef int64_t c_item_size = self.itemsize
         cdef int64_t c_dest_offset = dest_offset
         self._check_copy_into_bounds(&buffer, c_offset, c_length, dest_offset, c_item_size)
 
@@ -2010,7 +2031,7 @@ cdef class CBufferView:
         if length is None:
             length = self.n_elements
 
-        cdef int64_t bytes_to_copy = length * self.item_size
+        cdef int64_t bytes_to_copy = length * self.itemsize
         out = CBufferBuilder().set_data_type(self.data_type_id)
         out.reserve_bytes(bytes_to_copy)
         self.copy_into(out, offset, length)
@@ -2224,9 +2245,9 @@ cdef class CBuffer:
         return self._element_size_bits
 
     @property
-    def item_size(self):
+    def itemsize(self):
         self._assert_valid()
-        return self._view.item_size
+        return self._view.itemsize
 
     @property
     def format(self):
@@ -2339,6 +2360,13 @@ cdef class CBufferBuilder:
         """The number of bytes that have been written to this buffer"""
         return self._buffer.size_bytes
 
+    @property
+    def itemsize(self):
+        return self._buffer.itemsize
+
+    def __len__(self):
+        return self._buffer.size_bytes // self.itemsize
+
     @property
     def capacity_bytes(self):
         """The number of bytes allocated in the underlying buffer"""

python/src/nanoarrow/array.py

@@ -377,7 +377,7 @@ def to_columns(self) -> Tuple[str, Sequence]:
         >>> names
         ['col1']
         >>> columns
-        [[1, 2, 3]]
+        [nanoarrow.c_lib.CBuffer(int64[24 b] 1 2 3)]
         """
         return to_columns(self)
 

python/src/nanoarrow/array_stream.py

@@ -232,7 +232,7 @@ def to_columns(self) -> Tuple[str, Sequence]:
         >>> names
         ['col1']
         >>> columns
-        [[1, 2, 3]]
+        [nanoarrow.c_lib.CBuffer(int64[24 b] 1 2 3)]
         """
         return to_columns(self)
 

python/src/nanoarrow/iterator.py

@@ -156,11 +156,7 @@ def _object_label(self):
             return f"<unnamed {self._schema_view.type}>"
 
     def _contains_nulls(self):
-        return (
-            self._schema_view.nullable
-            and len(self._array_view.buffer(0))
-            and self._array_view.null_count != 0
-        )
+        return self._schema_view.nullable and self._array_view.null_count != 0
 
     def _set_array(self, array):
         self._array_view._set_array(array)

python/src/nanoarrow/visitor.py

@@ -15,10 +15,11 @@
 # specific language governing permissions and limitations
 # under the License.
 
-from typing import Any, List, Sequence, Tuple, Union
+from typing import Any, Callable, List, Sequence, Tuple, Union
 
-from nanoarrow._lib import CArrayView
+from nanoarrow._lib import CArrayView, CArrowType, CBuffer, CBufferBuilder
 from nanoarrow.c_array_stream import c_array_stream
+from nanoarrow.c_schema import c_schema_view
 from nanoarrow.iterator import ArrayViewBaseIterator, PyIterator
 from nanoarrow.schema import Type
 
@@ -49,7 +50,7 @@ def to_pylist(obj, schema=None) -> List:
     return ListBuilder.visit(obj, schema)
 
 
-def to_columns(obj, schema=None) -> Tuple[List[str], List[Sequence]]:
+def to_columns(obj, schema=None, handle_nulls=None) -> Tuple[List[str], List[Sequence]]:
     """Convert ``obj`` to a ``list()` of sequences
 
     Converts a stream of struct arrays into its column-wise representation
@@ -74,9 +75,60 @@ def to_columns(obj, schema=None) -> Tuple[List[str], List[Sequence]]:
     >>> names
     ['col1']
     >>> columns
-    [[1, 2, 3]]
+    [nanoarrow.c_lib.CBuffer(int64[24 b] 1 2 3)]
     """
-    return ColumnsBuilder.visit(obj, schema)
+    return ColumnsBuilder.visit(obj, schema, handle_nulls=handle_nulls)
+
+
+def nulls_forbid() -> Callable[[CBuffer, Sequence], Sequence]:
+    def handle(is_valid, data):
+        if len(is_valid) > 0:
+            raise ValueError("Null present with null_handler=nulls_forbid()")
+
+        return data
+
+    return handle
+
+
+def nulls_debug() -> Callable[[CBuffer, Sequence], Tuple[CBuffer, Sequence]]:
+    def handle(is_valid, data):
+        return is_valid, data
+
+    return handle
+
+
+def nulls_as_sentinel(sentinel=None):
+    from numpy import array, result_type
+
+    def handle(is_valid, data):
+        is_valid = array(is_valid, copy=False)
+        data = array(data, copy=False)
+
+        if len(is_valid) > 0:
+            out_type = result_type(data, sentinel)
+            data = array(data, dtype=out_type, copy=True)
+            data[~is_valid] = sentinel
+            return data
+        else:
+            return data
+
+    return handle
+
+
+def nulls_as_masked_array():
+    from numpy import array
+    from numpy.ma import masked_array
+
+    def handle(is_valid, data):
+        is_valid = array(is_valid, copy=False)
+        data = array(data, copy=False)
+
+        if len(is_valid) > 0:
+            return masked_array(data, ~is_valid)
+        else:
+            return data
+
+    return handle
 
 
 class ArrayStreamVisitor(ArrayViewBaseIterator):
@@ -144,7 +196,7 @@ def finish(self) -> List:
 
 
 class ColumnsBuilder(ArrayStreamVisitor):
-    def __init__(self, schema, *, array_view=None):
+    def __init__(self, schema, handle_nulls=None, *, array_view=None):
         super().__init__(schema, array_view=array_view)
 
         if self.schema.type != Type.STRUCT:
@@ -156,18 +208,23 @@ def __init__(self, schema, *, array_view=None):
             self._schema.children, self._array_view.children
         ):
             self._child_visitors.append(
-                self._resolve_child_visitor(child_schema, child_array_view)
+                self._resolve_child_visitor(
+                    child_schema, child_array_view, handle_nulls
+                )
             )
 
-    def _resolve_child_visitor(self, child_schema, child_array_view):
-        # TODO: Resolve more efficient column builders for single-buffer types
-        return ListBuilder(child_schema, array_view=child_array_view)
+    def _resolve_child_visitor(self, child_schema, child_array_view, handle_nulls):
+        cls, kwargs = _resolve_column_builder_cls(child_schema, handle_nulls)
+        return cls(child_schema, **kwargs, array_view=child_array_view)
 
     def begin(self, total_elements: Union[int, None] = None) -> None:
         for child_visitor in self._child_visitors:
             child_visitor.begin(total_elements)
 
     def visit_chunk_view(self, array_view: CArrayView) -> Any:
+        if array_view.null_count > 0:
+            raise ValueError("null_count > 0 encountered in ColumnsBuilder")
+
         for child_visitor, child_array_view in zip(
             self._child_visitors, array_view.children
         ):
@@ -177,3 +234,132 @@ def finish(self) -> Tuple[List[str], List[Sequence]]:
         return [v.schema.name for v in self._child_visitors], [
             v.finish() for v in self._child_visitors
         ]
+
+
+class BufferColumnBuilder(ArrayStreamVisitor):
+    def begin(self, total_elements: Union[int, None]):
+        self._builder = CBufferBuilder()
+        self._builder.set_format(self._schema_view.buffer_format)
+
+        if total_elements is not None:
+            element_size_bits = self._schema_view.layout.element_size_bits[1]
+            element_size_bytes = element_size_bits // 8
+            self._builder.reserve_bytes(total_elements * element_size_bytes)
+
+    def visit_chunk_view(self, array_view: CArrayView) -> None:
+        builder = self._builder
+        offset, length = array_view.offset, array_view.length
+        dst_bytes = length * builder.itemsize
+
+        builder.reserve_bytes(dst_bytes)
+        array_view.buffer(1).copy_into(builder, offset, length, len(builder))
+        builder.advance(dst_bytes)
+
+    def finish(self) -> Any:
+        return self._builder.finish()
+
+
+class BooleanColumnBuilder(ArrayStreamVisitor):
+    def begin(self, total_elements: Union[int, None]):
+        self._builder = CBufferBuilder()
+        self._builder.set_format("?")
+
+        if total_elements is not None:
+            self._builder.reserve_bytes(total_elements)
+
+    def visit_chunk_view(self, array_view: CArrayView) -> None:
+        builder = self._builder
+        offset, length = array_view.offset, array_view.length
+        builder.reserve_bytes(length)
+        array_view.buffer(1).unpack_bits_into(builder, offset, length, len(builder))
+        builder.advance(length)
+
+    def finish(self) -> Any:
+        return self._builder.finish()
+
+
+class NullableColumnBuilder(ArrayStreamVisitor):
+    def __init__(
+        self,
+        schema,
+        column_builder_cls=BufferColumnBuilder,
+        handle_nulls: Union[Callable[[CBuffer, Sequence], Any], None] = None,
+        *,
+        array_view=None
+    ):
+        super().__init__(schema, array_view=array_view)
+        self._column_builder = column_builder_cls(schema, array_view=self._array_view)
+
+        if handle_nulls is None:
+            self._handle_nulls = nulls_forbid()
+        else:
+            self._handle_nulls = handle_nulls
+
+    def begin(self, total_elements: Union[int, None]):
+        self._builder = CBufferBuilder()
+        self._builder.set_format("?")
+        self._length = 0
+
+        self._column_builder.begin(total_elements)
+
+    def visit_chunk_view(self, array_view: CArrayView) -> None:
+        offset, length = array_view.offset, array_view.length
+
+        builder = self._builder
+        chunk_contains_nulls = array_view.null_count != 0
+        bitmap_allocated = len(builder) > 0
+
+        if chunk_contains_nulls:
+            current_length = self._length
+            if not bitmap_allocated:
+                self._fill_valid(current_length)
+
+            builder.reserve_bytes(length)
+            array_view.buffer(0).unpack_bits_into(
+                builder, offset, length, current_length
+            )
+            builder.advance(length)
+
+        elif bitmap_allocated:
+            self._fill_valid(length)
+
+        self._length += length
+        self._column_builder.visit_chunk_view(array_view)
+
+    def finish(self) -> Any:
+        is_valid = self._builder.finish()
+        column = self._column_builder.finish()
+        return self._handle_nulls(is_valid, column)
+
+    def _fill_valid(self, length):
+        builder = self._builder
+        builder.reserve_bytes(length)
+        out_start = len(builder)
+        memoryview(builder)[out_start : out_start + length] = b"\x01" * length
+        builder.advance(length)
+
+
+def _resolve_column_builder_cls(schema, handle_nulls=None):
+    schema_view = c_schema_view(schema)
+
+    if schema_view.nullable:
+        if schema_view.type_id == CArrowType.BOOL:
+            return NullableColumnBuilder, {
+                "column_builder_cls": BooleanColumnBuilder,
+                "handle_nulls": handle_nulls,
+            }
+        elif schema_view.buffer_format is not None:
+            return NullableColumnBuilder, {
+                "column_builder_cls": BufferColumnBuilder,
+                "handle_nulls": handle_nulls,
+            }
+        else:
+            return ListBuilder, {}
+    else:
+
+        if schema_view.type_id == CArrowType.BOOL:
+            return BooleanColumnBuilder, {}
+        elif schema_view.buffer_format is not None:
+            return BufferColumnBuilder, {}
+        else:
+            return ListBuilder, {}

python/tests/test_array_stream.py

@@ -82,8 +82,8 @@ def test_array_stream_to_columns():
     stream = na.ArrayStream(c_array)
     names, columns = stream.to_columns()
     assert names == ["col1", "col2"]
-    assert columns[0] == [1, 2, 3]
-    assert columns[1] == ["a", "b", "c"]
+    assert list(columns[0]) == [1, 2, 3]
+    assert list(columns[1]) == ["a", "b", "c"]
 
 
 def test_array_stream_read_all():