refactor(python/sedonadb): Extract common read functionality into read accessor (#968)

Author: paleolimbot

python/sedonadb-zarr/README.md

@@ -22,15 +22,16 @@
 Zarr support for [SedonaDB](https://sedona.apache.org/) as an opt-in plugin package. Reads Zarr v3 groups (with sharding, vlen-utf8 dims, etc.) as a column of N-D rasters:
 
 ```python
-import sedonadb
-import sedonadb_zarr
+import sedona.db
+from sedonadb_zarr import ZarrExtension
 
-con = sedonadb.connect()
-con.read_format(sedonadb_zarr.ZarrFormatSpec(), "file:///path/to/foo.zarr").show()
+sd = sedona.db.connect()
+sd.register(ZarrExtension())
+sd.read("file:///path/to/foo.zarr").show()
 ```
 
 The main `sedonadb` package does not bundle Zarr support — applications that don't import `sedonadb_zarr` pay no runtime cost.
 
 ## Architecture
 
-A maturin-built mixed Rust/Python package. The Rust side is a thin PyO3 shim around `sedona-raster-zarr` exposing `PyZarrChunkReader` (implementing `__arrow_c_stream__`). The Python side defines `ZarrFormatSpec(ExternalFormatSpec)`, which sedonadb consumes via `con.read_format(spec, uri)`. The same plugin shape applies to future formats (`sedonadb-cog`, `sedonadb-icechunk`, …).
+A maturin-built mixed Rust/Python package. The Rust side is a thin PyO3 shim around `sedona-raster-zarr` exposing `PyZarrChunkReader` (implementing `__arrow_c_stream__`).

python/sedonadb-zarr/python/sedonadb_zarr/__init__.py

@@ -18,10 +18,12 @@
 """Zarr support for SedonaDB.
 
 ```python
-from sedonadb_zarr import Zarr
+import sedona.db
+from sedonadb_zarr import ZarrExtension
 
 sd = sedona.db.connect()
-sd.read_format(Zarr(), "file:///path/to/foo.zarr").show()
+sd.register(ZarrExtension())
+sd.read("file:///path/to/foo.zarr").show()
 ```
 
 Importing `sedonadb_zarr` is opt-in — applications that don't import
@@ -53,18 +55,19 @@ def __sedonadb_extension__(self, ctx: SedonaContext, **kwargs) -> None:
         if kwargs:
             raise ValueError("Registration options not supported for ZarrExtension")
 
-        # Register the Zarr() format as a FileFormatFactory for SQL support
+        # Register the Zarr() format as a FileFormatFactory for SQL and .read(..., format="zarr")
         ctx.register(Zarr())
 
 
 class Zarr(ExternalFormatSpec):
     """`ExternalFormatSpec` for Zarr groups.
 
     This is registered automatically when registering the module with a
-    SedonaContext. Use with `sd.read_format(spec, uri)`:
+    SedonaContext. Use with `sd.read(uri, format=Zarr())` or format="zarr"
+    after registering the extension:
 
     ```python
-    sd.read_format(Zarr(), "file:///path/to/foo.zarr")
+    sd.read("file:///path/to/foo.zarr", format="zarr")
     ```
 
     Args:

python/sedonadb-zarr/tests/conftest.py

@@ -51,7 +51,8 @@ def raster_con(zarr_group):
     Uses its own connection (not the shared module-level one) so the view
     doesn't leak into other tests.
     """
-    con = sedonadb.connect()
-    df = con.read_format(sedonadb_zarr.Zarr(), f"file://{zarr_group}")
+    sd = sedonadb.connect()
+    sd.register(sedonadb_zarr.ZarrExtension())
+    df = sd.read(f"file://{zarr_group}", format="zarr")
     df.to_view("rasters")
-    return con
+    return sd

python/sedonadb-zarr/tests/test_zarr.py

@@ -15,13 +15,6 @@
 # specific language governing permissions and limitations
 # under the License.
 
-"""Tests for the `sedonadb-zarr` plugin.
-
-Plugin surface: `ZarrFormatSpec(ExternalFormatSpec)` paired with
-`con.read_format(spec, uri)`. The SQL UDTF form (`sd_read_zarr`) is
-deferred to a follow-up PR.
-"""
-
 import numpy as np
 import pytest
 import sedonadb
@@ -48,9 +41,9 @@ def zarr_group(tmp_path):
     return tmp_path
 
 
-def test_format_spec_via_read_format(zarr_group):
+def test_format_spec_via_read(zarr_group):
     con = sedonadb.connect()
-    df = con.read_format(sedonadb_zarr.Zarr(), f"file://{zarr_group}")
+    df = con.read(f"file://{zarr_group}", format=sedonadb_zarr.Zarr())
     arrow_tab = df.to_arrow_table()
     assert arrow_tab.num_rows == 2
     assert arrow_tab.column_names == ["raster"]
@@ -99,7 +92,7 @@ def _zarr_with_attrs(tmp_path, group_attrs, *, dims=("y", "x")):
 def _envelope_bounds(con, path):
     """Read the single-chunk zarr and return RS_Envelope bounds of row 0."""
     shapely = pytest.importorskip("shapely")
-    df = con.read_format(sedonadb_zarr.Zarr(), f"file://{path}")
+    df = con.read(f"file://{path}", format=sedonadb_zarr.Zarr())
     raster = df.to_arrow_table()["raster"][0].as_py()
     wkt = (
         con.sql("SELECT ST_AsText(RS_Envelope($1)) AS wkt", params=(raster,))
@@ -182,7 +175,17 @@ def test_rs_envelope_honors_skew(tmp_path):
 def test_format_spec_with_arrays_option(zarr_group):
     con = sedonadb.connect()
     spec = sedonadb_zarr.Zarr().with_options({"arrays": ["temperature"]})
-    df = con.read_format(spec, f"file://{zarr_group}")
+
+    # Check via constructor options
+    df = con.read(f"file://{zarr_group}", format=spec)
+    assert df.to_arrow_table().num_rows == 2
+
+    # Check via read(..., options={...})
+    df = con.read(
+        f"file://{zarr_group}",
+        options={"arrays": ["temperature"]},
+        format=sedonadb_zarr.Zarr(),
+    )
     assert df.to_arrow_table().num_rows == 2
 
 
@@ -225,7 +228,7 @@ def test_dtype_mapping_roundtrips(tmp_path, numpy_dtype):
     arr[:] = np.ones((2, 2), dtype=numpy_dtype)
 
     con = sedonadb.connect()
-    df = con.read_format(sedonadb_zarr.Zarr(), f"file://{tmp_path}")
+    df = con.read(f"file://{tmp_path}", format=sedonadb_zarr.Zarr())
     tab = df.to_arrow_table()
     assert tab.num_rows == 2
 

python/sedonadb/python/sedonadb/context.py

@@ -15,7 +15,6 @@
 # specific language governing permissions and limitations
 # under the License.
 
-import json
 import os
 import sys
 from functools import cached_property
@@ -33,7 +32,7 @@
 )
 
 if TYPE_CHECKING:
-    from sedonadb.datasource import ExternalFormatSpec
+    from sedonadb.read import Read
 
 from sedonadb._lib import (
     InternalContext,
@@ -44,6 +43,7 @@
 from sedonadb._options import Options
 from sedonadb.dataframe import DataFrame, _create_data_frame
 from sedonadb.functions import Functions
+
 from sedonadb.expr.expression import (
     Expr,
     col as col_expr,
@@ -193,6 +193,12 @@ def drop_view(self, name: str) -> None:
         """
         self._impl.drop_view(name)
 
+    @cached_property
+    def read(self) -> "Read":
+        from sedonadb.read import Read
+
+        return Read(self)
+
     def read_parquet(
         self,
         table_paths: Union[str, Path, Iterable[str]],
@@ -275,27 +281,12 @@ def read_parquet(
             >>> sd.read_parquet(url)
             <sedonadb.dataframe.DataFrame object at ...>
         """
-        if isinstance(table_paths, (str, Path)):
-            table_paths = [table_paths]
-
-        if options is None:
-            options = {}
-
-        if geometry_columns is not None and not isinstance(geometry_columns, str):
-            geometry_columns = json.dumps(geometry_columns)
-
-        if isinstance(partitioning, str):
-            partitioning = [partitioning]
-
-        return DataFrame(
-            self,
-            self._impl.read_parquet(
-                [str(path) for path in table_paths],
-                options,
-                geometry_columns,
-                validate,
-                None if partitioning is None else list(partitioning),
-            ),
+        return self.read.parquet(
+            table_paths,
+            options=options,
+            geometry_columns=geometry_columns,
+            validate=validate,
+            partitioning=partitioning,
         )
 
     def read_pyogrio(
@@ -362,83 +353,8 @@ def read_pyogrio(
             └──────────────┘
 
         """
-        from sedonadb.datasource import PyogrioFormatSpec
-
-        if isinstance(table_paths, (str, Path)):
-            table_paths = [table_paths]
-
-        spec = PyogrioFormatSpec(extension)
-        if options is not None:
-            spec = spec.with_options(options)
-
-        if isinstance(partitioning, str):
-            partitioning = [partitioning]
-
-        return DataFrame(
-            self,
-            self._impl.read_external_format(
-                spec,
-                [str(path) for path in table_paths],
-                False,
-                None if partitioning is None else list(partitioning),
-            ),
-        )
-
-    def read_format(
-        self,
-        spec: "ExternalFormatSpec",
-        table_paths: Union[str, Path, Iterable[str]],
-        check_extension: bool = False,
-        partitioning: Union[str, Iterable[str], None] = None,
-    ) -> DataFrame:
-        """Read one or more paths using a Python-defined `ExternalFormatSpec`.
-
-        This is the plugin entry point: a format-specific package (e.g.
-        `sedonadb-zarr`) defines an `ExternalFormatSpec` subclass and the
-        user reads through it via this method. Built-in formats have
-        their own dedicated readers (`read_parquet`, `read_pyogrio`).
-
-        Format-specific options are passed via the spec itself using
-        `spec.with_options({...})`, which returns a configured copy.
-        Unlike `read_pyogrio`, this method has no `options=` keyword —
-        each spec class documents its own supported keys.
-
-        Args:
-            spec: An `ExternalFormatSpec` instance describing how to open
-                the underlying source.
-            table_paths: A str, Path, or iterable of paths/URLs.
-            check_extension: When `True`, error if a non-collection path
-                doesn't end in the spec's `extension`. Defaults to `False`.
-            partitioning:
-                Optional list of column names for hive-style partitioning. When reading
-                from a directory with paths like `/col=value/file.ext`, partition
-                column names are auto-discovered by default (`partitioning=None`).
-                Explicitly specify column names (e.g., `["col"]`) to override
-                auto-discovery, or pass an empty list `[]` to disable partitioning
-                entirely.
-
-        Examples:
-            >>> import sedonadb_zarr  # doctest: +SKIP
-            >>> sd = sedona.db.connect()
-            >>> spec = sedonadb_zarr.ZarrFormatSpec().with_options(  # doctest: +SKIP
-            ...     {"arrays": ["temperature"]}
-            ... )
-            >>> sd.read_format(spec, "file:///path/to/foo.zarr").show()  # doctest: +SKIP
-        """
-        if isinstance(table_paths, (str, Path)):
-            table_paths = [table_paths]
-
-        if isinstance(partitioning, str):
-            partitioning = [partitioning]
-
-        return DataFrame(
-            self,
-            self._impl.read_external_format(
-                spec,
-                [str(path) for path in table_paths],
-                check_extension,
-                None if partitioning is None else list(partitioning),
-            ),
+        return self.read.pyogrio(
+            table_paths, options=options, extension=extension, partitioning=partitioning
         )
 
     def sql(
@@ -543,6 +459,11 @@ def register(self, component: Any, **kwargs: Any) -> None:
             component.__sedonadb_extension__(self, **kwargs)
             return
 
+        # If this is an external format, register it so that sd.read(..., format="ext")
+        # works
+        if hasattr(component, "__sedonadb_external_format__") and component.extension:
+            self.read._register_external_format(component.extension, component)
+
         supported_interfaces = (
             "__sedonadb_internal_udf__",
             "__sedonadb_internal_aggregate_udf__",