GH-33986: [Python] Add a minimal protocol for datasets

docs/source/python/integration.rst

@@ -38,3 +38,4 @@ This allows to easily integrate PyArrow with other languages and technologies.
  integration/python_java
  integration/extending
  integration/cuda
+ integration/dataset

docs/source/python/integration/dataset.rst

@@ -0,0 +1,124 @@
+.. Licensed to the Apache Software Foundation (ASF) under one
+.. or more contributor license agreements. See the NOTICE file
+.. distributed with this work for additional information
+.. regarding copyright ownership. The ASF licenses this file
+.. to you under the Apache License, Version 2.0 (the
+.. "License"); you may not use this file except in compliance
+.. with the License. You may obtain a copy of the License at
+
+.. http://www.apache.org/licenses/LICENSE-2.0
+
+.. Unless required by applicable law or agreed to in writing,
+.. software distributed under the License is distributed on an
+.. "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+.. KIND, either express or implied. See the License for the
+.. specific language governing permissions and limitations
+.. under the License.
+
+Extending PyArrow Datasets
+==========================
+
+.. warn::
+
+ This protocol is currently experimental.
+
+PyArrow provides a core protocol for datasets, so third-party libraries can both
+produce and consume classes that conform to useful subset of the PyArrow dataset
+API. This subset provides enough functionality to provide projection
+pushdown. The subset of the API is contained in ``pyarrow.dataset.protocol``.
+
+.. image:: pyarrow_dataset_protocol.svg
+ :alt: A diagram showing the workflow for using the PyArrow Dataset protocol.
+ There are two flows shown, one for stream and one for tasks. The stream
+ case shows a linear flow from a producer class, to a dataset, to a 
+ scanner, and finally to a RecordBatchReader. The tasks case shows a
+ similar diagram, except the dataset is split into fragments, which are
+ then distributed to tasks, which each create their own scanner and
+ RecordBatchReader.
+
+Producers are responsible for outputting a class that conforms to the protocol.
+
+Consumers are responsible for calling methods on the protocol to get the data
+out of the dataset. The protocol supports getting data as a single stream or
+as a series of tasks which may be distributed.
+
+As an example, from the perspective of the user this is what the code looks like
+to retrieve a Delta Lake table as a dataset and use it in DuckDB: 
+
+.. code-block:: python
+ :emphasize-lines: 2,6
+
+ from deltalake import DeltaTable
+ table = DeltaTable("path/to/table")
+ dataset = table.to_pyarrow_dataset()
+
+ import duckdb
+ df = duckdb.arrow(dataset)
+ df.project("y")
+
+Here, the DuckDB would pass the the projection of ``y`` down to the producer
+through the dataset protocol. The deltalake scanner would then only read the
+column ``y``. Thus, the user gets to enjoy the performance benefits of pushing
+down projections while being able to specify those in their preferred query engine.
+
+
+Dataset Producers
+-----------------
+
+If you are a library implementing a new data source, you'll want to be able to
+produce a PyArrow-compatible dataset. Your dataset could be backed by the classes
+implemented in PyArrow or you could implement your own classes. Either way, you
+should implement the protocol below.
+
+Dataset Consumers
+-----------------
+
+If you are a query engine, you'll want to be able to
+consume any PyArrow datasets. To make sure your integration is compatible
+with any dataset, you should only call methods that are included in the 
+protocol. Dataset implementations provided by PyArrow implements additional
+options and methods beyond those, but they should not be relied upon without
+checking for specific classes.
+
+There are two general patterns for consuming PyArrow datasets: reading a single
+stream or creating a scan task per fragment.
+
+If you have a streaming execution model, you can receive a single stream
+of data by calling ``dataset.scanner(columns=...).to_reader()``.
+This will return a RecordBatchReader, which can be exported over the 
+:ref:`C Stream Interface <c-stream-interface>`. The record batches yield 
+from the stream can then be passed to worker threads for parallelism.
+
+If you are using a task-based model, you can split the dataset into fragments 
+and then distribute those fragments into tasks that create their own scanners
+and readers. In this case, the code looks more like:
+
+.. code-block:: python
+
+ fragments = list(dataset.get_fragments(columns=...))
+
+ def scan_partition(i):
+ fragment = fragments[i]
+ scanner = fragment.scanner()
+ return reader = scanner.to_reader()
+
+Fragments are pickleable, so they can be passed to remote workers in a 
+distributed system.
+
+If your engine supports projection (column) pushdown,
+you can pass those down to the dataset by passing them to the ``scanner``.
+Column pushdown is limited to selecting a subset of columns from the schema.
+Some implementations, including PyArrow may also support projecting and
+renaming columns, but this is not part of the protocol.
+
+
+The protocol
+------------
+
+This module can be imported starting in PyArrow ``13.0.0`` at
+``pyarrow.dataset.protocol``. The protocol is defined with ``typing.Protocol``
+classes. They can be checked at runtime with ``isinstance`` but can also be
+checked statically with Python type checkers like ``mypy``.
+
+.. literalinclude:: ../../../../python/pyarrow/dataset/protocol.py
+ :language: python

