feat(python): read manifests and manifest lists as pyarrow RecordBatchReaders

bindings/python/src/manifest.rs

@@ -17,10 +17,17 @@
 
 use std::sync::Arc;
 
+use arrow::array::RecordBatch;
+use arrow::datatypes::SchemaRef;
+use arrow::error::ArrowError;
+use arrow::pyarrow::IntoPyArrow;
+use arrow::record_batch::{RecordBatchIterator, RecordBatchReader};
 use iceberg::spec::{
-    FieldSummary, FormatVersion, Manifest, ManifestEntry, ManifestFile, ManifestList,
+    FieldSummary, FormatVersion, Manifest, ManifestContentType, ManifestEntry, ManifestFile,
+    ManifestList,
 };
 use pyo3::prelude::*;
+use pyo3::pybacked::PyBackedBytes;
 
 use crate::data_file::PyDataFile;
 use crate::error::to_py_err;
@@ -226,11 +233,98 @@ pub fn read_manifest_list(bs: &[u8]) -> PyResult<PyManifestList> {
     })
 }
 
+/// Reads the entries of a manifest as a `pyarrow.RecordBatchReader`.
+///
+/// The other arguments are the manifest's entry in the manifest list, which
+/// the entries inherit snapshot IDs, sequence numbers, and first row IDs from.
+/// Raises `ValueError` for an invalid manifest header. An error while reading
+/// raises `pyarrow.ArrowInvalid`, which subclasses `ValueError`. The reader
+/// holds a reference to `bs` until it's closed or collected, and the next call
+/// into `pyiceberg_core` releases it.
+#[pyfunction]
+#[pyo3(signature = (
+    bs,
+    *,
+    manifest_path,
+    manifest_length,
+    partition_spec_id,
+    content,
+    sequence_number,
+    min_sequence_number,
+    added_snapshot_id,
+    first_row_id = None,
+))]
+#[allow(clippy::too_many_arguments)]
+pub fn read_manifest_entries_as_arrow(
+    py: Python<'_>,
+    bs: PyBackedBytes,
+    manifest_path: String,
+    manifest_length: i64,
+    partition_spec_id: i32,
+    content: i32,
+    sequence_number: i64,
+    min_sequence_number: i64,
+    added_snapshot_id: i64,
+    first_row_id: Option<u64>,
+) -> PyResult<Bound<'_, PyAny>> {
+    // Inheritance reads none of the optional fields left as None.
+    let manifest_file = ManifestFile {
+        manifest_path,
+        manifest_length,
+        partition_spec_id,
+        content: ManifestContentType::try_from(content).map_err(to_py_err)?,
+        sequence_number,
+        min_sequence_number,
+        added_snapshot_id,
+        added_files_count: None,
+        existing_files_count: None,
+        deleted_files_count: None,
+        added_rows_count: None,
+        existing_rows_count: None,
+        deleted_rows_count: None,
+        partitions: None,
+        key_metadata: None,
+        first_row_id,
+    };
+    let reader = iceberg::arrow::read_manifest_entries(bs, &manifest_file).map_err(to_py_err)?;
+    into_pyarrow_reader(py, reader.schema(), reader)
+}
+
+/// Reads the entries of a manifest list as a `pyarrow.RecordBatchReader`.
+///
+/// Raises `ValueError` for an invalid manifest list header. An error while
+/// reading raises `pyarrow.ArrowInvalid`, which subclasses `ValueError`. The
+/// reader holds a reference to `bs` until it's closed or collected, and the
+/// next call into `pyiceberg_core` releases it.
+#[pyfunction]
+pub fn read_manifest_list_as_arrow(
+    py: Python<'_>,
+    bs: PyBackedBytes,
+) -> PyResult<Bound<'_, PyAny>> {
+    let reader = iceberg::arrow::read_manifest_list(bs).map_err(to_py_err)?;
+    into_pyarrow_reader(py, reader.schema(), reader)
+}
+
+/// Exports `batches` through the Arrow C stream interface.
+fn into_pyarrow_reader<'py>(
+    py: Python<'py>,
+    schema: SchemaRef,
+    batches: impl Iterator<Item = iceberg::Result<RecordBatch>> + Send + 'static,
+) -> PyResult<Bound<'py, PyAny>> {
+    let batches =
+        batches.map(|batch| batch.map_err(|err| ArrowError::ExternalError(Box::new(err))));
+    let reader: Box<dyn RecordBatchReader + Send> =
+        Box::new(RecordBatchIterator::new(batches, schema));
+    reader.into_pyarrow(py)
+}
+
 pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     let this = PyModule::new(py, "manifest")?;
 
     this.add_function(wrap_pyfunction!(read_manifest_entries, &this)?)?;
     this.add_function(wrap_pyfunction!(read_manifest_list, &this)?)?;
