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 6f2bc61..ecc139f 100644 --- a/arro3-io/python/arro3/io/_parquet.pyi +++ b/arro3-io/python/arro3/io/_parquet.pyi @@ -31,6 +31,83 @@ 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`, 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_arrow(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..fd8f123 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,118 @@ 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", subclass, 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 { + // `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(page_index_policy); + let meta = ArrowReaderMetadata::load(&mut file, options)?; + Ok(Self { meta }) + } + + /// The Arrow schema of this Parquet file. + #[getter] + fn schema_arrow(&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() + } + + 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 + /// `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())?; + + // 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(), + !missing_null_counts_as_zero, + ), + ])); + 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..6790616 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,163 @@ def test_copy_parquet_kv_metadata(): reader = read_parquet(pq_path) assert reader.schema.metadata[b"hello"] == b"world" + + +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_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 + 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(): + table = pa.table({"a": list(range(10))}) + with TemporaryDirectory() as tmp_path: + pq_path = Path(tmp_path) / "test.parquet" + write_parquet(table, pq_path, max_row_group_size=4) + + pf = ParquetFile.open(pq_path) + 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 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_parquet(table, pq_path, max_row_group_size=4) + + pf = ParquetFile.open(pq_path) + stats = pf.statistics("a") + assert stats.num_rows == 2 + 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_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(): + """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_parquet(table, pq_path, max_row_group_size=4) + + pf = ParquetFile.open(pq_path) + stats = pf.statistics("timestamp") + 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: + pq_path = Path(tmp_path) / "test.parquet" + write_parquet(table, pq_path) + + pf = ParquetFile.open(pq_path) + with pytest.raises(Exception, match="not_a_column"): + 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 = 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]