Import Arrow UTF-8 columns directly from buffers

Author: FrancescAlted

src/blosc2/_utf8_array.py

@@ -678,6 +678,59 @@ def extend(self, values: Iterable[Any]) -> None:
                     self._pending_chars += len(v)
             self._flush_if_needed()
 
+    def extend_arrow(self, array) -> None:
+        """Append an Arrow string array directly from offsets and UTF-8 buffers.
+
+        No Python strings are materialized. Null slots use the same physical
+        fill as ``extend``; CTable remains responsible for validity sidecars.
+        Sliced arrays are rebased, and view-based strings are cast within Arrow.
+        """
+        import pyarrow as pa
+        import pyarrow.compute as pc
+
+        self._check_writable()
+        if not isinstance(array, (pa.Array, pa.ChunkedArray)):
+            raise TypeError("extend_arrow requires an Arrow string array")
+        is_view = getattr(pa.types, "is_string_view", lambda _: False)(array.type)
+        if not (pa.types.is_string(array.type) or pa.types.is_large_string(array.type) or is_view):
+            raise TypeError("extend_arrow requires an Arrow string array")
+        if isinstance(array, pa.ChunkedArray):
+            for chunk in array.chunks:
+                self.extend_arrow(chunk)
+            return
+        if is_view:
+            array = array.cast(pa.large_string())
+        array.validate(full=True)
+        if not len(array):
+            return
+        if array.null_count:
+            array = pc.fill_null(array, self._coerce(None))
+        _, offsets_buffer, data_buffer = array.buffers()
+        dtype = np.dtype("<i8" if pa.types.is_large_string(array.type) else "<i4")
+        offsets = np.frombuffer(
+            offsets_buffer, dtype=dtype, count=len(array) + 1, offset=array.offset * dtype.itemsize
+        )
+        first, last = int(offsets[0]), int(offsets[-1])
+        data = (
+            np.frombuffer(data_buffer, dtype=np.uint8, count=last - first, offset=first)
+            if last > first
+            else np.empty(0, dtype=np.uint8)
+        )
+        # Preserve the ordering of any earlier Python-string appends.
+        self.flush()
+        start = self._bytes_used
+        rows = self._persisted_rows + len(array)
+        used = start + len(data)
+        self._data.resize((max(used, 1),))
+        if len(data):
+            self._data[start:used] = data
+        self._offsets.resize((rows + 1,))
+        self._offsets[self._persisted_rows + 1 : rows + 1] = (
+            offsets[1:].astype(np.int64, copy=False) - first + start
+        )
+        self._persisted_rows = rows
+        self._bytes_used_cache = used
+
     def flush(self) -> None:
         """Write pending rows to the backing offsets/data NDArrays."""
         self._check_open()

src/blosc2/ctable.py