python/pyarrow/dataset/protocol.py

@@ -0,0 +1,158 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Protocol definitions for pyarrow.dataset
+
+These provide the abstract interface for a dataset. Other libraries may implement
+this interface to expose their data, without having to extend PyArrow's classes.
+
+Applications and libraries that want to consume datasets should accept datasets
+that implement these protocols, rather than requiring the specific
+PyArrow classes.
+
+The pyarrow.dataset.Dataset class itself implements this protocol.
+
+See Extending PyArrow Datasets for more information:
+
+https://arrow.apache.org/docs/python/integration/dataset.html
+"""
+from abc import abstractmethod
+from typing import Iterator, List, Optional
+
+# TODO: remove once we drop support for Python 3.7
+if sys.version_info >= (3, 8):
+ from typing import Protocol, runtime_checkable
+else:
+ from typing_extensions import Protocol, runtime_checkable
+
+from pyarrow.dataset import Expression
+from pyarrow import Table, RecordBatchReader, Schema
+
+
+@runtime_checkable
+class Scanner(Protocol):
+ """
+ A scanner implementation for a dataset.
+
+ This may be a scan of a whole dataset, or a scan of a single fragment.
+ """
+ @abstractmethod
+ def count_rows(self) -> int:
+ """
+ Count the number of rows in this dataset or fragment.
+
+ Implementors may provide optimized code paths that compute this from metadata.
+
+ Returns
+ -------
+ int
+ The number of rows in the dataset or fragment.
+ """
+ ...
+
+ @abstractmethod
+ def head(self, num_rows: int) -> Table:
+ """
+ Get the first ``num_rows`` rows of the dataset or fragment.
+
+ Parameters
+ ----------
+ num_rows : int
+ The number of rows to return.
+
+ Returns
+ -------
+ Table
+ A table containing the first ``num_rows`` rows of the dataset or fragment.
+ """
+ ...
+
+ @abstractmethod
+ def to_reader(self) -> RecordBatchReader:
+ """
+ Create a Record Batch Reader for this scan.
+
+ This is used to read the data in chunks.
+
+ Returns
+ -------
+ RecordBatchReader
+ """
+ ...
+
+
+@runtime_checkable
+class Scannable(Protocol):
+ @abstractmethod
+ def scanner(self, columns: Optional[List[str]] = None,
+ batch_size: Optional[int] = None,
+ use_threads: bool = True,
+ **kwargs) -> Scanner:
+ """Create a scanner for this dataset.
+
+ Parameters
+ ----------
+ columns : List[str], optional
+ Names of columns to include in the scan. If None, all columns are
+ included.
+ batch_size : int, optional
+ The number of rows to include in each batch. If None, the default
+ value is used. The default value is implementation specific.
+ use_threads : bool, default True
+ Whether to use multiple threads to read the rows. Often consumers
+ reading a whole dataset in one scanner will keep this
+ as True, while consumers reading a single fragment per worker will
+ set this to False.
+ """
+ ...
+
+
+@runtime_checkable
+class Fragment(Scannable, Protocol):
+ """A fragment of a dataset.
+
+ This might be a partition, a file, a file chunk, etc.
+
+ This class should be pickleable so that it can be used in a distributed scan."""
+ ...
+
+
+@runtime_checkable
+class Dataset(Scannable, Protocol):
+ @abstractmethod
+ def get_fragments(
+ self, **kwargs
+ ) -> Iterator[Fragment]:
+ """Get the fragments of this dataset.
+
+ Parameters
+ ----------
+ **kwargs : dict
+ Additional arguments to pass to underlying implementation.
+ """
+ ...
+
+ @property
+ @abstractmethod
+ def schema(self) -> Schema:
+ """
+ Get the schema of this dataset.
+
+ Returns
+ -------
+ Schema
+ """
+ ...

python/pyarrow/tests/test_dataset_protocol.py

@@ -0,0 +1,29 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Test that PyArrow datasets conform to the protocol."""
+import pyarrow.dataset.protocol as protocol
+import pyarrow.dataset as ds
+
+
+def test_dataset_protocol():
+ assert isinstance(ds.Dataset, protocol.Dataset)
+ assert isinstance(ds.Fragment, protocol.Fragment)
+
+ assert isinstance(ds.Dataset, protocol.Scannable)
+ assert isinstance(ds.Fragment, protocol.Scannable)
+
+ assert isinstance(ds.Scanner, protocol.Scanner)