From 9200eec6a24fa044c89b34463af976b1c2589ccf Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 14 Apr 2026 20:51:34 +0000 Subject: [PATCH 1/5] feat(io): add ParquetFile with row-group statistics() Expose per-row-group Parquet statistics (min/max/null_count) as an Arrow RecordBatch, via a new `arro3.io.ParquetFile` class. The class loads only the Parquet footer (no data pages) via `ArrowReaderMetadata::load` and wraps `arrow-rs`'s `StatisticsConverter`. Implements the design sketched in kylebarron/arro3#315. Struct columns are not yet supported upstream (apache/arrow-rs#7364). --- arro3-io/python/arro3/io/_parquet.pyi | 75 +++++++++++++++++++ arro3-io/src/lib.rs | 1 + arro3-io/src/parquet.rs | 101 +++++++++++++++++++++++++- docs/api/io/parquet.md | 1 + tests/io/test_parquet.py | 100 ++++++++++++++++++++++++- 5 files changed, 273 insertions(+), 5 deletions(-) diff --git a/arro3-io/python/arro3/io/_parquet.pyi b/arro3-io/python/arro3/io/_parquet.pyi index 6f2bc61..184bef9 100644 --- a/arro3-io/python/arro3/io/_parquet.pyi +++ b/arro3-io/python/arro3/io/_parquet.pyi @@ -31,6 +31,81 @@ ParquetEncoding = Literal[ ] """Allowed Parquet encodings.""" +class ParquetFile: + """A Parquet file opened for metadata inspection. + + This loads only the Parquet footer (no data pages) and exposes row-group + level metadata, including per-row-group statistics via + [`statistics`][arro3.io.ParquetFile.statistics]. + """ + + @classmethod + def open( + cls, + file: IO[bytes] | Path | str, + *, + skip_arrow_metadata: bool = False, + page_index: bool = False, + ) -> "ParquetFile": + """Open a Parquet file and read its footer metadata. + + Args: + file: The input Parquet file path or binary buffer. + + Keyword Args: + skip_arrow_metadata: If `True`, do not decode the embedded Arrow + schema stored in the Parquet key-value metadata. Defaults to + `False`. + page_index: If `True`, load the Parquet page index as part of the + footer read. Defaults to `False`. + + Returns: + A new `ParquetFile` wrapping the file's footer metadata. + """ + + @property + def schema(self) -> core.Schema: + """The Arrow schema of this Parquet file.""" + + @property + def num_rows(self) -> int: + """The total number of rows in this Parquet file.""" + + @property + def num_row_groups(self) -> int: + """The number of row groups in this Parquet file.""" + + @property + def num_columns(self) -> int: + """The number of columns in this Parquet file.""" + + def statistics( + self, + column_name: str, + *, + missing_null_counts_as_zero: bool = True, + ) -> core.RecordBatch: + """Row-group statistics for a single column. + + Args: + column_name: The name of the column to read statistics for. + + Keyword Args: + missing_null_counts_as_zero: If `True`, row groups that do not + report a null count in their statistics are reported as zero. + If `False`, a `null` value is emitted instead. Defaults to + `True`. + + Returns: + A `RecordBatch` with one row per row group and three columns: + `min`, `max` and `null_count`. `min` and `max` have the Arrow + data type of the source column; `null_count` is a `UInt64` + column. Row order matches the Parquet row-group order. + + Note: + Struct columns are not yet supported (see apache/arrow-rs#7364). + """ + def read_parquet(file: IO[bytes] | Path | str) -> core.RecordBatchReader: """Read a Parquet file to an Arrow RecordBatchReader diff --git a/arro3-io/src/lib.rs b/arro3-io/src/lib.rs index 2a33cc2..b6afc02 100644 --- a/arro3-io/src/lib.rs +++ b/arro3-io/src/lib.rs @@ -60,6 +60,7 @@ fn _io(py: Python, m: &Bound) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(parquet::read_parquet))?; m.add_wrapped(wrap_pyfunction!(parquet::read_parquet_async))?; m.add_wrapped(wrap_pyfunction!(parquet::write_parquet))?; + m.add_class::()?; Ok(()) } diff --git a/arro3-io/src/parquet.rs b/arro3-io/src/parquet.rs index 13bd1fc..fe25f73 100644 --- a/arro3-io/src/parquet.rs +++ b/arro3-io/src/parquet.rs @@ -2,20 +2,25 @@ use std::collections::HashMap; use std::str::FromStr; use std::sync::Arc; -use arrow_array::{RecordBatchIterator, RecordBatchReader}; -use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use arrow_array::{Array, RecordBatch, RecordBatchIterator, RecordBatchReader}; +use arrow_schema::{Field, Schema}; +use parquet::arrow::arrow_reader::statistics::StatisticsConverter; +use parquet::arrow::arrow_reader::{ + ArrowReaderMetadata, ArrowReaderOptions, ParquetRecordBatchReaderBuilder, +}; use parquet::arrow::arrow_writer::ArrowWriterOptions; use parquet::arrow::async_reader::ParquetObjectReader; use parquet::arrow::ArrowWriter; use parquet::basic::{Compression, Encoding}; -use parquet::file::metadata::KeyValue; +use parquet::file::metadata::{KeyValue, PageIndexPolicy}; use parquet::file::properties::{WriterProperties, WriterVersion}; use parquet::schema::types::ColumnPath; use pyo3::exceptions::{PyTypeError, PyValueError}; use pyo3::prelude::*; use pyo3::pybacked::PyBackedStr; +use pyo3::types::PyType; use pyo3_arrow::error::PyArrowResult; -use pyo3_arrow::export::Arro3RecordBatchReader; +use pyo3_arrow::export::{Arro3RecordBatch, Arro3RecordBatchReader, Arro3Schema}; use pyo3_arrow::input::AnyRecordBatch; use pyo3_arrow::{PyRecordBatchReader, PyTable}; use pyo3_object_store::PyObjectStore; @@ -78,6 +83,94 @@ async fn read_parquet_async_inner( Ok(PyTable::try_new(batches, arrow_schema)?) } +/// A Parquet file opened for metadata inspection. +/// +/// This loads only the Parquet footer (no data pages) and exposes row-group +/// level metadata, including per-row-group statistics via [`PyParquetFile::statistics`]. +#[pyclass(module = "arro3.io", name = "ParquetFile", frozen)] +pub(crate) struct PyParquetFile { + meta: ArrowReaderMetadata, +} + +#[pymethods] +impl PyParquetFile { + /// Open a Parquet file and read its footer metadata. + #[classmethod] + #[pyo3(signature = (file, *, skip_arrow_metadata = false, page_index = false))] + fn open( + _cls: &Bound, + mut file: FileReader, + skip_arrow_metadata: bool, + page_index: bool, + ) -> Arro3IoResult { + let options = ArrowReaderOptions::new() + .with_skip_arrow_metadata(skip_arrow_metadata) + .with_page_index_policy(PageIndexPolicy::from(page_index)); + let meta = ArrowReaderMetadata::load(&mut file, options)?; + Ok(Self { meta }) + } + + /// The Arrow schema of this Parquet file. + #[getter] + fn schema(&self) -> Arro3Schema { + self.meta.schema().clone().into() + } + + /// The total number of rows in this Parquet file. + #[getter] + fn num_rows(&self) -> i64 { + self.meta.metadata().file_metadata().num_rows() + } + + /// The number of row groups in this Parquet file. + #[getter] + fn num_row_groups(&self) -> usize { + self.meta.metadata().num_row_groups() + } + + /// The number of columns in this Parquet file. + #[getter] + fn num_columns(&self) -> usize { + self.meta.schema().fields().len() + } + + /// Row-group statistics for a single column. + /// + /// Returns a `RecordBatch` with one row per row group and the columns + /// `min`, `max` and `null_count`. `min` and `max` take the Arrow data + /// type of the source column; `null_count` is a `UInt64` column. + /// + /// Note: struct columns are not yet supported upstream + /// (see apache/arrow-rs#7364). + #[pyo3(signature = (column_name, *, missing_null_counts_as_zero = true))] + fn statistics( + &self, + column_name: &str, + missing_null_counts_as_zero: bool, + ) -> Arro3IoResult { + let parquet_meta = self.meta.metadata(); + let converter = StatisticsConverter::try_new( + column_name, + self.meta.schema(), + self.meta.parquet_schema(), + )? + .with_missing_null_counts_as_zero(missing_null_counts_as_zero); + + let min_values = converter.row_group_mins(parquet_meta.row_groups())?; + let max_values = converter.row_group_maxes(parquet_meta.row_groups())?; + let null_counts = converter.row_group_null_counts(parquet_meta.row_groups())?; + + let schema = Arc::new(Schema::new(vec![ + Field::new("min", min_values.data_type().clone(), true), + Field::new("max", max_values.data_type().clone(), true), + Field::new("null_count", null_counts.data_type().clone(), true), + ])); + let batch = + RecordBatch::try_new(schema, vec![min_values, max_values, Arc::new(null_counts)])?; + Ok(batch.into()) + } +} + pub(crate) struct PyWriterVersion(WriterVersion); impl<'py> FromPyObject<'_, 'py> for PyWriterVersion { diff --git a/docs/api/io/parquet.md b/docs/api/io/parquet.md index edd33f1..478d58d 100644 --- a/docs/api/io/parquet.md +++ b/docs/api/io/parquet.md @@ -1,5 +1,6 @@ # Parquet +::: arro3.io.ParquetFile ::: arro3.io.read_parquet ::: arro3.io.read_parquet_async ::: arro3.io.write_parquet diff --git a/tests/io/test_parquet.py b/tests/io/test_parquet.py index 3daa824..c5477e5 100644 --- a/tests/io/test_parquet.py +++ b/tests/io/test_parquet.py @@ -4,7 +4,8 @@ import pyarrow as pa import pyarrow.parquet as pq -from arro3.io import read_parquet, write_parquet +import pytest +from arro3.io import ParquetFile, read_parquet, write_parquet def test_parquet_round_trip(): @@ -51,3 +52,100 @@ def test_copy_parquet_kv_metadata(): reader = read_parquet(pq_path) assert reader.schema.metadata[b"hello"] == b"world" + + +def _write_multi_row_group(path: Path, table: pa.Table, rows_per_group: int) -> None: + write_parquet(table, path, max_row_group_size=rows_per_group) + + +def test_parquet_file_accessors(): + table = pa.table({"a": list(range(10)), "b": [float(i) for i in range(10)]}) + with TemporaryDirectory() as tmp_path: + pq_path = Path(tmp_path) / "test.parquet" + _write_multi_row_group(pq_path, table, rows_per_group=4) + + pf = ParquetFile.open(pq_path) + assert pf.num_rows == 10 + assert pf.num_row_groups == 3 + assert pf.num_columns == 2 + arrow_schema = pa.schema(pf.schema) + assert arrow_schema.names == ["a", "b"] + + +def test_parquet_file_statistics_int(): + table = pa.table({"a": list(range(10))}) + with TemporaryDirectory() as tmp_path: + pq_path = Path(tmp_path) / "test.parquet" + _write_multi_row_group(pq_path, table, rows_per_group=4) + + pf = ParquetFile.open(pq_path) + stats = pa.record_batch(pf.statistics("a")) + assert stats.schema.names == ["min", "max", "null_count"] + assert stats.num_rows == pf.num_row_groups + + # Ground-truth: pyarrow footer metadata + pa_pf = pq.ParquetFile(pq_path) + col_idx = pa_pf.schema_arrow.get_field_index("a") + for rg_idx in range(pa_pf.num_row_groups): + col_stats = pa_pf.metadata.row_group(rg_idx).column(col_idx).statistics + assert stats["min"][rg_idx].as_py() == col_stats.min + assert stats["max"][rg_idx].as_py() == col_stats.max + assert stats["null_count"][rg_idx].as_py() == col_stats.null_count + + +def test_parquet_file_statistics_float_with_nulls(): + table = pa.table({"a": [1.0, None, 3.0, None, 5.0, 6.0, None, 8.0]}) + with TemporaryDirectory() as tmp_path: + pq_path = Path(tmp_path) / "test.parquet" + _write_multi_row_group(pq_path, table, rows_per_group=4) + + pf = ParquetFile.open(pq_path) + stats = pa.record_batch(pf.statistics("a")) + assert stats.num_rows == 2 + assert stats["min"].to_pylist() == [1.0, 5.0] + assert stats["max"].to_pylist() == [3.0, 8.0] + assert stats["null_count"].to_pylist() == [2, 1] + + +def test_parquet_file_statistics_first_last_row_group_bounds(): + """The downstream use case: cheaply read the min/max of a column by + slicing the first row group's min and the last row group's max.""" + timestamps = pa.array( + [pa.scalar(i, type=pa.timestamp("us")).as_py() for i in range(12)], + type=pa.timestamp("us"), + ) + table = pa.table({"timestamp": timestamps}) + with TemporaryDirectory() as tmp_path: + pq_path = Path(tmp_path) / "test.parquet" + _write_multi_row_group(pq_path, table, rows_per_group=4) + + pf = ParquetFile.open(pq_path) + stats = pa.record_batch(pf.statistics("timestamp")) + lo = stats["min"][0].as_py() + hi = stats["max"][-1].as_py() + assert lo == timestamps[0].as_py() + assert hi == timestamps[-1].as_py() + + +def test_parquet_file_statistics_unknown_column_raises(): + table = pa.table({"a": [1, 2, 3]}) + with TemporaryDirectory() as tmp_path: + pq_path = Path(tmp_path) / "test.parquet" + write_parquet(table, pq_path) + + pf = ParquetFile.open(pq_path) + with pytest.raises(Exception): + pf.statistics("not_a_column") + + +def test_parquet_file_open_from_file_like(): + table = pa.table({"a": [1, 2, 3, 4]}) + with BytesIO() as bio: + write_parquet(table, bio) + bio.seek(0) + pf = ParquetFile.open(bio) + assert pf.num_rows == 4 + assert pf.num_row_groups == 1 + stats = pa.record_batch(pf.statistics("a")) + assert stats["min"][0].as_py() == 1 + assert stats["max"][0].as_py() == 4 From c41035ec936897ba76292a4d40860991c96594de Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 14 Apr 2026 20:56:15 +0000 Subject: [PATCH 2/5] test(io): cover arro3-native ParquetFile.statistics bounds pattern Add a companion test to `test_parquet_file_statistics_first_last_row_group_bounds` that exercises `stats.column("min")[0].as_py()` / `stats.column("max")[-1].as_py()` directly on the arro3 RecordBatch, with no pyarrow wrapping on the result side. This locks down the recommended downstream usage pattern for cheap footer-only min/max reads. --- tests/io/test_parquet.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/tests/io/test_parquet.py b/tests/io/test_parquet.py index c5477e5..15dcae8 100644 --- a/tests/io/test_parquet.py +++ b/tests/io/test_parquet.py @@ -127,6 +127,27 @@ def test_parquet_file_statistics_first_last_row_group_bounds(): assert hi == timestamps[-1].as_py() +def test_parquet_file_statistics_bounds_arro3_native(): + """Same as the `first_last_row_group_bounds` test, but exercising the + arro3-native API with no pyarrow wrapping on the result side — this is + the recommended downstream usage pattern.""" + timestamps = pa.array( + [pa.scalar(i, type=pa.timestamp("us")).as_py() for i in range(12)], + type=pa.timestamp("us"), + ) + table = pa.table({"timestamp": timestamps}) + with TemporaryDirectory() as tmp_path: + pq_path = Path(tmp_path) / "test.parquet" + _write_multi_row_group(pq_path, table, rows_per_group=4) + + pf = ParquetFile.open(pq_path) + stats = pf.statistics("timestamp") # arro3.core.RecordBatch + lo = stats.column("min")[0].as_py() + hi = stats.column("max")[-1].as_py() + assert lo == timestamps[0].as_py() + assert hi == timestamps[-1].as_py() + + def test_parquet_file_statistics_unknown_column_raises(): table = pa.table({"a": [1, 2, 3]}) with TemporaryDirectory() as tmp_path: From 0028449822a3a2e6b6ad7b43cc89554804284e2b Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 14 Apr 2026 21:00:55 +0000 Subject: [PATCH 3/5] test(io): simplify ParquetFile tests to arro3-native API - Drop `pa.record_batch(...)` wrapping on statistics results; use `stats.column(name)` / `.num_rows` / `.schema.names` / `.to_pylist()` directly on the returned arro3 RecordBatch. - Inline `write_parquet(..., max_row_group_size=4)` calls and remove the `_write_multi_row_group` helper. - Merge the two bounds tests (the pyarrow-wrapped and arro3-native variants were testing the same thing) into a single arro3-native test. --- tests/io/test_parquet.py | 63 +++++++++++++--------------------------- 1 file changed, 20 insertions(+), 43 deletions(-) diff --git a/tests/io/test_parquet.py b/tests/io/test_parquet.py index 15dcae8..91c6f11 100644 --- a/tests/io/test_parquet.py +++ b/tests/io/test_parquet.py @@ -54,57 +54,55 @@ def test_copy_parquet_kv_metadata(): assert reader.schema.metadata[b"hello"] == b"world" -def _write_multi_row_group(path: Path, table: pa.Table, rows_per_group: int) -> None: - write_parquet(table, path, max_row_group_size=rows_per_group) - - def test_parquet_file_accessors(): table = pa.table({"a": list(range(10)), "b": [float(i) for i in range(10)]}) with TemporaryDirectory() as tmp_path: pq_path = Path(tmp_path) / "test.parquet" - _write_multi_row_group(pq_path, table, rows_per_group=4) + write_parquet(table, pq_path, max_row_group_size=4) pf = ParquetFile.open(pq_path) assert pf.num_rows == 10 assert pf.num_row_groups == 3 assert pf.num_columns == 2 - arrow_schema = pa.schema(pf.schema) - assert arrow_schema.names == ["a", "b"] + assert pf.schema.names == ["a", "b"] def test_parquet_file_statistics_int(): table = pa.table({"a": list(range(10))}) with TemporaryDirectory() as tmp_path: pq_path = Path(tmp_path) / "test.parquet" - _write_multi_row_group(pq_path, table, rows_per_group=4) + write_parquet(table, pq_path, max_row_group_size=4) pf = ParquetFile.open(pq_path) - stats = pa.record_batch(pf.statistics("a")) + stats = pf.statistics("a") assert stats.schema.names == ["min", "max", "null_count"] assert stats.num_rows == pf.num_row_groups # Ground-truth: pyarrow footer metadata pa_pf = pq.ParquetFile(pq_path) col_idx = pa_pf.schema_arrow.get_field_index("a") + min_col = stats.column("min") + max_col = stats.column("max") + null_col = stats.column("null_count") for rg_idx in range(pa_pf.num_row_groups): col_stats = pa_pf.metadata.row_group(rg_idx).column(col_idx).statistics - assert stats["min"][rg_idx].as_py() == col_stats.min - assert stats["max"][rg_idx].as_py() == col_stats.max - assert stats["null_count"][rg_idx].as_py() == col_stats.null_count + assert min_col[rg_idx].as_py() == col_stats.min + assert max_col[rg_idx].as_py() == col_stats.max + assert null_col[rg_idx].as_py() == col_stats.null_count def test_parquet_file_statistics_float_with_nulls(): table = pa.table({"a": [1.0, None, 3.0, None, 5.0, 6.0, None, 8.0]}) with TemporaryDirectory() as tmp_path: pq_path = Path(tmp_path) / "test.parquet" - _write_multi_row_group(pq_path, table, rows_per_group=4) + write_parquet(table, pq_path, max_row_group_size=4) pf = ParquetFile.open(pq_path) - stats = pa.record_batch(pf.statistics("a")) + stats = pf.statistics("a") assert stats.num_rows == 2 - assert stats["min"].to_pylist() == [1.0, 5.0] - assert stats["max"].to_pylist() == [3.0, 8.0] - assert stats["null_count"].to_pylist() == [2, 1] + assert stats.column("min").to_pylist() == [1.0, 5.0] + assert stats.column("max").to_pylist() == [3.0, 8.0] + assert stats.column("null_count").to_pylist() == [2, 1] def test_parquet_file_statistics_first_last_row_group_bounds(): @@ -117,31 +115,10 @@ def test_parquet_file_statistics_first_last_row_group_bounds(): table = pa.table({"timestamp": timestamps}) with TemporaryDirectory() as tmp_path: pq_path = Path(tmp_path) / "test.parquet" - _write_multi_row_group(pq_path, table, rows_per_group=4) - - pf = ParquetFile.open(pq_path) - stats = pa.record_batch(pf.statistics("timestamp")) - lo = stats["min"][0].as_py() - hi = stats["max"][-1].as_py() - assert lo == timestamps[0].as_py() - assert hi == timestamps[-1].as_py() - - -def test_parquet_file_statistics_bounds_arro3_native(): - """Same as the `first_last_row_group_bounds` test, but exercising the - arro3-native API with no pyarrow wrapping on the result side — this is - the recommended downstream usage pattern.""" - timestamps = pa.array( - [pa.scalar(i, type=pa.timestamp("us")).as_py() for i in range(12)], - type=pa.timestamp("us"), - ) - table = pa.table({"timestamp": timestamps}) - with TemporaryDirectory() as tmp_path: - pq_path = Path(tmp_path) / "test.parquet" - _write_multi_row_group(pq_path, table, rows_per_group=4) + write_parquet(table, pq_path, max_row_group_size=4) pf = ParquetFile.open(pq_path) - stats = pf.statistics("timestamp") # arro3.core.RecordBatch + stats = pf.statistics("timestamp") lo = stats.column("min")[0].as_py() hi = stats.column("max")[-1].as_py() assert lo == timestamps[0].as_py() @@ -167,6 +144,6 @@ def test_parquet_file_open_from_file_like(): pf = ParquetFile.open(bio) assert pf.num_rows == 4 assert pf.num_row_groups == 1 - stats = pa.record_batch(pf.statistics("a")) - assert stats["min"][0].as_py() == 1 - assert stats["max"][0].as_py() == 4 + stats = pf.statistics("a") + assert stats.column("min")[0].as_py() == 1 + assert stats.column("max")[0].as_py() == 4 From eb2bcf71c6545b6fc2788e49c36bccbe51ff1c18 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 14 Apr 2026 21:08:15 +0000 Subject: [PATCH 4/5] test(io): tighten unknown-column assertion with match pattern MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `pytest.raises(Exception)` is too loose — it catches any error and would mask unrelated failures. Add `match="not_a_column"` so we verify the expected "column not found" error, not just any exception. --- tests/io/test_parquet.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/io/test_parquet.py b/tests/io/test_parquet.py index 91c6f11..d432530 100644 --- a/tests/io/test_parquet.py +++ b/tests/io/test_parquet.py @@ -132,7 +132,7 @@ def test_parquet_file_statistics_unknown_column_raises(): write_parquet(table, pq_path) pf = ParquetFile.open(pq_path) - with pytest.raises(Exception): + with pytest.raises(Exception, match="not_a_column"): pf.statistics("not_a_column") From 81cbcf859d65d0d98995ad62a604abd382cdf6ce Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 14 Apr 2026 21:21:00 +0000 Subject: [PATCH 5/5] fix(io): address ParquetFile audit findings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Export `ParquetFile` from `_io.pyi` (both the import line and `__all__`) so type checkers and mkdocstrings can see it. - Rename `schema` getter to `schema_arrow` to disambiguate from a future Parquet-schema accessor and match pyarrow's `pq.ParquetFile.schema_arrow`. - Map `page_index=True` to `PageIndexPolicy::Optional` (load if present) instead of `Required` (error if missing) — the latter is a footgun for users opting in on files that happen to lack a page index. - Add `__repr__` showing `num_rows`, `num_row_groups`, `num_columns`, in line with other arro3 pyclasses. - Mark the `null_count` field as non-nullable when the default `missing_null_counts_as_zero=True` path is taken (the array is guaranteed to have no nulls), and nullable otherwise. - Add `subclass` to the pyclass attribute to match pyo3-arrow convention. - Tests: - Cover the `missing_null_counts_as_zero=False` branch. - Cover the `skip_arrow_metadata=True` kwarg. - Cover the `page_index=True` kwarg on a file without a page index, asserting it does not raise (pins the Optional-vs-Required choice). - Cover `__repr__`. - Minor: generalize a test docstring ("Downstream use case:" instead of "The downstream use case:"). --- arro3-io/python/arro3/io/_io.pyi | 3 +- arro3-io/python/arro3/io/_parquet.pyi | 8 ++-- arro3-io/src/parquet.rs | 32 +++++++++++-- tests/io/test_parquet.py | 69 ++++++++++++++++++++++++++- 4 files changed, 102 insertions(+), 10 deletions(-) diff --git a/arro3-io/python/arro3/io/_io.pyi b/arro3-io/python/arro3/io/_io.pyi index 3d92a94..bdbc727 100644 --- a/arro3-io/python/arro3/io/_io.pyi +++ b/arro3-io/python/arro3/io/_io.pyi @@ -1,7 +1,7 @@ from ._csv import infer_csv_schema, read_csv, write_csv from ._ipc import read_ipc, read_ipc_stream, write_ipc, write_ipc_stream from ._json import infer_json_schema, read_json, write_json, write_ndjson -from ._parquet import read_parquet, read_parquet_async, write_parquet +from ._parquet import ParquetFile, read_parquet, read_parquet_async, write_parquet __all__ = [ "infer_csv_schema", @@ -15,6 +15,7 @@ __all__ = [ "read_ipc_stream", "write_ipc", "write_ipc_stream", + "ParquetFile", "read_parquet", "read_parquet_async", "write_parquet", diff --git a/arro3-io/python/arro3/io/_parquet.pyi b/arro3-io/python/arro3/io/_parquet.pyi index 184bef9..ecc139f 100644 --- a/arro3-io/python/arro3/io/_parquet.pyi +++ b/arro3-io/python/arro3/io/_parquet.pyi @@ -56,15 +56,17 @@ class ParquetFile: skip_arrow_metadata: If `True`, do not decode the embedded Arrow schema stored in the Parquet key-value metadata. Defaults to `False`. - page_index: If `True`, load the Parquet page index as part of the - footer read. Defaults to `False`. + page_index: If `True`, opportunistically load the Parquet page + index as part of the footer read (loaded if present, skipped + otherwise). Defaults to `False`. Returns: A new `ParquetFile` wrapping the file's footer metadata. """ + def __repr__(self) -> str: ... @property - def schema(self) -> core.Schema: + def schema_arrow(self) -> core.Schema: """The Arrow schema of this Parquet file.""" @property diff --git a/arro3-io/src/parquet.rs b/arro3-io/src/parquet.rs index fe25f73..fd8f123 100644 --- a/arro3-io/src/parquet.rs +++ b/arro3-io/src/parquet.rs @@ -87,7 +87,7 @@ async fn read_parquet_async_inner( /// /// This loads only the Parquet footer (no data pages) and exposes row-group /// level metadata, including per-row-group statistics via [`PyParquetFile::statistics`]. -#[pyclass(module = "arro3.io", name = "ParquetFile", frozen)] +#[pyclass(module = "arro3.io", name = "ParquetFile", subclass, frozen)] pub(crate) struct PyParquetFile { meta: ArrowReaderMetadata, } @@ -103,16 +103,25 @@ impl PyParquetFile { skip_arrow_metadata: bool, page_index: bool, ) -> Arro3IoResult { + // `PageIndexPolicy::Optional` matches the user expectation of + // "load the page index if present, otherwise skip". The upstream + // `From` impl maps `true` to `Required`, which errors if the + // index is missing — a surprising footgun for users opting in. + let page_index_policy = if page_index { + PageIndexPolicy::Optional + } else { + PageIndexPolicy::Skip + }; let options = ArrowReaderOptions::new() .with_skip_arrow_metadata(skip_arrow_metadata) - .with_page_index_policy(PageIndexPolicy::from(page_index)); + .with_page_index_policy(page_index_policy); let meta = ArrowReaderMetadata::load(&mut file, options)?; Ok(Self { meta }) } /// The Arrow schema of this Parquet file. #[getter] - fn schema(&self) -> Arro3Schema { + fn schema_arrow(&self) -> Arro3Schema { self.meta.schema().clone().into() } @@ -134,6 +143,15 @@ impl PyParquetFile { self.meta.schema().fields().len() } + fn __repr__(&self) -> String { + format!( + "arro3.io.ParquetFile(num_rows={}, num_row_groups={}, num_columns={})", + self.num_rows(), + self.num_row_groups(), + self.num_columns(), + ) + } + /// Row-group statistics for a single column. /// /// Returns a `RecordBatch` with one row per row group and the columns @@ -160,10 +178,16 @@ impl PyParquetFile { let max_values = converter.row_group_maxes(parquet_meta.row_groups())?; let null_counts = converter.row_group_null_counts(parquet_meta.row_groups())?; + // When `missing_null_counts_as_zero` is true, `null_counts` is + // guaranteed to contain no nulls, so the field is non-nullable. let schema = Arc::new(Schema::new(vec![ Field::new("min", min_values.data_type().clone(), true), Field::new("max", max_values.data_type().clone(), true), - Field::new("null_count", null_counts.data_type().clone(), true), + Field::new( + "null_count", + null_counts.data_type().clone(), + !missing_null_counts_as_zero, + ), ])); let batch = RecordBatch::try_new(schema, vec![min_values, max_values, Arc::new(null_counts)])?; diff --git a/tests/io/test_parquet.py b/tests/io/test_parquet.py index d432530..6790616 100644 --- a/tests/io/test_parquet.py +++ b/tests/io/test_parquet.py @@ -64,7 +64,10 @@ def test_parquet_file_accessors(): assert pf.num_rows == 10 assert pf.num_row_groups == 3 assert pf.num_columns == 2 - assert pf.schema.names == ["a", "b"] + assert pf.schema_arrow.names == ["a", "b"] + assert repr(pf) == ( + "arro3.io.ParquetFile(num_rows=10, num_row_groups=3, num_columns=2)" + ) def test_parquet_file_statistics_int(): @@ -105,8 +108,34 @@ def test_parquet_file_statistics_float_with_nulls(): assert stats.column("null_count").to_pylist() == [2, 1] +def test_parquet_file_statistics_missing_null_counts_flag(): + """Both branches of the `missing_null_counts_as_zero` flag: when `True` + (default) the `null_count` field is non-nullable; when `False` it is + nullable.""" + table = pa.table({"a": [1, 2, 3, 4]}) + with TemporaryDirectory() as tmp_path: + pq_path = Path(tmp_path) / "test.parquet" + write_parquet(table, pq_path) + + pf = ParquetFile.open(pq_path) + + stats_true = pf.statistics("a") + field_true = stats_true.schema.field("null_count") + assert field_true.nullable is False + + stats_false = pf.statistics("a", missing_null_counts_as_zero=False) + field_false = stats_false.schema.field("null_count") + assert field_false.nullable is True + # The values themselves should be identical for a file whose + # stats include null counts. + assert ( + stats_true.column("null_count").to_pylist() + == stats_false.column("null_count").to_pylist() + ) + + def test_parquet_file_statistics_first_last_row_group_bounds(): - """The downstream use case: cheaply read the min/max of a column by + """Downstream use case: cheaply read the min/max of a column by slicing the first row group's min and the last row group's max.""" timestamps = pa.array( [pa.scalar(i, type=pa.timestamp("us")).as_py() for i in range(12)], @@ -147,3 +176,39 @@ def test_parquet_file_open_from_file_like(): stats = pf.statistics("a") assert stats.column("min")[0].as_py() == 1 assert stats.column("max")[0].as_py() == 4 + + +def test_parquet_file_open_skip_arrow_metadata(): + """With `skip_arrow_metadata=True`, the Arrow schema is reconstructed + from the Parquet schema only, and the embedded `ARROW:schema` KV + metadata is not decoded into the Arrow schema's metadata map.""" + table = pa.table({"a": [1, 2, 3]}) + with TemporaryDirectory() as tmp_path: + pq_path = Path(tmp_path) / "test.parquet" + # Write the file *with* the embedded Arrow schema (default). + write_parquet(table, pq_path) + + pf_default = ParquetFile.open(pq_path) + pf_skipped = ParquetFile.open(pq_path, skip_arrow_metadata=True) + + # Both surface the same logical schema (one int64 column). + assert pf_default.schema_arrow.names == ["a"] + assert pf_skipped.schema_arrow.names == ["a"] + # And both let us read statistics. + assert pf_skipped.statistics("a").column("min")[0].as_py() == 1 + + +def test_parquet_file_open_page_index_optional(): + """`page_index=True` uses an Optional policy: if the file has no page + index (which is the case for a default `write_parquet` call), `open` + must not error.""" + table = pa.table({"a": list(range(8))}) + with TemporaryDirectory() as tmp_path: + pq_path = Path(tmp_path) / "test.parquet" + write_parquet(table, pq_path, max_row_group_size=4) + + # Would raise if we mapped `page_index=True` to Required and the + # file has no page index. + pf = ParquetFile.open(pq_path, page_index=True) + assert pf.num_row_groups == 2 + assert pf.statistics("a").column("min").to_pylist() == [0, 4]