diff --git a/Cargo.lock b/Cargo.lock index 6492f27..5d06ba0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -116,6 +116,7 @@ dependencies = [ "pyo3-file", "pyo3-object_store", "thiserror 1.0.69", + "tokio", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 45611b5..c3bdd96 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -42,6 +42,12 @@ pyo3-async-runtimes = { version = "0.24", features = ["tokio-runtime"] } pyo3-file = "0.12" pyo3-object_store = "0.2" thiserror = "1.0.63" +tokio = { version = "1.40", features = [ + "macros", + "rt", + "rt-multi-thread", + "sync", +] } [profile.release] lto = true diff --git a/arro3-io/Cargo.toml b/arro3-io/Cargo.toml index f81732f..59fb6a1 100644 --- a/arro3-io/Cargo.toml +++ b/arro3-io/Cargo.toml @@ -47,3 +47,4 @@ pyo3-async-runtimes = { workspace = true, features = [ pyo3-file = { workspace = true } pyo3-object_store = { workspace = true, optional = true } thiserror = { workspace = true } +tokio = { workspace = true } diff --git a/arro3-io/python/arro3/io/__init__.py b/arro3-io/python/arro3/io/__init__.py index f6a853d..fa891fb 100644 --- a/arro3-io/python/arro3/io/__init__.py +++ b/arro3-io/python/arro3/io/__init__.py @@ -1,4 +1,12 @@ +from . import _io from ._io import * from ._io import ___version, store __version__: str = ___version() + +__all__ = [ + "__version__", + # "exceptions", + "store", +] +__all__ += _io.__all__ diff --git a/arro3-io/python/arro3/io/_io.pyi b/arro3-io/python/arro3/io/_io.pyi index 3d92a94..c9b5b05 100644 --- a/arro3-io/python/arro3/io/_io.pyi +++ b/arro3-io/python/arro3/io/_io.pyi @@ -1,7 +1,18 @@ 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 ( + ParquetColumnPath, + ParquetCompression, + ParquetEncoding, + ParquetFile, + ParquetOpenOptions, + ParquetPredicate, + ParquetReadOptions, + read_parquet, + read_parquet_async, + write_parquet, +) __all__ = [ "infer_csv_schema", @@ -18,4 +29,13 @@ __all__ = [ "read_parquet", "read_parquet_async", "write_parquet", + "ParquetColumnPath", + "ParquetCompression", + "ParquetEncoding", + "ParquetFile", + "ParquetOpenOptions", + "ParquetPredicate", + "ParquetReadOptions", ] + +def ___version() -> str: ... diff --git a/arro3-io/python/arro3/io/_parquet.pyi b/arro3-io/python/arro3/io/_parquet.pyi index 6f2bc61..f96c6ac 100644 --- a/arro3-io/python/arro3/io/_parquet.pyi +++ b/arro3-io/python/arro3/io/_parquet.pyi @@ -1,13 +1,21 @@ +import sys from pathlib import Path -from typing import IO, Literal, Sequence +from typing import IO, Literal, Protocol, Sequence, TypedDict # Note: importing with # `from arro3.core import Array` # will cause Array to be included in the generated docs in this module. import arro3.core as core import arro3.core.types as types +from obstore.store import ObjectStore as ObstoreStore from ._pyo3_object_store import ObjectStore +from ._stream import RecordBatchStream + +if sys.version_info >= (3, 11): + from typing import Unpack +else: + from typing_extensions import Unpack ParquetColumnPath = str | Sequence[str] """Allowed types to refer to a Parquet Column.""" @@ -51,6 +59,221 @@ async def read_parquet_async(path: str, *, store: ObjectStore) -> core.Table: The loaded Arrow data. """ +class ParquetPredicate(Protocol): + """A predicate operating on RecordBatch. + + See RowFilter for more information on the use of this trait. + """ + + @property + def columns(self) -> Sequence[str]: + """Returns the columns required to evaluate this predicate. + + All projected columns will be provided in the batch passed to evaluate. + """ + def evaluate(self, batch: core.RecordBatch) -> types.ArrowArrayExportable: + """Evaluate this predicate for the given `RecordBatch`. + + Only the columns identified by `Self.columns` will be provided in the batch. + + Must return a boolean-typed `Array` that has the same length as the input batch + where each row indicates whether the row should be returned: + + - `True`: the row should be returned + - `False` or `null`: the row should not be returned + """ + +class ParquetOpenOptions(TypedDict, total=False): + """Options passed when opening `ParquetFile`.""" + + store: ObjectStore | ObstoreStore | None + """ + A store to use when opening Parquet files from cloud storage. + """ + + skip_arrow_metadata: bool + """If `True`, skip decoding the embedded arrow metadata. + + Parquet files generated by some writers may contain embedded arrow schema and + metadata. This may not be correct or compatible with your system, for example: + [ARROW-16184](https://issues.apache.org/jira/browse/ARROW-16184). + """ + + schema: core.Schema | None + """Provide a schema to use when reading the Parquet file. + + If provided it takes precedence over the schema inferred from the file or the schema + defined in the file's metadata. If the schema is not compatible with the file's + schema an error will be returned when constructing the builder. + + This option is only required if you want to cast columns to a different type. For + example, if you wanted to cast from an `Int64` in the Parquet file to a `Timestamp` + in the Arrow schema. + + The supplied schema must have the same number of columns as the Parquet schema and + the column names need to be the same. + """ + + page_index: bool + """Enable reading [`PageIndex`], if present. + + The `PageIndex` can be used to push down predicates to the parquet scan, potentially + eliminating unnecessary IO, by some query engines. + + [PageIndex]: https://github.com/apache/parquet-format/blob/master/PageIndex.md + """ + +class ParquetReadOptions(TypedDict, total=False): + """Options passed to read calls of `ParquetFile`.""" + + batch_size: int | None + """Set the size of `RecordBatch` to produce. + """ + row_groups: Sequence[int] | None + """Only read data from the provided row group indexes. + + This is also called row group filtering. + """ + columns: Sequence[str] | None + """Only read data from the provided column indexes. + """ + filters: ParquetPredicate | Sequence[ParquetPredicate] | None + """Provide one or more filters to skip decoding rows. + + These filters are applied during the parquet read process **after** row group + selection. + + Filters are applied in order after decoding only the columns + required. As filters eliminate rows, fewer rows from subsequent columns + may be required, thus potentially reducing IO and decode. + + `filters` consists of one or more [`ParquetPredicate`s][arro3.io.ParquetPredicate]. + Only the rows for which all the predicates evaluate to `true` will be returned. + + + + This design has a couple of implications: + + - filters can be used to skip entire pages, and thus IO, in addition to CPU decode overheads + - Columns may be decoded multiple times if they appear in multiple predicates + - IO will be deferred until needed by columns required by a predicate. + + As such there is a trade-off between a single large predicate, or multiple + predicates, that will depend on the shape of the data. Whilst multiple smaller + predicates may minimise the amount of data scanned/decoded, it may not be faster + overall. + + For example, if a predicate that needs a single column of data filters out all but + 1% of the rows, applying it as one of the early predicates will likely significantly + improve performance. + + As a counter example, if a predicate needs several columns of data to evaluate but + leaves 99% of the rows, it may be better to not filter the data from parquet and + apply the filter after the `RecordBatch` has been fully decoded. + + Additionally, even if a predicate eliminates a moderate number of rows, it may still + be faster to filter the data after the RecordBatch has been fully decoded, if the + eliminated rows are not contiguous. + + !!! note + + It is recommended to enable reading the page index if using this functionality, + to allow more efficient skipping over data pages. See + [`ParquetOpenOptions.page_index`][arro3.io.ParquetOpenOptions.page_index]. + """ + + limit: int | None + """Provide a limit to the number of rows to be read. + + + + The limit will be applied after any `filters` allowing it to limit the final set of + rows decoded after any pushed down predicates. + + !!! note + + It is recommended to enable reading the page index if using this functionality, + to allow more efficient skipping over data pages. See + [`ParquetOpenOptions.page_index`][arro3.io.ParquetOpenOptions.page_index]. + """ + offset: int | None + """Provide an offset to skip over the given number of rows. + + The offset will be applied after any `filters` allowing it to skip rows after any + pushed down predicates. + + !!! note + + It is recommended to enable reading the page index if using this functionality, + to allow more efficient skipping over data pages. See + [`ParquetOpenOptions.page_index`][arro3.io.ParquetOpenOptions.page_index]. + """ + +class ParquetFile: + @classmethod + def open( + cls, + file: IO[bytes] | Path | str, + **kwargs: Unpack[ParquetOpenOptions], + ) -> ParquetFile: + """Open a Parquet file.""" + + @classmethod + async def open_async( + cls, + file: IO[bytes] | Path | str, + *, + store: ObjectStore | ObstoreStore | None = None, + skip_arrow_metadata: bool = False, + schema: core.Schema | None = None, + page_index: bool = False, + ) -> ParquetFile: + """Open a Parquet file.""" + + @property + def num_row_groups(self) -> int: + """Return the number of row groups in the Parquet file.""" + + def read(self, **kwargs: Unpack[ParquetReadOptions]) -> core.RecordBatchReader: + """Read the Parquet file to an Arrow RecordBatchReader. + + Keyword Args: + batch_size: The number of rows to read in each batch. + row_groups: The row groups to read. + columns: The columns to read. + limit: The number of rows to read. + offset: The number of rows to skip. + + Returns: + The loaded Arrow data. + """ + def read_async(self, **kwargs: Unpack[ParquetReadOptions]) -> RecordBatchStream: + """Read the Parquet file to an Arrow async RecordBatchStream. + + Note that this method itself is **not async**, but returns an async iterable + that yields RecordBatches. + + Keyword Args: + batch_size: The number of rows to read in each batch. + row_groups: The row groups to read. + columns: The columns to read. + limit: The number of rows to read. + offset: The number of rows to skip. + + Returns: + The loaded Arrow data. + """ + @property + def schema_arrow(self) -> core.Schema: + """Return the Arrow schema of the Parquet file.""" + def write_parquet( data: types.ArrowStreamExportable | types.ArrowArrayExportable, file: IO[bytes] | Path | str, diff --git a/arro3-io/python/arro3/io/_stream.pyi b/arro3-io/python/arro3/io/_stream.pyi new file mode 100644 index 0000000..dc76ff1 --- /dev/null +++ b/arro3-io/python/arro3/io/_stream.pyi @@ -0,0 +1,12 @@ +# Note: importing with +# `from arro3.core import Array` +# will cause Array to be included in the generated docs in this module. +import arro3.core as core + +class RecordBatchStream: + def __aiter__(self) -> RecordBatchStream: + """Return `Self` as an async iterator.""" + async def __anext__(self) -> core.RecordBatch: + """Return the next record batch in the stream.""" + async def collect_async(self) -> core.Table: + """Collect the stream into a single table.""" diff --git a/arro3-io/src/csv.rs b/arro3-io/src/csv.rs index d0f77f0..d44cfaf 100644 --- a/arro3-io/src/csv.rs +++ b/arro3-io/src/csv.rs @@ -8,7 +8,7 @@ use pyo3_arrow::export::{Arro3RecordBatchReader, Arro3Schema}; use pyo3_arrow::input::AnyRecordBatch; use pyo3_arrow::{PyRecordBatchReader, PySchema}; -use crate::utils::{FileReader, FileWriter}; +use crate::source::{FileWriter, SyncReader}; /// Infer a CSV file's schema #[pyfunction] @@ -25,7 +25,7 @@ use crate::utils::{FileReader, FileWriter}; ))] #[allow(clippy::too_many_arguments)] pub fn infer_csv_schema( - file: FileReader, + file: SyncReader, has_header: Option, max_records: Option, delimiter: Option, @@ -76,7 +76,7 @@ pub fn infer_csv_schema( ))] #[allow(clippy::too_many_arguments)] pub fn read_csv( - file: FileReader, + file: SyncReader, schema: PySchema, has_header: Option, batch_size: Option, diff --git a/arro3-io/src/error.rs b/arro3-io/src/error.rs index b7a957a..da21a14 100644 --- a/arro3-io/src/error.rs +++ b/arro3-io/src/error.rs @@ -1,6 +1,6 @@ //! Contains the [`Arro3IoError`], the Error returned by most fallible functions in this crate. -use pyo3::exceptions::{PyException, PyValueError}; +use pyo3::exceptions::{PyException, PyIOError, PyValueError}; use pyo3::prelude::*; use pyo3::DowncastError; use thiserror::Error; @@ -13,6 +13,10 @@ pub enum Arro3IoError { #[error(transparent)] ArrowError(#[from] arrow_schema::ArrowError), + /// A wrapped [std::io::Error] + #[error(transparent)] + IOError(#[from] std::io::Error), + /// A wrapped [object_store::Error] #[error(transparent)] ObjectStoreError(#[from] object_store::Error), @@ -29,10 +33,11 @@ pub enum Arro3IoError { impl From for PyErr { fn from(error: Arro3IoError) -> Self { match error { - Arro3IoError::PyErr(err) => err, Arro3IoError::ArrowError(err) => PyException::new_err(err.to_string()), + Arro3IoError::IOError(err) => PyIOError::new_err(err.to_string()), Arro3IoError::ObjectStoreError(err) => PyException::new_err(err.to_string()), Arro3IoError::ParquetError(err) => PyException::new_err(err.to_string()), + Arro3IoError::PyErr(err) => err, } } } diff --git a/arro3-io/src/ipc.rs b/arro3-io/src/ipc.rs index 0e143ad..723ca61 100644 --- a/arro3-io/src/ipc.rs +++ b/arro3-io/src/ipc.rs @@ -9,11 +9,11 @@ use pyo3_arrow::export::Arro3RecordBatchReader; use pyo3_arrow::input::AnyRecordBatch; use pyo3_arrow::PyRecordBatchReader; -use crate::utils::{FileReader, FileWriter}; +use crate::source::{FileWriter, SyncReader}; /// Read an Arrow IPC file to an Arrow RecordBatchReader #[pyfunction] -pub fn read_ipc(file: FileReader) -> PyArrowResult { +pub fn read_ipc(file: SyncReader) -> PyArrowResult { let builder = FileReaderBuilder::new(); let buf_file = BufReader::new(file); let reader = builder.build(buf_file)?; @@ -22,7 +22,7 @@ pub fn read_ipc(file: FileReader) -> PyArrowResult { /// Read an Arrow IPC Stream file to an Arrow RecordBatchReader #[pyfunction] -pub fn read_ipc_stream(file: FileReader) -> PyArrowResult { +pub fn read_ipc_stream(file: SyncReader) -> PyArrowResult { let reader = StreamReader::try_new(file, None)?; Ok(PyRecordBatchReader::new(Box::new(reader)).into()) } diff --git a/arro3-io/src/json.rs b/arro3-io/src/json.rs index 4b8b8b1..aec1954 100644 --- a/arro3-io/src/json.rs +++ b/arro3-io/src/json.rs @@ -8,7 +8,7 @@ use pyo3_arrow::export::{Arro3RecordBatchReader, Arro3Schema}; use pyo3_arrow::input::AnyRecordBatch; use pyo3_arrow::{PyRecordBatchReader, PySchema}; -use crate::utils::{FileReader, FileWriter}; +use crate::source::{FileWriter, SyncReader}; /// Infer a JSON file's schema #[pyfunction] @@ -18,7 +18,7 @@ use crate::utils::{FileReader, FileWriter}; max_records=None, ))] pub fn infer_json_schema( - file: FileReader, + file: SyncReader, max_records: Option, ) -> PyArrowResult { let buf_file = BufReader::new(file); @@ -35,7 +35,7 @@ pub fn infer_json_schema( batch_size=None, ))] pub fn read_json( - file: FileReader, + file: SyncReader, schema: PySchema, batch_size: Option, ) -> PyArrowResult { diff --git a/arro3-io/src/lib.rs b/arro3-io/src/lib.rs index 025f9d8..3c16841 100644 --- a/arro3-io/src/lib.rs +++ b/arro3-io/src/lib.rs @@ -8,7 +8,8 @@ mod error; mod ipc; mod json; mod parquet; -mod utils; +mod runtime; +mod source; const VERSION: &str = env!("CARGO_PKG_VERSION"); @@ -56,9 +57,10 @@ fn _io(py: Python, m: &Bound) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(ipc::write_ipc))?; m.add_wrapped(wrap_pyfunction!(ipc::write_ipc_stream))?; - 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_wrapped(wrap_pyfunction!(parquet::reader::read_parquet))?; + m.add_wrapped(wrap_pyfunction!(parquet::reader::read_parquet_async))?; + m.add_wrapped(wrap_pyfunction!(parquet::writer::write_parquet))?; + m.add_class::()?; Ok(()) } diff --git a/arro3-io/src/parquet/mod.rs b/arro3-io/src/parquet/mod.rs new file mode 100644 index 0000000..9d92585 --- /dev/null +++ b/arro3-io/src/parquet/mod.rs @@ -0,0 +1,2 @@ +pub(crate) mod reader; +pub(crate) mod writer; diff --git a/arro3-io/src/parquet/reader/file.rs b/arro3-io/src/parquet/reader/file.rs new file mode 100644 index 0000000..8997408 --- /dev/null +++ b/arro3-io/src/parquet/reader/file.rs @@ -0,0 +1,252 @@ +use std::sync::Arc; + +use arrow_array::{RecordBatch, RecordBatchReader}; +use arrow_schema::ArrowError; +use futures::StreamExt; +use object_store::ObjectStore; +use parquet::arrow::arrow_reader::{ + ArrowReaderMetadata, ArrowReaderOptions, ParquetRecordBatchReaderBuilder, +}; +use parquet::arrow::async_reader::{ + AsyncFileReader, ParquetObjectReader, ParquetRecordBatchStream, +}; +use parquet::arrow::ParquetRecordBatchStreamBuilder; +use pyo3::prelude::*; +use pyo3::types::PyType; +use pyo3::IntoPyObjectExt; +use pyo3_arrow::export::{Arro3RecordBatchReader, Arro3Schema}; +use pyo3_arrow::{PyRecordBatchReader, PySchema}; +use pyo3_async_runtimes::tokio::future_into_py; +use pyo3_object_store::AnyObjectStore; + +use crate::error::Arro3IoResult; +use crate::parquet::reader::options::PyParquetOptions; +use crate::parquet::reader::stream::PyRecordBatchStream; +use crate::runtime::get_runtime; +use crate::source::{AsyncReader, SyncReader}; + +enum ParquetSource { + Sync(SyncReader), + Async(AsyncReader), +} + +impl ParquetSource { + fn new(store: Option>, file: Bound) -> Arro3IoResult { + if let Some(store) = store { + let reader = ParquetObjectReader::new(store, file.extract::<&str>()?.into()); + Ok(ParquetSource::Async(AsyncReader::ObjectStore(reader))) + } else { + Ok(ParquetSource::Sync(file.extract()?)) + } + } +} + +impl From for ParquetSource { + fn from(value: SyncReader) -> Self { + Self::Sync(value) + } +} + +impl From for ParquetSource { + fn from(value: AsyncReader) -> Self { + Self::Async(value) + } +} + +impl SyncReader { + fn open_parquet(&self, options: ArrowReaderOptions) -> Arro3IoResult { + Ok(ArrowReaderMetadata::load(self, options)?) + } +} + +impl AsyncReader { + async fn open_parquet( + &mut self, + options: ArrowReaderOptions, + ) -> Arro3IoResult { + Ok(ArrowReaderMetadata::load_async(self, options).await?) + } +} + +impl ParquetSource { + async fn open_parquet( + &mut self, + options: ArrowReaderOptions, + ) -> Arro3IoResult { + match self { + ParquetSource::Sync(sync_source) => { + Ok(ArrowReaderMetadata::load(sync_source, options)?) + } + ParquetSource::Async(async_source) => { + Ok(ArrowReaderMetadata::load_async(async_source, options).await?) + } + } + } +} + +/// Reader interface for a single Parquet file. +#[pyclass(module = "arro3.io", frozen)] +pub struct ParquetFile { + meta: ArrowReaderMetadata, + source: ParquetSource, +} + +#[pymethods] +impl ParquetFile { + #[classmethod] + #[pyo3(signature = (file, *, store=None, skip_arrow_metadata=false, schema=None, page_index=false))] + pub(crate) fn open( + _cls: &Bound, + py: Python, + file: Bound, + store: Option, + skip_arrow_metadata: bool, + schema: Option, + page_index: bool, + ) -> Arro3IoResult { + let mut options = ArrowReaderOptions::default() + .with_skip_arrow_metadata(skip_arrow_metadata) + .with_page_index(page_index); + if let Some(schema) = schema { + options = options.with_schema(schema.into_inner()); + } + + let runtime = get_runtime(py)?; + + let mut source = ParquetSource::new(store.map(|s| s.into()), file)?; + let meta = runtime.block_on(source.open_parquet(options))?; + Ok(Self { meta, source }) + } + + #[classmethod] + #[pyo3(signature = (file, *, store=None, skip_arrow_metadata=false, schema=None, page_index=false))] + pub(crate) fn open_async<'py>( + _cls: &Bound, + py: Python<'py>, + file: Bound<'py, PyAny>, + store: Option, + skip_arrow_metadata: bool, + schema: Option, + page_index: bool, + ) -> PyResult> { + let mut options = ArrowReaderOptions::default() + .with_skip_arrow_metadata(skip_arrow_metadata) + .with_page_index(page_index); + if let Some(schema) = schema { + options = options.with_schema(schema.into_inner()); + } + + if let Some(store) = store { + let store = store.into_dyn(); + let mut reader = AsyncReader::ObjectStore(ParquetObjectReader::new( + store, + file.extract::<&str>()?.into(), + )); + future_into_py(py, async move { + let meta = reader.open_parquet(options).await?; + Ok(Self { + meta, + source: reader.into(), + }) + }) + } else { + let reader = file.extract::()?; + let meta = reader.open_parquet(options)?; + let slf = Self { + meta, + source: reader.into(), + }; + slf.into_bound_py_any(py) + } + } + + #[getter] + fn num_row_groups(&self) -> usize { + self.meta.metadata().num_row_groups() + } + + #[pyo3(signature = (**kwargs))] + fn read(&self, kwargs: Option) -> Arro3IoResult { + let options = kwargs.unwrap_or_default(); + match &self.source { + ParquetSource::Sync(sync_source) => { + let sync_reader_builder = ParquetRecordBatchReaderBuilder::new_with_metadata( + sync_source.try_clone()?, + self.meta.clone(), + ); + let record_batch_reader = options + .apply_to_reader_builder(sync_reader_builder, &self.meta) + .build()?; + Ok(PyRecordBatchReader::new(Box::new(record_batch_reader)).into()) + } + ParquetSource::Async(async_source) => { + let async_reader_builder = ParquetRecordBatchStreamBuilder::new_with_metadata( + async_source.clone(), + self.meta.clone(), + ); + let record_batch_stream = options + .apply_to_reader_builder(async_reader_builder, &self.meta) + .build()?; + let blocking_record_batch_reader = BlockingAsyncParquetReader(record_batch_stream); + Ok(PyRecordBatchReader::new(Box::new(blocking_record_batch_reader)).into()) + } + } + } + + #[pyo3(signature = (**kwargs))] + fn read_async(&self, kwargs: Option) -> Arro3IoResult { + let options = kwargs.unwrap_or_default(); + match &self.source { + ParquetSource::Sync(sync_source) => { + let async_reader_builder = ParquetRecordBatchStreamBuilder::new_with_metadata( + Box::new(sync_source.try_clone()?) as _, + self.meta.clone(), + ); + let record_batch_stream = options + .apply_to_reader_builder(async_reader_builder, &self.meta) + .build()?; + Ok(PyRecordBatchStream::new(record_batch_stream)) + } + ParquetSource::Async(async_source) => { + let async_reader_builder = ParquetRecordBatchStreamBuilder::new_with_metadata( + Box::new(async_source.clone()) as _, + self.meta.clone(), + ); + let record_batch_stream = options + .apply_to_reader_builder(async_reader_builder, &self.meta) + .build()?; + Ok(PyRecordBatchStream::new(record_batch_stream)) + } + } + } + + #[getter] + fn schema_arrow(&self) -> Arro3Schema { + self.meta.schema().clone().into() + } +} + +struct BlockingAsyncParquetReader( + ParquetRecordBatchStream, +); + +impl Iterator for BlockingAsyncParquetReader { + type Item = std::result::Result; + + fn next(&mut self) -> Option { + Python::with_gil(|py| { + let runtime = get_runtime(py).unwrap(); + runtime + .block_on(self.0.next()) + .map(|maybe_batch| maybe_batch.map_err(|err| err.into())) + }) + } +} + +impl RecordBatchReader + for BlockingAsyncParquetReader +{ + fn schema(&self) -> arrow_schema::SchemaRef { + self.0.schema().clone() + } +} diff --git a/arro3-io/src/parquet/reader/functional.rs b/arro3-io/src/parquet/reader/functional.rs new file mode 100644 index 0000000..0e57ae0 --- /dev/null +++ b/arro3-io/src/parquet/reader/functional.rs @@ -0,0 +1,68 @@ +use std::sync::Arc; + +use arrow_array::{RecordBatchIterator, RecordBatchReader}; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use parquet::arrow::async_reader::ParquetObjectReader; +use pyo3::prelude::*; +use pyo3_arrow::error::PyArrowResult; +use pyo3_arrow::export::Arro3RecordBatchReader; +use pyo3_arrow::{PyRecordBatchReader, PyTable}; +use pyo3_object_store::PyObjectStore; + +use crate::error::Arro3IoResult; +use crate::source::SyncReader; + +#[pyfunction] +pub fn read_parquet(file: SyncReader) -> PyArrowResult { + let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap(); + + let metadata = builder.schema().metadata().clone(); + let reader = builder.build().unwrap(); + + // Add source schema metadata onto reader's schema. The original schema is not valid + // with a given column projection, but we want to persist the source's metadata. + let arrow_schema = Arc::new(reader.schema().as_ref().clone().with_metadata(metadata)); + + // Create a new iterator with the arrow schema specifically + // + // Passing ParquetRecordBatchReader directly to PyRecordBatchReader::new loses schema + // metadata + // + // https://docs.rs/parquet/latest/parquet/arrow/arrow_reader/struct.ParquetRecordBatchReader.html#method.schema + // https://github.com/apache/arrow-rs/pull/5135 + let iter = Box::new(RecordBatchIterator::new(reader, arrow_schema)); + Ok(PyRecordBatchReader::new(iter).into()) +} + +#[pyfunction] +#[pyo3(signature = (path, *, store))] +pub fn read_parquet_async( + py: Python, + path: String, + store: PyObjectStore, +) -> PyArrowResult { + let fut = pyo3_async_runtimes::tokio::future_into_py(py, async move { + Ok(read_parquet_async_inner(store.into_inner(), path).await?) + })?; + + Ok(fut.into()) +} + +async fn read_parquet_async_inner( + store: Arc, + path: String, +) -> Arro3IoResult { + use futures::TryStreamExt; + use parquet::arrow::ParquetRecordBatchStreamBuilder; + + let object_reader = ParquetObjectReader::new(store, path.into()); + let builder = ParquetRecordBatchStreamBuilder::new(object_reader).await?; + + let metadata = builder.schema().metadata().clone(); + let reader = builder.build()?; + + let arrow_schema = Arc::new(reader.schema().as_ref().clone().with_metadata(metadata)); + + let batches = reader.try_collect::>().await?; + Ok(PyTable::try_new(batches, arrow_schema)?) +} diff --git a/arro3-io/src/parquet/reader/mod.rs b/arro3-io/src/parquet/reader/mod.rs new file mode 100644 index 0000000..480d65e --- /dev/null +++ b/arro3-io/src/parquet/reader/mod.rs @@ -0,0 +1,7 @@ +mod file; +mod functional; +mod options; +mod stream; + +pub(crate) use file::ParquetFile; +pub(crate) use functional::{read_parquet, read_parquet_async}; diff --git a/arro3-io/src/parquet/reader/options.rs b/arro3-io/src/parquet/reader/options.rs new file mode 100644 index 0000000..1c84997 --- /dev/null +++ b/arro3-io/src/parquet/reader/options.rs @@ -0,0 +1,191 @@ +use arrow::array::AsArray; +use arrow_array::{BooleanArray, RecordBatch}; +use arrow_schema::{ArrowError, DataType}; +use parquet::arrow::arrow_reader::{ + ArrowPredicate, ArrowReaderBuilder, ArrowReaderMetadata, RowFilter, +}; +use parquet::arrow::ProjectionMask; +use parquet::schema::types::SchemaDescriptor; +use pyo3::exceptions::PyTypeError; +use pyo3::intern; +use pyo3::prelude::*; +use pyo3_arrow::export::Arro3RecordBatch; +use pyo3_arrow::PyArray; + +/// Note that these are a list of columns and are not yet resolved until [resolve] is called. +#[derive(Debug, FromPyObject)] +struct PyProjectionMask(Vec); + +impl PyProjectionMask { + fn resolve(&self, schema: &SchemaDescriptor) -> ProjectionMask { + ProjectionMask::columns(schema, self.0.iter().map(|s| s.as_str())) + } +} + +#[derive(Debug)] +struct PyInputPredicate { + evaluate: PyObject, + projection: PyProjectionMask, +} + +impl PyInputPredicate { + fn into_arrow_predicate(self, schema: &SchemaDescriptor) -> Box { + Box::new(SchemaResolvedPredicate { + evaluate: self.evaluate, + projection: self.projection.resolve(schema), + }) + } +} + +impl<'py> FromPyObject<'py> for PyInputPredicate { + fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { + let py = ob.py(); + let columns = ob.getattr(intern!(py, "columns"))?.extract()?; + + let evaluate_callback = ob.getattr(intern!(py, "evaluate"))?; + + // Check that the callback object is actually callable + let builtins_mod = py.import(intern!(py, "builtins"))?; + if !builtins_mod + .call_method1(intern!(py, "callable"), (&evaluate_callback,))? + .extract::()? + { + return Err(PyTypeError::new_err("evaluate must be callable")); + } + + Ok(Self { + projection: columns, + evaluate: evaluate_callback.unbind(), + }) + } +} + +struct SchemaResolvedPredicate { + evaluate: PyObject, + projection: ProjectionMask, +} + +impl SchemaResolvedPredicate { + fn evaluate_inner(&mut self, py: Python, batch: RecordBatch) -> PyResult { + let result = self + .evaluate + .bind(py) + .call1((Arro3RecordBatch::from(batch),))?; + result.extract::() + } +} + +impl ArrowPredicate for SchemaResolvedPredicate { + fn projection(&self) -> &ProjectionMask { + &self.projection + } + + fn evaluate(&mut self, batch: RecordBatch) -> std::result::Result { + let py_array = Python::with_gil(|py| self.evaluate_inner(py, batch)) + .map_err(|err| ArrowError::ExternalError(Box::new(err)))?; + let (array, _field) = py_array.into_inner(); + if !matches!(array.data_type(), DataType::Boolean) { + Err(ArrowError::SchemaError(format!( + "Expected a boolean array, got {:?}", + array.data_type() + ))) + } else { + Ok(array.as_boolean().clone()) + } + } +} + +#[derive(Debug, Default)] +pub(crate) struct PyParquetOptions { + batch_size: Option, + row_groups: Option>, + columns: Option, + filters: Option>, + // selection: Option, + limit: Option, + offset: Option, +} + +impl<'py> FromPyObject<'py> for PyParquetOptions { + fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { + let py = ob.py(); + + let mut batch_size = None; + let mut row_groups = None; + let mut columns = None; + let mut filters = None; + let mut limit = None; + let mut offset = None; + if let Ok(val) = ob.get_item(intern!(py, "batch_size")) { + batch_size = Some(val.extract()?); + } + if let Ok(val) = ob.get_item(intern!(py, "row_groups")) { + row_groups = Some(val.extract()?); + } + if let Ok(val) = ob.get_item(intern!(py, "columns")) { + columns = Some(val.extract()?); + } + if let Ok(val) = ob.get_item(intern!(py, "filters")) { + // Can either be a sequence of filters or a single filter + let mut local_filters = vec![]; + match val.try_iter() { + Ok(iter) => { + for item in iter { + local_filters.push(item?.extract()?); + } + } + Err(_) => { + local_filters.push(val.extract()?); + } + } + filters = Some(local_filters); + } + if let Ok(val) = ob.get_item(intern!(py, "limit")) { + limit = Some(val.extract()?); + } + if let Ok(val) = ob.get_item(intern!(py, "offset")) { + offset = Some(val.extract()?); + } + + Ok(Self { + batch_size, + row_groups, + columns, + filters, + limit, + offset, + }) + } +} + +impl PyParquetOptions { + pub(crate) fn apply_to_reader_builder( + self, + mut builder: ArrowReaderBuilder, + metadata: &ArrowReaderMetadata, + ) -> ArrowReaderBuilder { + if let Some(batch_size) = self.batch_size { + builder = builder.with_batch_size(batch_size); + } + if let Some(row_groups) = self.row_groups { + builder = builder.with_row_groups(row_groups); + } + if let Some(columns) = self.columns { + builder = builder.with_projection(columns.resolve(metadata.parquet_schema())); + } + if let Some(filters) = self.filters { + let predicates = filters + .into_iter() + .map(|predicate| predicate.into_arrow_predicate(metadata.parquet_schema())) + .collect(); + builder = builder.with_row_filter(RowFilter::new(predicates)); + } + if let Some(limit) = self.limit { + builder = builder.with_limit(limit); + } + if let Some(offset) = self.offset { + builder = builder.with_offset(offset); + } + builder + } +} diff --git a/arro3-io/src/parquet/reader/stream.rs b/arro3-io/src/parquet/reader/stream.rs new file mode 100644 index 0000000..1308d10 --- /dev/null +++ b/arro3-io/src/parquet/reader/stream.rs @@ -0,0 +1,84 @@ +use std::sync::Arc; + +use arrow_schema::SchemaRef; +use futures::StreamExt; +use parquet::arrow::async_reader::{AsyncFileReader, ParquetRecordBatchStream}; +use pyo3::exceptions::{PyStopAsyncIteration, PyStopIteration}; +use pyo3::prelude::*; +use pyo3_arrow::export::{Arro3RecordBatch, Arro3Table}; +use pyo3_arrow::PyTable; +use tokio::sync::Mutex; + +use crate::error::Arro3IoError; + +#[pyclass(name = "RecordBatchStream", frozen)] +pub(crate) struct PyRecordBatchStream { + stream: Arc>>>, + schema: SchemaRef, +} + +impl PyRecordBatchStream { + pub(crate) fn new( + stream: ParquetRecordBatchStream>, + ) -> Self { + let schema = stream.schema().clone(); + Self { + stream: Arc::new(Mutex::new(stream)), + schema, + } + } +} + +#[pymethods] +impl PyRecordBatchStream { + fn __aiter__(slf: Py) -> Py { + slf + } + + fn __anext__<'py>(&'py self, py: Python<'py>) -> PyResult> { + let stream = self.stream.clone(); + pyo3_async_runtimes::tokio::future_into_py(py, next_stream(stream, false)) + } + + fn collect_async<'py>(&'py self, py: Python<'py>) -> PyResult> { + let stream = self.stream.clone(); + pyo3_async_runtimes::tokio::future_into_py(py, collect_stream(stream, self.schema.clone())) + } +} + +async fn next_stream( + stream: Arc>>>, + is_sync: bool, +) -> PyResult { + let mut stream = stream.lock().await; + match stream.next().await { + Some(Ok(batch)) => Ok(batch.into()), + Some(Err(err)) => Err(Arro3IoError::from(err).into()), + None => { + // Depending on whether the iteration is sync or not, we raise either a + // StopIteration or a StopAsyncIteration + if is_sync { + Err(PyStopIteration::new_err("stream exhausted")) + } else { + Err(PyStopAsyncIteration::new_err("stream exhausted")) + } + } + } +} + +async fn collect_stream( + stream: Arc>>>, + schema: SchemaRef, +) -> PyResult { + let mut stream = stream.lock().await; + let mut batches: Vec<_> = vec![]; + loop { + match stream.next().await { + Some(Ok(batch)) => { + batches.push(batch); + } + Some(Err(err)) => return Err(Arro3IoError::ParquetError(err).into()), + None => return Ok(PyTable::try_new(batches, schema)?.into()), + }; + } +} diff --git a/arro3-io/src/parquet.rs b/arro3-io/src/parquet/writer.rs similarity index 74% rename from arro3-io/src/parquet.rs rename to arro3-io/src/parquet/writer.rs index 8d27bc2..076ec88 100644 --- a/arro3-io/src/parquet.rs +++ b/arro3-io/src/parquet/writer.rs @@ -1,11 +1,8 @@ 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::RecordBatchReader; use parquet::arrow::arrow_writer::ArrowWriterOptions; -use parquet::arrow::async_reader::ParquetObjectReader; use parquet::arrow::ArrowWriter; use parquet::basic::{Compression, Encoding}; use parquet::file::properties::{WriterProperties, WriterVersion}; @@ -14,68 +11,9 @@ use parquet::schema::types::ColumnPath; use pyo3::exceptions::{PyTypeError, PyValueError}; use pyo3::prelude::*; use pyo3_arrow::error::PyArrowResult; -use pyo3_arrow::export::Arro3RecordBatchReader; use pyo3_arrow::input::AnyRecordBatch; -use pyo3_arrow::{PyRecordBatchReader, PyTable}; -use pyo3_object_store::PyObjectStore; -use crate::error::Arro3IoResult; -use crate::utils::{FileReader, FileWriter}; - -#[pyfunction] -pub fn read_parquet(file: FileReader) -> PyArrowResult { - let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap(); - - let metadata = builder.schema().metadata().clone(); - let reader = builder.build().unwrap(); - - // Add source schema metadata onto reader's schema. The original schema is not valid - // with a given column projection, but we want to persist the source's metadata. - let arrow_schema = Arc::new(reader.schema().as_ref().clone().with_metadata(metadata)); - - // Create a new iterator with the arrow schema specifically - // - // Passing ParquetRecordBatchReader directly to PyRecordBatchReader::new loses schema - // metadata - // - // https://docs.rs/parquet/latest/parquet/arrow/arrow_reader/struct.ParquetRecordBatchReader.html#method.schema - // https://github.com/apache/arrow-rs/pull/5135 - let iter = Box::new(RecordBatchIterator::new(reader, arrow_schema)); - Ok(PyRecordBatchReader::new(iter).into()) -} - -#[pyfunction] -#[pyo3(signature = (path, *, store))] -pub fn read_parquet_async( - py: Python, - path: String, - store: PyObjectStore, -) -> PyArrowResult { - let fut = pyo3_async_runtimes::tokio::future_into_py(py, async move { - Ok(read_parquet_async_inner(store.into_inner(), path).await?) - })?; - - Ok(fut.into()) -} - -async fn read_parquet_async_inner( - store: Arc, - path: String, -) -> Arro3IoResult { - use futures::TryStreamExt; - use parquet::arrow::ParquetRecordBatchStreamBuilder; - - let object_reader = ParquetObjectReader::new(store, path.into()); - let builder = ParquetRecordBatchStreamBuilder::new(object_reader).await?; - - let metadata = builder.schema().metadata().clone(); - let reader = builder.build()?; - - let arrow_schema = Arc::new(reader.schema().as_ref().clone().with_metadata(metadata)); - - let batches = reader.try_collect::>().await?; - Ok(PyTable::try_new(batches, arrow_schema)?) -} +use crate::source::FileWriter; pub(crate) struct PyWriterVersion(WriterVersion); diff --git a/arro3-io/src/runtime.rs b/arro3-io/src/runtime.rs new file mode 100644 index 0000000..6bbdda8 --- /dev/null +++ b/arro3-io/src/runtime.rs @@ -0,0 +1,38 @@ +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::sync::GILOnceCell; +use tokio::runtime::Runtime; + +static RUNTIME: GILOnceCell = GILOnceCell::new(); +static PID: GILOnceCell = GILOnceCell::new(); + +/// Construct a tokio runtime for sync requests +/// +/// This constructs a runtime with default tokio settings (e.g. [`Runtime::new`]). +/// +/// This runtime can possibly be used in the store creation process (e.g. in the AWS case, for +/// finding shared credentials), and thus any downstream applications may wish to reuse the same +/// runtime. +/// +/// Downstream consumers may explicitly want to depend on tokio and add `rt-multi-thread` as a +/// tokio feature flag to opt-in to the multi-threaded tokio runtime. +pub fn get_runtime(py: Python<'_>) -> PyResult<&'static Runtime> { + let pid = std::process::id(); + let runtime_pid = *PID.get_or_init(py, || pid); + if pid != runtime_pid { + panic!( + "Forked process detected - current PID is {} but the tokio runtime was created by {}. The tokio \ + runtime does not support forked processes https://github.com/tokio-rs/tokio/issues/4301. If you are \ + seeing this message while using Python multithreading make sure to use the `spawn` or `forkserver` \ + mode.", + pid, runtime_pid + ); + } + + let runtime = RUNTIME.get_or_try_init(py, || { + Runtime::new().map_err(|err| { + PyValueError::new_err(format!("Could not create tokio runtime. {}", err)) + }) + })?; + Ok(runtime) +} diff --git a/arro3-io/src/source/async.rs b/arro3-io/src/source/async.rs new file mode 100644 index 0000000..6d3b52c --- /dev/null +++ b/arro3-io/src/source/async.rs @@ -0,0 +1,77 @@ +use std::ops::Range; +use std::sync::Arc; + +use bytes::Bytes; +use futures::future::BoxFuture; +use parquet::arrow::async_reader::ParquetObjectReader; +use parquet::file::metadata::ParquetMetaData; + +// TODO: come back to this to implement for arbitrary async Python backends +// struct PyObspecReader(PyObject); + +// impl Clone for PyObspecReader { +// fn clone(&self) -> Self { +// Self(Python::with_gil(|py| self.0.clone_ref(py))) +// } +// } + +// impl parquet::arrow::async_reader::AsyncFileReader for PyObspecReader { +// fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, parquet::errors::Result> { +// todo!() +// } + +// fn get_byte_ranges( +// &mut self, +// ranges: Vec>, +// ) -> BoxFuture<'_, parquet::errors::Result>> { +// todo!() +// } + +// fn get_metadata(&mut self) -> BoxFuture<'_, parquet::errors::Result>> { +// Box::pin(async move { +// let file_size = self.meta.size; +// let metadata = ParquetMetaDataReader::new() +// .with_column_indexes(self.preload_column_index) +// .with_offset_indexes(self.preload_offset_index) +// .with_prefetch_hint(self.metadata_size_hint) +// .load_and_finish(self, file_size) +// .await?; +// Ok(Arc::new(metadata)) +// }) +// } +// } + +#[derive(Clone)] +pub(crate) enum AsyncReader { + // Python(PyObspecReader), + ObjectStore(ParquetObjectReader), +} + +impl parquet::arrow::async_reader::AsyncFileReader for AsyncReader { + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, parquet::errors::Result> { + match self { + // Self::Python(reader) => reader.get_bytes(range), + Self::ObjectStore(reader) => reader.get_bytes(range), + } + } + + fn get_byte_ranges( + &mut self, + ranges: Vec>, + ) -> BoxFuture<'_, parquet::errors::Result>> { + match self { + // Self::Python(reader) => reader.get_byte_ranges(ranges), + Self::ObjectStore(reader) => reader.get_byte_ranges(ranges), + } + } + + fn get_metadata<'a>( + &'a mut self, + options: Option<&'a parquet::arrow::arrow_reader::ArrowReaderOptions>, + ) -> BoxFuture<'a, parquet::errors::Result>> { + match self { + // Self::Python(reader) => reader.get_metadata(), + Self::ObjectStore(reader) => reader.get_metadata(options), + } + } +} diff --git a/arro3-io/src/source/mod.rs b/arro3-io/src/source/mod.rs new file mode 100644 index 0000000..6407643 --- /dev/null +++ b/arro3-io/src/source/mod.rs @@ -0,0 +1,5 @@ +mod r#async; +mod sync; + +pub(crate) use r#async::AsyncReader; +pub(crate) use sync::{FileWriter, SyncReader}; diff --git a/arro3-io/src/utils.rs b/arro3-io/src/source/sync.rs similarity index 72% rename from arro3-io/src/utils.rs rename to arro3-io/src/source/sync.rs index 0d1b8e0..a88d568 100644 --- a/arro3-io/src/utils.rs +++ b/arro3-io/src/source/sync.rs @@ -1,21 +1,28 @@ use bytes::Bytes; +use futures::future::BoxFuture; +use futures::FutureExt; +use parquet::arrow::async_reader::AsyncFileReader; +use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataReader}; use parquet::file::reader::{ChunkReader, Length}; +use pyo3::exceptions::PyTypeError; +use pyo3::intern; use pyo3_file::PyFileLikeObject; use pyo3::prelude::*; use std::fs::File; use std::io::{BufReader, Read, Seek, SeekFrom, Write}; use std::path::PathBuf; +use std::sync::Arc; /// Represents either a path `File` or a file-like object `FileLike` #[derive(Debug)] -pub enum FileReader { +pub enum SyncReader { File(File), FileLike(PyFileLikeObject), } -impl FileReader { - fn try_clone(&self) -> std::io::Result { +impl SyncReader { + pub(crate) fn try_clone(&self) -> std::io::Result { match self { Self::File(f) => Ok(Self::File(f.try_clone()?)), Self::FileLike(f) => Ok(Self::FileLike(f.clone())), @@ -23,13 +30,14 @@ impl FileReader { } } -impl<'py> FromPyObject<'py> for FileReader { +impl<'py> FromPyObject<'py> for SyncReader { fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { + let py = ob.py(); if let Ok(path) = ob.extract::() { Ok(Self::File(File::open(path)?)) } else if let Ok(path) = ob.extract::() { Ok(Self::File(File::open(path)?)) - } else { + } else if ob.hasattr(intern!(py, "read"))? && ob.hasattr(intern!(py, "seek"))? { Ok(Self::FileLike(PyFileLikeObject::py_with_requirements( ob.clone(), true, @@ -37,11 +45,15 @@ impl<'py> FromPyObject<'py> for FileReader { true, false, )?)) + } else { + Err(PyTypeError::new_err( + "Expected a file path or a file-like object", + )) } } } -impl Read for FileReader { +impl Read for SyncReader { fn read(&mut self, buf: &mut [u8]) -> std::io::Result { match self { Self::File(f) => f.read(buf), @@ -50,7 +62,7 @@ impl Read for FileReader { } } -impl Seek for FileReader { +impl Seek for SyncReader { fn seek(&mut self, pos: std::io::SeekFrom) -> std::io::Result { match self { Self::File(f) => f.seek(pos), @@ -59,7 +71,7 @@ impl Seek for FileReader { } } -impl Length for FileReader { +impl Length for SyncReader { fn len(&self) -> u64 { match self { Self::File(f) => f.len(), @@ -80,8 +92,8 @@ impl Length for FileReader { } } -impl ChunkReader for FileReader { - type T = BufReader; +impl ChunkReader for SyncReader { + type T = BufReader; fn get_read(&self, start: u64) -> parquet::errors::Result { let mut reader = self.try_clone()?; @@ -105,6 +117,30 @@ impl ChunkReader for FileReader { } } +// This impl allows us to use SyncReader in `ParquetFile::read_async` +impl AsyncFileReader for SyncReader { + fn get_bytes( + &mut self, + range: std::ops::Range, + ) -> BoxFuture<'_, parquet::errors::Result> { + async move { ChunkReader::get_bytes(self, range.start, (range.end - range.start) as usize) } + .boxed() + } + + fn get_metadata<'a>( + &'a mut self, + _options: Option<&'a parquet::arrow::arrow_reader::ArrowReaderOptions>, + ) -> BoxFuture<'a, parquet::errors::Result>> { + async move { + let metadata = ParquetMetaDataReader::new() + .with_page_indexes(true) + .parse_and_finish(self)?; + Ok(Arc::new(metadata)) + } + .boxed() + } +} + /// Represents either a path `File` or a file-like object `FileLike` #[derive(Debug)] pub enum FileWriter { diff --git a/docs/api/io/parquet.md b/docs/api/io/parquet.md index edd33f1..a191513 100644 --- a/docs/api/io/parquet.md +++ b/docs/api/io/parquet.md @@ -1,5 +1,20 @@ # Parquet -::: arro3.io.read_parquet + +::: arro3.io + options: + filters: + - ".*arquet.*" + + + diff --git a/mkdocs.yml b/mkdocs.yml index fbaa4fa..4b3d5be 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -48,6 +48,7 @@ watch: - docs theme: + language: en name: material palette: # Palette toggle for automatic mode diff --git a/pyproject.toml b/pyproject.toml index 95c5dc0..5e4ad19 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,17 +12,20 @@ dev-dependencies = [ "boto3>=1.35.38", "geoarrow-types>=0.2.0", "griffe-inherited-docstrings>=1.0.1", + "griffe>=1.6.0", "ipykernel>=6.29.5", "maturin>=1.7.4", "mike>=2.1.3", "mkdocs-material[imaging]>=9.6.7", "mkdocs-redirects>=1.2.2", "mkdocs>=1.6.1", - "mkdocstrings[python]>=0.28.3", + "mkdocstrings-python>=1.13.0", + "mkdocstrings>=0.27.0", "pandas>=2.2.3", "pip>=24.2", "pyarrow>=17.0.0", "pytest>=8.3.3", + "obstore>=0.6.0", ] [tool.ruff] diff --git a/uv.lock b/uv.lock index 4b220ad..fd88fb2 100644 --- a/uv.lock +++ b/uv.lock @@ -24,6 +24,7 @@ dev = [ { name = "black" }, { name = "boto3" }, { name = "geoarrow-types" }, + { name = "griffe" }, { name = "griffe-inherited-docstrings" }, { name = "ipykernel" }, { name = "maturin" }, @@ -31,7 +32,9 @@ dev = [ { name = "mkdocs" }, { name = "mkdocs-material", extra = ["imaging"] }, { name = "mkdocs-redirects" }, - { name = "mkdocstrings", extra = ["python"] }, + { name = "mkdocstrings" }, + { name = "mkdocstrings-python" }, + { name = "obstore" }, { name = "pandas" }, { name = "pip" }, { name = "pyarrow" }, @@ -45,6 +48,7 @@ dev = [ { name = "black", specifier = ">=24.10.0" }, { name = "boto3", specifier = ">=1.35.38" }, { name = "geoarrow-types", specifier = ">=0.2.0" }, + { name = "griffe", specifier = ">=1.6.0" }, { name = "griffe-inherited-docstrings", specifier = ">=1.0.1" }, { name = "ipykernel", specifier = ">=6.29.5" }, { name = "maturin", specifier = ">=1.7.4" }, @@ -52,7 +56,9 @@ dev = [ { name = "mkdocs", specifier = ">=1.6.1" }, { name = "mkdocs-material", extras = ["imaging"], specifier = ">=9.6.7" }, { name = "mkdocs-redirects", specifier = ">=1.2.2" }, - { name = "mkdocstrings", extras = ["python"], specifier = ">=0.28.3" }, + { name = "mkdocstrings", specifier = ">=0.27.0" }, + { name = "mkdocstrings-python", specifier = ">=1.13.0" }, + { name = "obstore", specifier = ">=0.6.0" }, { name = "pandas", specifier = ">=2.2.3" }, { name = "pip", specifier = ">=24.2" }, { name = "pyarrow", specifier = ">=17.0.0" }, @@ -810,11 +816,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/15/47/eb876dfd84e48f31ff60897d161b309cf6a04ca270155b0662aae562b3fb/mkdocstrings-0.29.0-py3-none-any.whl", hash = "sha256:8ea98358d2006f60befa940fdebbbc88a26b37ecbcded10be726ba359284f73d", size = 1630824 }, ] -[package.optional-dependencies] -python = [ - { name = "mkdocstrings-python" }, -] - [[package]] name = "mkdocstrings-python" version = "1.16.8" @@ -895,6 +896,55 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3e/05/eb7eec66b95cf697f08c754ef26c3549d03ebd682819f794cb039574a0a6/numpy-2.2.4-cp313-cp313t-win_amd64.whl", hash = "sha256:188dcbca89834cc2e14eb2f106c96d6d46f200fe0200310fc29089657379c58d", size = 12739119 }, ] +[[package]] +name = "obstore" +version = "0.6.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions", marker = "python_full_version < '3.12'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/80/2f/678fbcfe54381b41e079cdfff3a371fd12e678e655cd9951df7fa738accc/obstore-0.6.0-cp311-cp311-macosx_10_12_x86_64.whl", hash = "sha256:d2da9f9ed59b52ef389c71096745d6fbf5551796282b4d324517a63a94c0e64d", size = 3557198 }, + { url = "https://files.pythonhosted.org/packages/c5/77/a4bd4063bd9769f4dd713e767cce70ba5ac90428ab799fdabb84866f6ef5/obstore-0.6.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:c67ca501c5402b6cd32e5cce9f8b08552145fca4faf699729c752f56dc4ad0bd", size = 3305127 }, + { url = "https://files.pythonhosted.org/packages/d2/b5/71c5513eb306080d069879b11209e021e76d6191f02efe6c8f741af116a1/obstore-0.6.0-cp311-cp311-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:d347b56ffbabda1392f072c7dfdab2c6b73d4757590459916b02b3ab45bf72ff", size = 3441233 }, + { url = "https://files.pythonhosted.org/packages/4a/d0/d1b43469701165d9ad32911133c28a05fa3b4e6b359cc22e52617191df39/obstore-0.6.0-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c2c154a1574a6fad07c19ea48d7416e5e22e00c165f544ba8395869a88861e91", size = 3571541 }, + { url = "https://files.pythonhosted.org/packages/9c/00/4851f0fecd2f33b01cc19f8b8f5e0e28fb8a03e5dd4df75d5984115bd3ae/obstore-0.6.0-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:a8dc7d1178f38e189d2c0b5bb95e07825635e536be9003d413185530b40a62d1", size = 3772633 }, + { url = "https://files.pythonhosted.org/packages/75/15/6b89c5cf61dbe83b2305a75547249ae796075b59d5e74a7e1b894a2f2098/obstore-0.6.0-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:69523e513ba9ff72a05e1ea8ece9becbfaa3b1b0614a1c87f97a0db01ac461f3", size = 4657456 }, + { url = "https://files.pythonhosted.org/packages/4e/b9/f00ea75626eb948de1c70cc873d30cf708aee16306cd5a0919e8399c8c2e/obstore-0.6.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:960e9ee1065951d917c1a278f32f7e39333ef9f61800c18bb6a81e8cc4fb1995", size = 3676862 }, + { url = "https://files.pythonhosted.org/packages/e7/13/68c9f3468960764aeef3b477da9d2fe93ee768534073cb4441af2a7a8614/obstore-0.6.0-cp311-cp311-manylinux_2_24_aarch64.whl", hash = "sha256:7e4a901ef233e86639114303f204fa66fb8a5b4b5d5ae8761c762531e1526cd4", size = 3483716 }, + { url = "https://files.pythonhosted.org/packages/97/c9/c55007c1134eb2d745f2dc24af4b06c0f185b2db0bfaefa8e5c413594412/obstore-0.6.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:037c78c4f366c07e91f6fdc2bc83c949bf954e2ddb9706a51afa1c0f303af148", size = 3643034 }, + { url = "https://files.pythonhosted.org/packages/70/1b/62ac5907a4621f19032358f145e0bf273a95aa3d0c5824b00b03375f931a/obstore-0.6.0-cp311-cp311-musllinux_1_2_armv7l.whl", hash = "sha256:22c32e96f1fc173638fcbd8e75e61d5a2a2bc8a823e908dffa7abbf345455e9b", size = 3666755 }, + { url = "https://files.pythonhosted.org/packages/8f/3f/bd61fae0228672d4f0816bee8c92fdd2ee15f78162b788cd609e0733fb22/obstore-0.6.0-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:4bfc1e15a3cc438af6f840a73dacba1a6458755e5c3a556bf70aa60e45adffb8", size = 3650990 }, + { url = "https://files.pythonhosted.org/packages/ff/73/a80de87b19eeca28d55f15b3bc6690dd512a87e7f052cfca708d7c9ae20d/obstore-0.6.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:5adffb65220c3aa9799f0119af1f6847dbc7b1f0ae8de57580b2c1226a2f9587", size = 3849268 }, + { url = "https://files.pythonhosted.org/packages/0a/af/25e90cc26dde481ebd9450dd86a34e6ec31507d6df2ae20826b18ef470c2/obstore-0.6.0-cp311-cp311-win_amd64.whl", hash = "sha256:84b3e5e076ac4ec514a644afb914f6bb1ec55c1a558f9332896f162a139a0ffa", size = 3913027 }, + { url = "https://files.pythonhosted.org/packages/b3/9c/ef9e34967b5bb1bbbbc88241b490f1e8867490c96d022db902862c44ffad/obstore-0.6.0-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:df26abb21acb61eb3d62731c4b0dd5f89e84ec71de2463cb6b357cd60fdeca11", size = 3550189 }, + { url = "https://files.pythonhosted.org/packages/3d/cf/578dc236cfc08e3770dfac80a8692e2306f802abdc0bb59a452e8ed051a2/obstore-0.6.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:a92549e6e3f56733125d789cac65ef1e9c27b15262fd73e354636f6b1d9ea9a0", size = 3300414 }, + { url = "https://files.pythonhosted.org/packages/ec/53/cf50d1cb603202e29b404d654fa03c6375767a2bdd2361555c2e187322ed/obstore-0.6.0-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:65db8c3bee4884593629e9287028f106a7690ed88bd11b60eb892574414b6b96", size = 3444238 }, + { url = "https://files.pythonhosted.org/packages/e0/9a/53dec7fb8c099a8dca3739cf9113c4b7672aa818cd8d648bbfcf43305ae2/obstore-0.6.0-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0b2adaadd1f80aaadfb93ceaa713f890aeb4fe95f07343055eb7da70aaeff259", size = 3572839 }, + { url = "https://files.pythonhosted.org/packages/77/c7/dd569d12d3dcb541232b489bfe9409dd6bb3f131af7545f86af4a53e5570/obstore-0.6.0-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:aa0ae7b848b161f03d91c081724a3b5f2560411d3b7abbffcab76bc59ef1c265", size = 3771373 }, + { url = "https://files.pythonhosted.org/packages/65/18/659fe5d68a591b23aad4a65e36a042acde1760ef5850e45d5da00e0ad626/obstore-0.6.0-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b9c093a2bced27ab9bf190faa4703a4993506e3a18393d57a9f77b9fdc19d481", size = 4660514 }, + { url = "https://files.pythonhosted.org/packages/d8/69/e92c78faa31badf9f8cee5f5250af93a3bc5a32b8b32771699103ec34e31/obstore-0.6.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bbef509f49a7443f6acddb7313aee8ce6983d9ed189d012df0d3c07a4f45ac70", size = 3674811 }, + { url = "https://files.pythonhosted.org/packages/39/56/56ffe23a797da8e1157732eafbd6478f4c2e9f46a34926a04788b6000966/obstore-0.6.0-cp312-cp312-manylinux_2_24_aarch64.whl", hash = "sha256:18cedb9adaceed5f9bcc75f1ec4da9f8753c6eed716775bb0235f036aba6315c", size = 3478676 }, + { url = "https://files.pythonhosted.org/packages/10/a9/ffcd5de970c83cbb0456d9518737c28f06a5b012e9d6808461a904147081/obstore-0.6.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:1c10e9db12195b00c96d0f5845a78d8bb9e0f4646d446a0d31da73e64f787034", size = 3638500 }, + { url = "https://files.pythonhosted.org/packages/74/37/8c449fed33a9af87f4f06e0382ebcbecffe37dc96ff4626ba7ca86209c20/obstore-0.6.0-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:2b02b4e939b9095783a5582a947ed07687d1b3acc31b154d8e9a9ee40b4ef6fd", size = 3668122 }, + { url = "https://files.pythonhosted.org/packages/f7/54/42c5e83c1a83e60b406b40e15e3263cc667bd4e3c174b64cc0f80cbc39ca/obstore-0.6.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:c67007c226c6c0e9f3376c087c330ab76510ad11d365939d3039b32954e78cd9", size = 3651329 }, + { url = "https://files.pythonhosted.org/packages/3f/48/503745f22379ce6b62cf03a5850e42bfc2d9e94612f6755836dc36a1864a/obstore-0.6.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:d8ce9eaaacaa262c57255fd452ecfab311c6e5783632d4dd749dfb3f19b282d0", size = 3846452 }, + { url = "https://files.pythonhosted.org/packages/f5/f2/8ffcf2ea988f716d8abbe083563da1dd020830a6e831991e7c189108b0d7/obstore-0.6.0-cp312-cp312-win_amd64.whl", hash = "sha256:e97e8d03d2010979a2d773363009826bd6aea3e41cd03ed4a23d91ffd6cf861b", size = 3926101 }, + { url = "https://files.pythonhosted.org/packages/42/b5/c2c90680adc3b0b3504e13cb7625e3a0fc10cda6eb40327982fdc8461b49/obstore-0.6.0-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:49a2a8eb340a9f1d827ed16000e1f16775166e2291646c6d339dd5549461b4f5", size = 3549577 }, + { url = "https://files.pythonhosted.org/packages/d9/4a/4e9e7b11f2f0174c29614173155c8af382f8df4a0cf3fc497d07db6f6014/obstore-0.6.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:127853549da668ab89ee8d7ccb4e671faf7781e879d05deb8ca635a5ca12d929", size = 3299999 }, + { url = "https://files.pythonhosted.org/packages/c1/ea/e88079bfe991f273b22fa6bf085cfe072a90fc1162b4b5fc66ab6555ca64/obstore-0.6.0-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:52080174fc48e4b1ba8860b6a2b0d7937d903b133e42c3a3784adc744ffce9ad", size = 3443750 }, + { url = "https://files.pythonhosted.org/packages/c9/b9/8d702fb26820bef0a4c095e8b83930eec8d14e7b01c68f5213a4361e6152/obstore-0.6.0-cp313-cp313-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ea887fb0d995e500db7398288684f7fd08088f204e376afb6c9b57b0e2037f7c", size = 3572210 }, + { url = "https://files.pythonhosted.org/packages/80/48/be234b8f24be81522801514b9bbe68bb90ea8f8d9e2af9c8370f14ec04dd/obstore-0.6.0-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:bfba02d8259d4415cddfa7091b0a4f964678f54b17752374b416995c2354a8bb", size = 3770956 }, + { url = "https://files.pythonhosted.org/packages/bc/62/07954937e39b9fb022cc810cab63bcbaa4d47b708f814152b14d06dedaa1/obstore-0.6.0-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:fc2864ef21b157f9c543b2cfb3bd77b2734adfeb91b3b2875604a2f1bb51af35", size = 4650804 }, + { url = "https://files.pythonhosted.org/packages/05/9e/9a45bc0c7ca3f587516ae97be70135b2a5da80e33c316e9e1e3f6777123e/obstore-0.6.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:adf294b232927697297d7142649444e8ed74284b8d6a8b4bbba1d4b5ed726854", size = 3674650 }, + { url = "https://files.pythonhosted.org/packages/31/22/1f309fa9e4acc2270bf335c38b188e3839488c1c62e6ae9e51e7efd41c5b/obstore-0.6.0-cp313-cp313-manylinux_2_24_aarch64.whl", hash = "sha256:044d4b125a5cd63fd3fe8772d7ada7dc19b61b5763ec3d7f757a77425cab80bb", size = 3478362 }, + { url = "https://files.pythonhosted.org/packages/27/d7/68b8aabac54fd4277a88093f68658d6ab2b17a6e718429f7c77a6c1a4d54/obstore-0.6.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:fead1c226d8c4dfcaaec8441d9cf713cd46ce774561a7f12f512c60cbb319ccf", size = 3637940 }, + { url = "https://files.pythonhosted.org/packages/df/40/ed7d0024096e7eb0610cd82b716e9a234eeafaab8395d0c2c7791ee8fdc1/obstore-0.6.0-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:6bd2f93287240d6812abcfb65207b81e6fd53e1aa641d4aa638053a2e34d991d", size = 3668000 }, + { url = "https://files.pythonhosted.org/packages/20/8c/45d0f483f26774ac7dccbf28a684b97ea694862d7c29512e854bb7237391/obstore-0.6.0-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:8d6f7d607599ff9e019433736b940497f320f5cac0956a1ff360b15deed8af1f", size = 3650840 }, + { url = "https://files.pythonhosted.org/packages/bc/57/018c53c8e8f42cba8e41a8d1bab16efd1891791ba1d586ba7c6c3a4bc0cb/obstore-0.6.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:40e659ed0b1b4db447dd41b1b47f06f7ca2df1df119200ce815f7aa64616d083", size = 3846072 }, + { url = "https://files.pythonhosted.org/packages/a7/54/969edfbb288b58e096e3089d02535a83715617e62255fb83909989434b00/obstore-0.6.0-cp313-cp313-win_amd64.whl", hash = "sha256:bec5e38903ea6e8f24a327c07db45b3d3ce8979a98170e6097636bb5dfc07f75", size = 3925354 }, +] + [[package]] name = "packaging" version = "24.2"