+    this.add_function(wrap_pyfunction!(read_manifest_entries_as_arrow, &this)?)?;
+    this.add_function(wrap_pyfunction!(read_manifest_list_as_arrow, &this)?)?;
 
     m.add_submodule(&this)?;
     py.import("sys")?

bindings/python/tests/test_manifest.py

@@ -15,22 +15,33 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import struct
+import sys
 from pathlib import Path
 from typing import Any
 
+import pyarrow as pa
 import pytest
 from pyiceberg.io import FileIO
 from pyiceberg.io.pyarrow import PyArrowFileIO
 from pyiceberg.manifest import (
     DataFile,
     DataFileContent,
     FileFormat,
+    ManifestContent,
     ManifestEntry,
     ManifestEntryStatus,
     ManifestFile,
     _manifests,
+    read_manifest_list,
+    write_manifest,
     write_manifest_list,
 )
+from pyiceberg.partitioning import PartitionField, PartitionSpec
+from pyiceberg.schema import Schema
+from pyiceberg.transforms import IdentityTransform
+from pyiceberg.typedef import Record
+from pyiceberg.types import LongType, NestedField, StringType
 
 
 @pytest.fixture(autouse=True)
@@ -235,3 +246,244 @@ def test_manifest_file_without_partition_summaries(tmp_path: Path) -> None:
     manifest_file = manifest.read_manifest_list(bs).entries()[0]
 
     assert manifest_file.partitions is None
