Streaming Arrow reader (#206)

Author: kylebarron

CHANGES.md

@@ -12,9 +12,14 @@
     specifying a mask manually for missing values in `write` (#219)
 -   Standardized 3-dimensional geometry type labels from "2.5D <type>" to
     "<type> Z" for consistency with well-known text (WKT) formats (#234)
+-   Failure error messages from GDAL are no longer printed to stderr (they were
+    already translated into Python exceptions as well) (#236).
 -   Failure and warning error messages from GDAL are no longer printed to
     stderr: failures were already translated into Python exceptions
     and warning messages are now translated into Python warnings (#236, #242).
+-   Add access to low-level pyarrow `RecordBatchReader` via
+    `pyogrio.raw.open_arrow`, which allows iterating over batches of Arrow
+    tables (#205).
 
 ### Packaging
 

docs/source/api.rst

@@ -12,3 +12,9 @@ GeoPandas integration
 
 .. autofunction:: pyogrio.read_dataframe
 .. autofunction:: pyogrio.write_dataframe
+
+Reading as Arrow data
+---------------------
+
+.. autofunction:: pyogrio.raw.read_arrow
+.. autofunction:: pyogrio.raw.open_arrow

pyogrio/_io.pyx

@@ -4,6 +4,7 @@
 """
 
 
+import contextlib
 import datetime
 import locale
 import logging
@@ -1027,8 +1028,8 @@ def ogr_read(
         field_data
     )
 
-
-def ogr_read_arrow(
+@contextlib.contextmanager
+def ogr_open_arrow(
     str path,
     dataset_kwargs,
     object layer=None,
@@ -1043,7 +1044,8 @@ def ogr_read_arrow(
     object fids=None,
     str sql=None,
     str sql_dialect=None,
-    int return_fids=False):
+    int return_fids=False,
+    int batch_size=0):
 
     cdef int err = 0
     cdef const char *path_c = NULL
@@ -1074,6 +1076,7 @@ def ogr_read_arrow(
     if sql is not None and layer is not None:
         raise ValueError("'sql' paramater cannot be combined with 'layer'")
 
+    reader = None
     try:
         dataset_options = dict_to_options(dataset_kwargs)
         ogr_dataset = ogr_open(path_c, 0, dataset_options)
@@ -1129,6 +1132,13 @@ def ogr_read_arrow(
         if not return_fids:
             options = CSLSetNameValue(options, "INCLUDE_FID", "NO")
 
+        if batch_size > 0:
+            options = CSLSetNameValue(
+                options,
+                "MAX_FEATURES_IN_BATCH",
+                str(batch_size).encode('UTF-8')
+            )
+
         # make sure layer is read from beginning
         OGR_L_ResetReading(ogr_layer)
 
@@ -1142,7 +1152,7 @@ def ogr_read_arrow(
 
         # stream has to be consumed before the Dataset is closed
         import pyarrow as pa
-        table = pa.RecordBatchStreamReader._import_from_c(stream_ptr).read_all()
+        reader = pa.RecordBatchStreamReader._import_from_c(stream_ptr)
 
         meta = {
             'crs': crs,
@@ -1152,7 +1162,12 @@ def ogr_read_arrow(
             'geometry_name': geometry_name,
         }
 
+        yield meta, reader
+
     finally:
+        if reader is not None:
+            # Mark reader as closed to prevent reading batches
+            reader.close()
 
         CSLDestroy(options)
         if fields_c != NULL:
@@ -1170,9 +1185,6 @@ def ogr_read_arrow(
             GDALClose(ogr_dataset)
             ogr_dataset = NULL
 
-    return meta, table
-
-
 def ogr_read_bounds(
     str path,
     object layer=None,

pyogrio/raw.py

@@ -6,7 +6,7 @@
 from pyogrio.util import get_vsi_path
 
 with GDALEnv():
-    from pyogrio._io import ogr_read, ogr_read_arrow, ogr_write
+    from pyogrio._io import ogr_open_arrow, ogr_read, ogr_write
     from pyogrio._ogr import (
         get_gdal_version,
         get_gdal_version_string,
@@ -179,7 +179,100 @@ def read_arrow(
     """
     Read OGR data source into a pyarrow Table.
 
-    See docstring of `read` for details.
+    See docstring of `read` for parameters.
+
+    Returns
+    -------
+    (dict, pyarrow.Table)
+
+        Returns a tuple of meta information about the data source in a dict,
+        and a pyarrow Table with data.
+
+        Meta is: {
+            "crs": "<crs>",
+            "fields": <ndarray of field names>,
+            "encoding": "<encoding>",
+            "geometry_type": "<geometry_type>",
+            "geometry_name": "<name of geometry column in arrow table>",
+        }
+    """
+    with open_arrow(
+        path_or_buffer,
+        layer=layer,
+        encoding=encoding,
+        columns=columns,
+        read_geometry=read_geometry,
+        force_2d=force_2d,
+        skip_features=skip_features,
+        max_features=max_features,
+        where=where,
+        bbox=bbox,
+        fids=fids,
+        sql=sql,
+        sql_dialect=sql_dialect,
+        return_fids=return_fids,
+        **kwargs,
+    ) as source:
+        meta, reader = source
+        table = reader.read_all()
+
+    return meta, table
+
+
+def open_arrow(
+    path_or_buffer,
+    /,
+    layer=None,
+    encoding=None,
+    columns=None,
+    read_geometry=True,
+    force_2d=False,
+    skip_features=0,
+    max_features=None,
+    where=None,
+    bbox=None,
+    fids=None,
+    sql=None,
+    sql_dialect=None,
+    return_fids=False,
+    batch_size=65_536,
+    **kwargs,
+):
+    """
+    Open OGR data source as a stream of pyarrow record batches.
+
+    See docstring of `read` for parameters.
+
+    The RecordBatchStreamReader is reading from a stream provided by OGR and must not be
+    accessed after the OGR dataset has been closed, i.e. after the context manager has
+    been closed.
+
+    Examples
+    --------
+
+    >>> from pyogrio.raw import open_arrow
+    >>> import pyarrow as pa
+    >>> import shapely
+    >>>
+    >>> with open_arrow(path) as source:
+    >>>     meta, reader = source
+    >>>     for table in reader:
+    >>>         geometries = shapely.from_wkb(table[meta["geometry_name"]])
+
+    Returns
+    -------
+    (dict, pyarrow.RecordBatchStreamReader)
+
+        Returns a tuple of meta information about the data source in a dict,
+        and a pyarrow RecordBatchStreamReader with data.
+
+        Meta is: {
+            "crs": "<crs>",
+            "fields": <ndarray of field names>,
+            "encoding": "<encoding>",
+            "geometry_type": "<geometry_type>",
+            "geometry_name": "<name of geometry column in arrow table>",
+        }
     """
     try:
         import pyarrow  # noqa
@@ -191,7 +284,7 @@ def read_arrow(
     dataset_kwargs = _preprocess_options_key_value(kwargs) if kwargs else {}
 
     try:
-        result = ogr_read_arrow(
+        return ogr_open_arrow(
             path,
             layer=layer,
             encoding=encoding,
@@ -207,13 +300,12 @@ def read_arrow(
             sql_dialect=sql_dialect,
             return_fids=return_fids,
             dataset_kwargs=dataset_kwargs,
+            batch_size=batch_size,
         )
     finally:
         if buffer is not None:
             remove_virtual_file(path)
 
-    return result
-
 
 def detect_driver(path):
     # try to infer driver from path

pyogrio/tests/test_arrow.py

@@ -1,7 +1,9 @@
+import math
+
 import pytest
 
 from pyogrio import __gdal_version__, read_dataframe
-from pyogrio.raw import read_arrow
+from pyogrio.raw import open_arrow, read_arrow
 
 try:
     import pandas as pd
@@ -74,3 +76,27 @@ def test_read_arrow_raw(naturalearth_lowres):
     meta, table = read_arrow(naturalearth_lowres)
     assert isinstance(meta, dict)
     assert isinstance(table, pyarrow.Table)
+
+
+def test_open_arrow(naturalearth_lowres):
+    with open_arrow(naturalearth_lowres) as (meta, reader):
+        assert isinstance(meta, dict)
+        assert isinstance(reader, pyarrow.RecordBatchReader)
+        assert isinstance(reader.read_all(), pyarrow.Table)
+
+
+def test_open_arrow_batch_size(naturalearth_lowres):
+    meta, table = read_arrow(naturalearth_lowres)
+    batch_size = math.ceil(len(table) / 2)
+
+    with open_arrow(naturalearth_lowres, batch_size=batch_size) as (meta, reader):
+        assert isinstance(meta, dict)
+        assert isinstance(reader, pyarrow.RecordBatchReader)
+        count = 0
+        tables = []
+        for table in reader:
+            tables.append(table)
+            count += 1
+
+        assert count == 2, "Should be two batches given the batch_size parameter"
+        assert len(tables[0]) == batch_size, "First table should match the batch size"