Fixes from three audits: safe path writes, XLSX metadata parsing, exact enums, and wide-header, pipeline and shared-string performance

.github/workflows/native-bindings.yml

@@ -219,7 +219,7 @@ jobs:
         run: chmod +x python/src/excelreader/_lib/*
 
       - name: Install the Python package
-        run: pip install -e "python[dev]" ${{ matrix.os == 'ubuntu-latest' && 'pytest-cov' || '' }}
+        run: pip install -e "python[dev]" ${{ matrix.os == 'ubuntu-latest' && 'pytest-cov pyarrow pandas polars' || '' }}
 
       - name: Run the Python tests
         run: >-

python/README.md

@@ -627,17 +627,17 @@ threads](#reading-a-large-sheet-on-several-threads).
 | `write_workbook()` → xlsb | 51.7 ms | 4.5 MB |
 | `write_workbook()` → xls | 59.2 ms | 13.4 MB |
 | `write_workbook()` → xlsx | 68.5 ms | 5.1 MB |
-| `write_workbook(polars_df)` → xlsx | 456.9 ms | 5.1 MB |
-| `write_workbook(pandas_df)` → xlsx | 465.5 ms | 5.1 MB |
-| `polars.DataFrame.write_excel()` | 4,663.7 ms | 5.6 MB |
-| `pandas.DataFrame.to_excel()` | 7,173.6 ms | 5.5 MB |
+| `write_workbook(polars_df)` → xlsx | 76.6 ms | 5.1 MB |
+| `write_workbook(pandas_df)` → xlsx | 82.5 ms | 5.1 MB |
+| `polars.DataFrame.write_excel()` | 4,714.9 ms | 5.6 MB |
+| `pandas.DataFrame.to_excel()` | 7,371.4 ms | 5.5 MB |
 
 The two DataFrame comparisons are matched work — same DataFrame in, xlsx out both times:
-`write_workbook(polars_df)` is ~10.2x faster than polars' own `write_excel()`, and
-`write_workbook(pandas_df)` ~15.4x faster than `to_excel()`. Both of ours pay a conversion the raw path does not: the DataFrame goes
-through Arrow and then a Python list before reaching the native columns, which is most of the gap
-between the ~460 ms rows and the 68.5 ms one. Handing `write_workbook()` buffers that are already
-columnar — what `parse()` returns — skips all of it.
+`write_workbook(polars_df)` is ~62x faster than polars' own `write_excel()`, and
+`write_workbook(pandas_df)` ~89x faster than `to_excel()`. Both of ours pay a conversion the raw path does not: the DataFrame goes
+through Arrow, and each column is cast to the layout the native side reads and copied out of Arrow's
+buffers, which is the gap between those rows and the 68.5 ms one. Handing `write_workbook()` buffers
+that are already columnar — what `parse()` returns — skips it.
 
 `write_workbook(xlsx)` at 68.5 ms lands within ~4 ms of the C++ binding's
 `xl::write_columns` on the same 14 columns (64.2 ms), which is the expected result: both are thin wrappers

python/benchmarks/bench_write.py

@@ -82,7 +82,7 @@ def benchmark_pandas(args: argparse.Namespace, tmp: Path):
     with excelreader.open_workbook(args.path, format="xlsb") as workbook:
         pandas_df = workbook.sheets[0].to_pandas(_FIXTURE_SCHEMA)
 
-    print("write_workbook(pandas) [pandas.DataFrame -> Arrow -> pylist -> native columns -> xlsx]:")
+    print("write_workbook(pandas) [pandas.DataFrame -> Arrow -> native columns -> xlsx]:")
     _time_and_assert(
         "  write_workbook(pandas)",
         lambda p: excelreader.write_workbook(p, pandas_df),
@@ -115,7 +115,7 @@ def benchmark_polars(args: argparse.Namespace, tmp: Path):
         schema = workbook.sheets[0].infer_schema(sample_size=10)
         polars_df = workbook.sheets[0].to_polars(schema)
 
-    print("write_workbook(polars) [polars.DataFrame -> Arrow -> pylist -> native columns -> xlsx]:")
+    print("write_workbook(polars) [polars.DataFrame -> Arrow -> native columns -> xlsx]:")
     _time_and_assert(
         "  write_workbook(polars)",
         lambda p: excelreader.write_workbook(p, polars_df),

python/src/excelreader/writer.py

@@ -10,7 +10,6 @@
 from __future__ import annotations
 
 import ctypes
-import datetime
 import io
 import os
 import sys
@@ -215,12 +214,17 @@ def _typed_table(data: Any, types: list[ColumnType] | None) -> tuple[TypedTable,
 
 def _write_all(destination: IO[bytes], data: bytes) -> None:
     """Calls `write()` until it has taken all of `data`. A `write()` that returns None is taken to
-    have written everything, as files that predate the `int` return convention do.
+    have written everything, as files that predate the `int` return convention do, except on a raw
+    stream, where None means the write would block and nothing was written.
     """
     view = memoryview(data)
     offset = 0
     while offset < len(data):
         written = destination.write(data if offset == 0 else view[offset:])
+        if written is None and isinstance(destination, io.RawIOBase):
+            raise BlockingIOError(
+                f"write() would block with {len(data) - offset} bytes left to write; use a blocking stream"
+            )
         if written is None:
             return
         if not 0 < written <= len(data) - offset:
@@ -283,18 +287,17 @@ def _write_memory(
     "timestamp[us]": ColumnType.TIMESTAMP,
 }
 
-_ARRAY_TYPECODES = {
-    ColumnType.I64: "q",
-    ColumnType.TIME: "q",
-    ColumnType.TIMESTAMP: "q",
-    ColumnType.F64: "d",
-    ColumnType.BOOL: "b",
-    ColumnType.DATE: "i",
+# The Arrow type each column is cast to before its values are read, and the array typecode of the
+# same layout: date32 is days, time64[us] and timestamp[us] are microseconds, as the native side takes them.
+_STORAGE = {
+    ColumnType.I64: ("int64", "q"),
+    ColumnType.TIME: ("int64", "q"),
+    ColumnType.TIMESTAMP: ("int64", "q"),
+    ColumnType.F64: ("float64", "d"),
+    ColumnType.BOOL: ("int8", "b"),
+    ColumnType.DATE: ("int32", "i"),
 }
 
-_EPOCH_DATE = datetime.date(1970, 1, 1)
-_EPOCH_DATETIME = datetime.datetime(1970, 1, 1, tzinfo=datetime.timezone.utc)
-
 
 def _arrow_column_type(field: Any) -> ColumnType:
     key = str(field.type)
@@ -307,68 +310,73 @@ def _arrow_column_type(field: Any) -> ColumnType:
         ) from None
 
 
-def _column_from_pylist(values: list[Any], column_type: ColumnType, row_count: int) -> tuple[Any, bytes | None]:
-    """Builds one column's native buffer layout from a plain Python list, as `to_pylist()` returns it.
-
-    Returns `(column, validity)`: `column` is a `StringColumn` for `ColumnType.STRING` (UTF-8 data plus
-    an `array("i")` of `row_count + 1` offsets), or an `array.array` of the typecode matching
-    `column_type` otherwise. `validity` is an Arrow-style LSB-first bitmap — `(row_count + 7) // 8`
-    bytes, bit `r` set when row `r` is valid — or `None` when no value in `values` was ever `None`,
-    matching `TypedTable.validity`'s "None means no nulls" convention. A `None` value contributes a
-    zero/empty placeholder to `column` and a cleared bit to the bitmap.
+def _slots(array_: Any, width: int, count: int, buffer_index: int = 1) -> memoryview:
+    """The bytes of `count` fixed-width slots of an Arrow array's buffer, from its slice offset on."""
+    buffer = array_.buffers()[buffer_index]
+    start = array_.offset * width
+    return memoryview(buffer)[start : start + count * width]
+
+
+def _validity_bitmap(column: Any) -> bytes:
+    """An LSB-first bitmap of `(len + 7) // 8` bytes, bit `r` set when row `r` is valid."""
+    import pyarrow.compute
+
+    valid = pyarrow.compute.is_valid(column)
+    row_count = len(column)
+    if valid.offset % 8:
+        bits = bytearray((row_count + 7) // 8)
+        for row, flag in enumerate(valid.to_pylist()):
+            bits[row // 8] |= flag << (row % 8)
+        return bytes(bits)
+    start = valid.offset // 8
+    bits = bytearray(memoryview(valid.buffers()[1])[start : start + (row_count + 7) // 8])
+    if row_count % 8:
+        bits[-1] &= (1 << (row_count % 8)) - 1
+    return bytes(bits)
+
+
+def _string_column(strings: Any) -> StringColumn:
+    """UTF-8 data plus `len + 1` offsets starting at 0, from an Arrow `string` array without nulls."""
+    row_count = len(strings)
+    if row_count == 0:
+        return StringColumn(array("i", [0]), b"")
+    offsets = array("i")
+    offsets.frombytes(_slots(strings, 4, row_count + 1))
+    first, last = offsets[0], offsets[-1]
+    data_buffer = strings.buffers()[2]
+    data = b"" if data_buffer is None else bytes(memoryview(data_buffer)[first:last])
+    if first:
+        offsets = array("i", [offset - first for offset in offsets])
+    return StringColumn(offsets, data)
+
+
+def _column_from_arrow(column: Any, column_type: ColumnType) -> tuple[Any, bytes | None]:
+    """Builds one column's native buffer layout and its validity bitmap, `None` when nothing is null.
+
+    Arrow first casts the column to the one layout this reads, and replaces nulls with an empty or
+    zero placeholder, so the producer's choices of string variant, null slots or bit packing never
+    reach the buffer reads.
     """
-    validity = bytearray((row_count + 7) // 8)
-    has_null = False
+    import pyarrow
+    import pyarrow.compute
 
+    validity = _validity_bitmap(column) if column.null_count else None
     if column_type is ColumnType.STRING:
-        offsets = array("i", [0])
-        chunks: list[bytes] = []
-        total = 0
-        for row, value in enumerate(values):
-            if value is None:
-                has_null = True
-            else:
-                validity[row // 8] |= 1 << (row % 8)
-                encoded = value.encode("utf-8")
-                chunks.append(encoded)
-                total += len(encoded)
-            offsets.append(total)
-        string_column: Any = StringColumn(offsets, b"".join(chunks))
-        return string_column, (bytes(validity) if has_null else None)
-
-    data_array = array(_ARRAY_TYPECODES[column_type])
-    for row, value in enumerate(values):
-        if value is None:
-            has_null = True
-            data_array.append(0)
-            continue
-        validity[row // 8] |= 1 << (row % 8)
-        if column_type is ColumnType.DATE:
-            data_array.append((value - _EPOCH_DATE).days)
-        elif column_type is ColumnType.TIME:
-            data_array.append(
-                value.hour * 3_600_000_000
-                + value.minute * 60_000_000
-                + value.second * 1_000_000
-                + value.microsecond
-            )
-        elif column_type is ColumnType.TIMESTAMP:
-            delta = value.replace(tzinfo=datetime.timezone.utc) - _EPOCH_DATETIME
-            data_array.append(delta.days * 86_400_000_000 + delta.seconds * 1_000_000 + delta.microseconds)
-        elif column_type is ColumnType.BOOL:
-            data_array.append(1 if value else 0)
-        else:
-            data_array.append(value)
-
-    return data_array, (bytes(validity) if has_null else None)
+        strings = column.cast(pyarrow.string())
+        if strings.null_count:
+            strings = pyarrow.compute.fill_null(strings, "")
+        return _string_column(strings), validity
+    storage, typecode = _STORAGE[column_type]
+    values = column.cast(storage)
+    if values.null_count:
+        values = pyarrow.compute.fill_null(values, 0)
+    data = array(typecode)
+    if len(values):
+        data.frombytes(_slots(values, data.itemsize, len(values)))
+    return data, validity
 
 
 def _table_from_arrow(batch: Any) -> tuple[TypedTable, list[ColumnType]]:
-    """Converts through Python lists rather than borrowing Arrow's buffers directly: Arrow's null
-    representation, offset conventions and chunking are the producer's choice, and reproducing all of
-    them faithfully is exactly the parser this package deliberately does not have. The conversion is
-    one pass and keeps the native side reading only layouts it produced itself.
-    """
     import pyarrow
 
     if isinstance(batch, pyarrow.Table):
@@ -379,7 +387,7 @@ def _table_from_arrow(batch: Any) -> tuple[TypedTable, list[ColumnType]]:
     columns: list[Any] = []
     validity: list[bytes | None] = []
     for index, column in enumerate(batch.columns):
-        built, mask = _column_from_pylist(column.to_pylist(), types[index], batch.num_rows)
+        built, mask = _column_from_arrow(column, types[index])
         columns.append(built)
         validity.append(mask)
 

python/tests/test_arrow_conversion.py

@@ -0,0 +1,132 @@
+import datetime
+
+import pytest
+
+pa = pytest.importorskip("pyarrow")
+
+from excelreader.types import ColumnType, StringColumn
+from excelreader.writer import _table_from_arrow
+
+_EPOCH_DATE = datetime.date(1970, 1, 1)
+_EPOCH = datetime.datetime(1970, 1, 1)
+
+
+def _decode(column, column_type, validity, row_count):
+    """Reads one converted column back into Python values, None where the validity bit is clear."""
+    values = []
+    for row in range(row_count):
+        if validity is not None and not validity[row // 8] >> (row % 8) & 1:
+            values.append(None)
+            continue
+        if column_type is ColumnType.STRING:
+            values.append(column.data[column.offsets[row] : column.offsets[row + 1]].decode("utf-8"))
+            continue
+        raw = column[row]
+        if column_type is ColumnType.DATE:
+            values.append(_EPOCH_DATE + datetime.timedelta(days=raw))
+        elif column_type is ColumnType.TIME:
+            values.append((datetime.datetime.min + datetime.timedelta(microseconds=raw)).time())
+        elif column_type is ColumnType.TIMESTAMP:
+            values.append(_EPOCH + datetime.timedelta(microseconds=raw))
+        elif column_type is ColumnType.BOOL:
+            values.append(bool(raw))
+        else:
+            values.append(raw)
+    return values
+
+
+def _assert_round_trips(data, expected_types):
+    table, types = _table_from_arrow(data)
+    batch = data.combine_chunks() if isinstance(data, pa.Table) else data
+    assert types == expected_types
+    assert table.names == list(batch.schema.names)
+    for index, column_type in enumerate(types):
+        column, validity = table.columns[index], table.validity[index]
+        expected = batch.column(index).to_pylist()
+        assert _decode(column, column_type, validity, table.row_count) == expected
+        assert (validity is None) == (None not in expected)
+        if validity is not None:
+            assert len(validity) == (table.row_count + 7) // 8
+            spare = table.row_count % 8
+            assert spare == 0 or validity[-1] >> spare == 0
+        if column_type is ColumnType.STRING:
+            assert isinstance(column, StringColumn)
+            assert column.offsets[0] == 0 and column.offsets[-1] == len(column.data)
+    return table
+
+
+_TYPED = pa.record_batch(
+    {
+        "text": pa.array(["a", None, "ção ✓", "", "last"]),
+        "large": pa.array(["x", "y", None, "zz", ""], pa.large_string()),
+        "view": pa.array(["v1", None, "v3", "v4", "v5"], pa.string_view()),
+        "i64": pa.array([1, None, -(2**63), 2**63 - 1, 0], pa.int64()),
+        "i32": pa.array([7, -7, None, 2**31 - 1, 0], pa.int32()),
+        "f64": pa.array([1.5, None, -0.0, 1e308, float("inf")], pa.float64()),
+        "f32": pa.array([0.25, 2.5, None, -1.0, 3.0], pa.float32()),
+        "flag": pa.array([True, False, None, True, False]),
+        "day": pa.array([datetime.date(2020, 1, 2), None, datetime.date(1900, 3, 1), _EPOCH_DATE, datetime.date(9999, 12, 31)]),
+        "clock": pa.array([datetime.time(0, 0), datetime.time(23, 59, 59, 999999), None, datetime.time(12, 30), datetime.time(1, 2, 3)]),
+        "stamp": pa.array(
+            [datetime.datetime(2024, 5, 6, 12, 30), None, datetime.datetime(1969, 12, 31, 23, 59, 59, 1), _EPOCH, datetime.datetime(2000, 1, 1)],
+            pa.timestamp("us"),
+        ),
+    }
+)
+_TYPES = [
+    ColumnType.STRING, ColumnType.STRING, ColumnType.STRING,
+    ColumnType.I64, ColumnType.I64, ColumnType.F64, ColumnType.F64,
+    ColumnType.BOOL, ColumnType.DATE, ColumnType.TIME, ColumnType.TIMESTAMP,
+]
+
+
+def test_every_supported_type_with_nulls_converts_value_for_value():
+    _assert_round_trips(_TYPED, _TYPES)
+
+
+@pytest.mark.parametrize(("offset", "length"), [(1, 3), (2, 3), (3, 2), (0, 0), (5, 0)])
+def test_a_sliced_batch_converts_only_its_own_rows(offset, length):
+    table = _assert_round_trips(_TYPED.slice(offset, length), _TYPES)
+    assert table.row_count == length
+
+
+def test_a_chunked_table_converts_as_one_batch():
+    chunked = pa.Table.from_batches([_TYPED.slice(0, 2), _TYPED.slice(2, 3), _TYPED.slice(0, 0)])
+    assert chunked.column(0).num_chunks > 1
+    _assert_round_trips(chunked, _TYPES)
+
+
+def test_an_empty_table_converts_to_zero_rows():
+    table = _assert_round_trips(_TYPED.schema.empty_table(), _TYPES)
+    assert table.row_count == 0
+
+
+def test_columns_without_nulls_carry_no_validity_bitmap():
+    batch = pa.record_batch({"n": pa.array(range(20), pa.int64()), "s": pa.array([str(i) for i in range(20)])})
+    table = _assert_round_trips(batch, [ColumnType.I64, ColumnType.STRING])
+    assert table.validity == [None, None]
+
+
+def test_an_all_null_column_and_row_counts_past_a_byte_boundary_convert():
+    batch = pa.record_batch(
+        {
+            "s": pa.array([None] * 13, pa.string()),
+            "b": pa.array([None if i % 3 == 0 else i % 2 == 0 for i in range(13)]),
+            "f": pa.array([None if i % 4 == 0 else i / 3 for i in range(13)]),
+        }
+    )
+    _assert_round_trips(batch, [ColumnType.STRING, ColumnType.BOOL, ColumnType.F64])
+
+
+def test_a_validity_mask_off_a_byte_boundary_converts(monkeypatch):
+    import pyarrow.compute
+
+    is_valid = pyarrow.compute.is_valid
+    monkeypatch.setattr(pyarrow.compute, "is_valid", lambda column: pa.concat_arrays([pa.array([False] * 3), is_valid(column)]).slice(3))
+    batch = pa.record_batch({"n": pa.array([None if i % 3 == 0 else i for i in range(11)], pa.int64())})
+    _assert_round_trips(batch, [ColumnType.I64])
+
+
+def test_an_unsupported_arrow_type_is_rejected():
+    with pytest.raises(ValueError, match="int8"):
+        _table_from_arrow(pa.record_batch({"small": pa.array([1], pa.int8())}))