+
+
+SNAPSHOT_ID = 25
+
+
+def _write_manifest(tmp_path: Path, num_files: int) -> tuple[bytes, ManifestFile]:
+    """Writes a v2 manifest whose entries inherit their sequence numbers."""
+    schema = Schema(
+        NestedField(1, "id", LongType(), required=False),
+        NestedField(2, "name", StringType(), required=False),
+    )
+    spec = PartitionSpec(
+        PartitionField(
+            source_id=1, field_id=1000, transform=IdentityTransform(), name="id"
+        ),
+        spec_id=0,
+    )
+    path = str(tmp_path / "manifest.avro")
+    io = PyArrowFileIO()
+    with write_manifest(
+        format_version=2,
+        spec=spec,
+        schema=schema,
+        output_file=io.new_output(path),
+        snapshot_id=SNAPSHOT_ID,
+        avro_compression="deflate",
+    ) as writer:
+        for i in range(num_files):
+            writer.add(
+                ManifestEntry.from_args(
+                    status=ManifestEntryStatus.ADDED,
+                    snapshot_id=SNAPSHOT_ID,
+                    data_file=DataFile.from_args(
+                        content=DataFileContent.DATA,
+                        file_path=f"s3://bucket/data/{i}.parquet",
+                        file_format=FileFormat.PARQUET,
+                        partition=Record(i),
+                        record_count=10 + i,
+                        file_size_in_bytes=100,
+                        column_sizes={1: 8, 2: 16},
+                        value_counts={1: 10 + i, 2: 10 + i},
+                        null_value_counts={1: 0, 2: i},
+                        nan_value_counts={},
+                        lower_bounds={1: struct.pack("<q", i)},
+                        upper_bounds={1: struct.pack("<q", i)},
+                        split_offsets=[4],
+                        sort_order_id=0,
+                    ),
+                )
+            )
+    manifest_file = ManifestFile.from_args(
+        manifest_path=path,
+        manifest_length=writer.to_manifest_file().manifest_length,
+        partition_spec_id=0,
+        content=ManifestContent.DATA,
+        sequence_number=5,
+        min_sequence_number=5,
+        added_snapshot_id=SNAPSHOT_ID,
+    )
+    return io.new_input(path).open().read(), manifest_file
+
+
+def _read_entries(
+    bs: bytes, manifest_file: ManifestFile, first_row_id: int | None = None
+) -> pa.RecordBatchReader:
+    from pyiceberg_core import manifest
+
+    return manifest.read_manifest_entries_as_arrow(
+        bs,
+        manifest_path=manifest_file.manifest_path,
+        manifest_length=manifest_file.manifest_length,
+        partition_spec_id=manifest_file.partition_spec_id,
+        content=manifest_file.content.value,
+        sequence_number=manifest_file.sequence_number,
+        min_sequence_number=manifest_file.min_sequence_number,
+        added_snapshot_id=manifest_file.added_snapshot_id,
+        first_row_id=first_row_id,
+    )
+
+
+def test_read_manifest_entries_as_arrow_matches_pyiceberg(tmp_path: Path) -> None:
+    bs, manifest_file = _write_manifest(tmp_path, 3)
+    reader = _read_entries(bs, manifest_file)
+    assert isinstance(reader, pa.RecordBatchReader)
+
+    rows = reader.read_all().to_pylist()
+    expected = manifest_file.fetch_manifest_entry(
+        PyArrowFileIO(), discard_deleted=False
+    )
+    assert [
+        {
+            "status": row["status"],
+            "snapshot_id": row["snapshot_id"],
+            "sequence_number": row["sequence_number"],
+            "file_sequence_number": row["file_sequence_number"],
+            "file_path": row["data_file"]["file_path"],
+            "partition": row["data_file"]["partition"]["id"],
+            "record_count": row["data_file"]["record_count"],
+            "null_value_counts": dict(row["data_file"]["null_value_counts"]),
+            "lower_bounds": dict(row["data_file"]["lower_bounds"]),
+        }
+        for row in rows
+    ] == [
+        {
+            "status": entry.status.value,
+            "snapshot_id": entry.snapshot_id,
+            "sequence_number": entry.sequence_number,
+            "file_sequence_number": entry.file_sequence_number,
+            "file_path": entry.data_file.file_path,
+            "partition": entry.data_file.partition[0],
+            "record_count": entry.data_file.record_count,
+            "null_value_counts": entry.data_file.null_value_counts,
+            "lower_bounds": entry.data_file.lower_bounds,
+        }
+        for entry in expected
+    ]
+
+
+def test_read_manifest_entries_as_arrow_yields_batches_of_1024_rows(
+    tmp_path: Path,
+) -> None:
+    bs, manifest_file = _write_manifest(tmp_path, 2500)
+    assert [batch.num_rows for batch in _read_entries(bs, manifest_file)] == [
+        1024,
+        1024,
+        452,
+    ]
+
+
+def test_read_manifest_entries_as_arrow_releases_input_on_close(
+    tmp_path: Path,
+) -> None:
+    from pyiceberg_core import transform
+
+    bs, manifest_file = _write_manifest(tmp_path, 3)
+    refcount = sys.getrefcount(bs)
+    reader = _read_entries(bs, manifest_file)
+    assert sys.getrefcount(bs) == refcount + 1
+
+    # pyarrow releases the stream without PyO3 attached to the interpreter, so
+    # PyO3 defers the decref to its next call.
+    reader.close()
+    transform.identity(pa.array([1]))
+    assert sys.getrefcount(bs) == refcount
+
+
+def test_read_manifest_entries_as_arrow_batches_outlive_reader(
+    tmp_path: Path,
+) -> None:
+    bs, manifest_file = _write_manifest(tmp_path, 3)
+    reader = _read_entries(bs, manifest_file)
+    batch = reader.read_next_batch()
+    reader.close()
+    del reader, bs
+
+    assert batch.column("data_file").field("file_path").to_pylist() == [
+        f"s3://bucket/data/{i}.parquet" for i in range(3)
+    ]
+
+
+def test_read_manifest_entries_as_arrow_raises_while_reading(tmp_path: Path) -> None:
+    bs, manifest_file = _write_manifest(tmp_path, 3)
+    # Assigning first row IDs to the second entry overflows i64.
+    reader = _read_entries(bs, manifest_file, first_row_id=2**63 - 15)
+
+    with pytest.raises(pa.ArrowInvalid, match="row ID overflow"):
+        reader.read_next_batch()
+
+
+@pytest.mark.parametrize(
+    "payload", [b"", b"not-an-avro-file", b"Obj\x01truncated-after-the-avro-magic"]
+)
+def test_read_manifest_entries_as_arrow_raises_on_invalid_avro(
+    tmp_path: Path, payload: bytes
+) -> None:
+    _, manifest_file = _write_manifest(tmp_path, 1)
+    with pytest.raises(ValueError):
+        _read_entries(payload, manifest_file)
+
+
+@pytest.mark.parametrize("format_version", [1, 2])
+def test_read_manifest_list_as_arrow_matches_pyiceberg(
+    tmp_path: Path, format_version: int
+) -> None:
+    from pyiceberg_core import manifest
+
+    path = str(tmp_path / "snap.avro")
+    io = PyArrowFileIO()
+    with write_manifest_list(
+        format_version=format_version,
+        output_file=io.new_output(path),
+        snapshot_id=SNAPSHOT_ID,
+        parent_snapshot_id=None,
+        sequence_number=7,
+        avro_compression="deflate",
+    ) as writer:
+        writer.add_manifests(
+            [
+                ManifestFile.from_args(
+                    manifest_path=f"s3://bucket/metadata/{i}.avro",
+                    manifest_length=1024,
+                    partition_spec_id=0,
+                    added_snapshot_id=SNAPSHOT_ID,
+                    sequence_number=-1,
+                    min_sequence_number=-1,
+                    added_files_count=i,
+                    existing_files_count=0,
+                    deleted_files_count=0,
+                    added_rows_count=10 * i,
+                    existing_rows_count=0,
+                    deleted_rows_count=0,
+                )
+                for i in range(2)
+            ]
+        )
+
+    reader = manifest.read_manifest_list_as_arrow(io.new_input(path).open().read())
+    assert isinstance(reader, pa.RecordBatchReader)
+    rows = reader.read_all().to_pylist()
+    expected = list(read_manifest_list(io.new_input(path)))
+    fields = [
+        "manifest_path",
+        "sequence_number",
+        "min_sequence_number",
+        "added_snapshot_id",
+        "added_files_count",
+        "added_rows_count",
+    ]
+    assert [{f: row[f] for f in fields} for row in rows] == [
+        {f: getattr(manifest_file, f) for f in fields} for manifest_file in expected
+    ]
+
+
+@pytest.mark.parametrize(
+    "payload", [b"", b"not-an-avro-file", b"Obj\x01truncated-after-the-avro-magic"]
+)
+def test_read_manifest_list_as_arrow_raises_on_invalid_avro(payload: bytes) -> None:
+    from pyiceberg_core import manifest
+
+    with pytest.raises(ValueError):
+        manifest.read_manifest_list_as_arrow(payload)