Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion arro3-io/python/arro3/io/_io.pyi
Original file line number Diff line number Diff line change
@@ -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",
Expand All @@ -15,6 +15,7 @@ __all__ = [
"read_ipc_stream",
"write_ipc",
"write_ipc_stream",
"ParquetFile",
"read_parquet",
"read_parquet_async",
"write_parquet",
Expand Down
77 changes: 77 additions & 0 deletions arro3-io/python/arro3/io/_parquet.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
1 change: 1 addition & 0 deletions arro3-io/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,7 @@ fn _io(py: Python, m: &Bound<PyModule>) -> 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::<parquet::PyParquetFile>()?;

Ok(())
}
125 changes: 121 additions & 4 deletions arro3-io/src/parquet.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<PyType>,
mut file: FileReader,
skip_arrow_metadata: bool,
page_index: bool,
) -> Arro3IoResult<Self> {
// `PageIndexPolicy::Optional` matches the user expectation of
// "load the page index if present, otherwise skip". The upstream
// `From<bool>` 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<Arro3RecordBatch> {
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 {
Expand Down
1 change: 1 addition & 0 deletions docs/api/io/parquet.md
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
# Parquet

::: arro3.io.ParquetFile
::: arro3.io.read_parquet
::: arro3.io.read_parquet_async
::: arro3.io.write_parquet
Loading