@@ -9392,7 +9392,10 @@ def _write_arrow_batch(
                     obj._ensure_null_mask(col.name)[pos : pos + m] = arrow_col.is_valid().to_numpy(
                         zero_copy_only=False
                     )
-                new_cols[col.name].extend(arrow_col.to_pylist())
+                if cls._is_utf8_column(col):
+                    new_cols[col.name].extend_arrow(arrow_col)
+                else:
+                    new_cols[col.name].extend(arrow_col.to_pylist())
             elif cls._is_dictionary_column(col):
                 import pyarrow as _pa
 

tests/ctable/test_arrow_interop.py

@@ -19,6 +19,32 @@
 pa = pytest.importorskip("pyarrow")
 
 
+@pytest.mark.parametrize("arrow_kind", ["string", "large_string", "string_view"])
+@pytest.mark.parametrize("null_storage", ["mask", "sentinel"])
+@pytest.mark.parametrize("persistent", [False, True])
+def test_arrow_utf8_import_avoids_python_strings(
+    arrow_kind, null_storage, persistent, tmp_path, monkeypatch
+):
+    from blosc2._utf8_array import UTF8Array
+
+    if not hasattr(pa, arrow_kind):
+        pytest.skip(f"Installed PyArrow has no {arrow_kind}")
+    arrow_type = getattr(pa, arrow_kind)()
+    values = ["ignored", "café", "", None, "日本語", "emoji 🎉", "a\0b", None, "ignored"]
+    array = pa.array(values, type=arrow_type).slice(1, 7)
+    source = pa.table({"name": array})
+    monkeypatch.setattr(UTF8Array, "extend", lambda *_: pytest.fail("Python-string ingestion used"))
+    options = {"urlpath": str(tmp_path / "utf8.b2z")} if persistent else {}
+    with CTable.from_arrow(
+        source.schema, iter(source.to_batches(max_chunksize=3)), null_storage=null_storage, **options
+    ) as table:
+        assert table.to_arrow().column("name").to_pylist() == values[1:8]
+        assert table["name"].null_count() == 2
+    if persistent:
+        with CTable.open(options["urlpath"]) as table:
+            assert table.to_arrow().column("name").to_pylist() == values[1:8]
+
+
 @pytest.mark.parametrize("first_rows", [0, 1, 7])
 @pytest.mark.parametrize("producer", ["schema", "reader"])
 def test_unknown_arrow_stream_grid_is_independent_of_first_batch(first_rows, producer, monkeypatch):

tests/ctable/test_utf8.py

@@ -54,6 +54,61 @@ def make_table(values=None, **kwargs):
     return CTable(Row, new_data={"name": list(values), "x": list(range(len(values)))}, **kwargs)
 
 
+@pytest.mark.parametrize("large", [False, True])
+@pytest.mark.parametrize("null_storage", ["mask", "sentinel"])
+def test_utf8_extend_arrow_buffers_and_pending_order(large, null_storage):
+    pa = pytest.importorskip("pyarrow")
+    from blosc2._utf8_array import UTF8Array
+
+    options = {"null_value": "<NULL>"} if null_storage == "sentinel" else {}
+    backend = UTF8Array(blosc2.utf8(nullable=True, null_storage=null_storage, **options))
+    backend.extend(["pending"])
+    values = ["ignored", "café", "", None, "日本語", "a\0b", "ignored"]
+    array = pa.array(values, type=pa.large_string() if large else pa.string()).slice(1, 5)
+    backend.extend_arrow(array)
+    backend.extend_arrow(pa.chunked_array([array.slice(0, 2), array.slice(2)]))
+    backend.append("last")
+    backend.flush()
+    expected = ["pending"] + [backend._coerce(value) for value in values[1:6]] * 2 + ["last"]
+    assert backend[:].tolist() == expected
+    assert len(backend) == len(expected)
+    empty = pa.array([], type=array.type)
+    backend.extend_arrow(empty)
+    assert backend[:].tolist() == expected
+
+
+@pytest.mark.parametrize("values", [[], ["", ""], [None, None]])
+@pytest.mark.parametrize("large", [False, True])
+def test_utf8_extend_arrow_empty_byte_spans(values, large):
+    pa = pytest.importorskip("pyarrow")
+    from blosc2._utf8_array import UTF8Array
+
+    backend = UTF8Array(blosc2.utf8(nullable=True, null_storage="mask"))
+    arrow_type = pa.large_string() if large else pa.string()
+    array = pa.array(["prefix", *values, "suffix"], type=arrow_type).slice(1, len(values))
+    backend.extend_arrow(array)
+    backend.extend_arrow(pa.array(["tail"], type=arrow_type))
+    assert backend[:].tolist() == [""] * len(values) + ["tail"]
+
+
+def test_utf8_extend_arrow_rejects_wrong_types_and_invalid_utf8():
+    pa = pytest.importorskip("pyarrow")
+    from blosc2._utf8_array import UTF8Array
+
+    backend = UTF8Array(blosc2.utf8())
+    for wrong in (["hello"], pa.array([1]), pa.chunked_array([], type=pa.int64())):
+        with pytest.raises(TypeError, match="Arrow string array"):
+            backend.extend_arrow(wrong)
+    with pytest.raises(TypeError, match="not nullable"):
+        backend.extend_arrow(pa.array([None], type=pa.string()))
+    invalid = pa.Array.from_buffers(
+        pa.string(), 1, [None, pa.py_buffer(np.array([0, 1], dtype="<i4")), pa.py_buffer(b"\xff")]
+    )
+    with pytest.raises(pa.ArrowInvalid):
+        backend.extend_arrow(invalid)
+    assert len(backend) == 0
+
+
 # ---------------------------------------------------------------------------
 # Schema spec
 # ---------------------------------------------------------------------------