paleolimbot commented on code in PR #1125:
URL: https://github.com/apache/sedona-db/pull/1125#discussion_r3725135994
##########
python/sedonadb-zarr/tests/test_zarr.py:
##########
@@ -65,6 +65,66 @@ def test_format_spec_via_read(zarr_group):
)
+def test_zarr_url_as_table(tmp_path):
+ """`SELECT * FROM '<.zarr url>'` reads a Zarr group with no explicit
format.
+
+ This is the motivating case for the URL-as-table resolver: a bare
+ `FROM '<url>'` keys off the `.zarr` extension and routes the directory
+ through the single-object table path, yielding the same rasters as the
+ explicit `read(url, format="zarr")`. DataFusion's default resolver instead
+ lists the directory's inner chunks and tries to parse one as the wrong
+ format, failing with `Json error: Not valid JSON: EOF ...`.
+ """
+ zarr = pytest.importorskip("zarr", minversion="3.0")
+ np = pytest.importorskip("numpy")
+ pytest.importorskip("sedonadb_zarr")
+
+ # The resolver keys off the path extension, so the group must live at a
+ # `.zarr`-suffixed directory (pytest's `tmp_path` itself has no suffix).
+ # Same 2x2 UInt8 / (1, 2)-chunk layout as the `zarr_group` fixture, so it
+ # reads as two OutDb raster rows (one per chunk).
+ zarr_path = tmp_path / "temperature.zarr"
+ root = zarr.open_group(str(zarr_path), mode="w")
+ arr = root.create_array(
+ "temperature",
+ shape=(2, 2),
+ chunks=(1, 2),
+ dtype="uint8",
+ dimension_names=["y", "x"],
+ )
+ arr[:] = np.array([[10, 11], [20, 21]], dtype=np.uint8)
+
+ con = sedonadb.connect()
+ con.register(sedonadb_zarr.ZarrExtension())
+
+ url = zarr_path.as_uri()
+
+ # Ground truth: the working explicit-format read of the same `.zarr` group.
+ expected = con.read(url, format="zarr").to_arrow_table()
+
+ def assert_matches_expected(table):
+ assert table.num_rows == expected.num_rows == 2
+ assert table.column_names == ["raster"]
+ # The column carries the RASTER extension type (rows round-trip to
+ # `Raster`), matching the explicit read.
+ assert table.schema.field("raster").type ==
expected.schema.field("raster").type
+ got = [table["raster"][i].as_py() for i in range(table.num_rows)]
+ exp = [expected["raster"][i].as_py() for i in range(expected.num_rows)]
+ assert all(isinstance(r, Raster) for r in got)
+ # Same georeferencing and OutDb chunk anchors as the explicit read
+ # proves the URL resolved to this group, not a listing over its chunks.
+ assert [r.transform for r in got] == [r.transform for r in exp]
+ assert [r.bands[0].outdb_uri for r in got] == [
+ r.bands[0].outdb_uri for r in exp
+ ]
Review Comment:
Probably just checking the number of rows and the fact that it executed is
sufficient here
##########
docs/reference/python-zarr.md:
##########
@@ -31,6 +31,18 @@ sd.register(sedonadb_zarr.ZarrExtension())
sd.read("file:///path/to/foo.zarr").show()
```
+Once the extension is registered, you can also query a Zarr group directly in
+SQL by its URL, without an explicit format:
+
+```python
+sd.sql("SELECT * FROM 'file:///path/to/foo.zarr'").show()
+```
Review Comment:
Should this highlight `'path/to/foo.zarr'`, since that's what you'd normally
type?
##########
rust/sedona-datasource/src/url_table.rs:
##########
@@ -0,0 +1,218 @@
+// 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.
+
+//! URL-as-table resolution for directory-shaped external formats.
+//!
+//! **Experimental**: the shape of this hook may change.
+//!
+//! DataFusion's [`enable_url_table`](SessionContext::enable_url_table)
+//! installs a resolver (`DynamicListTableFactory`) that always builds a
+//! `ListingTable` for a bare `FROM '<url>'`: it lists the files under the
+//! URL prefix, takes the first object it finds, and picks a file format by
+//! *that object's* extension. For a directory-shaped format like Zarr —
+//! where the "table" is the `.zarr` directory itself, not the files within
+//! it — this lists the directory contents and tries to parse an inner chunk
+//! (e.g. `zarr.json` or a raw binary chunk) as the wrong format, which fails.
+//!
+//! [`enable_sedona_url_table`] installs [`SedonaUrlTableFactory`] instead. It
+//! matches the URL's extension against the session's registered
+//! [`ExternalFormatSpec`](crate::spec::ExternalFormatSpec)s: when the match
+//! is a directory-shaped format
+//!
([`list_single_object`](crate::spec::ExternalFormatSpec::list_single_object)
+//! `== true`), it builds a
+//! [`SingleObjectExternalTable`](crate::provider::SingleObjectExternalTable)
+//! (via [`external_table`]) that passes the URL through untouched. Everything
+//! else — the file-shaped formats DataFusion already handles (GeoParquet,
+//! CSV, ...) — delegates to DataFusion's default resolver unchanged.
+
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use datafusion::{
+ catalog::TableProvider,
+ datasource::{dynamic_file::DynamicListTableFactory,
listing::ListingTableUrl},
+ execution::SessionState,
+ prelude::SessionContext,
+};
+use datafusion_catalog::{DynamicFileCatalog, UrlTableFactory};
+use datafusion_common::Result;
+use datafusion_session::SessionStore;
+
+use crate::{format::ExternalFormatFactory, provider::external_table};
+
+/// Install SedonaDB's URL-as-table resolver on `ctx`.
+///
+/// Drop-in replacement for DataFusion's
+/// [`SessionContext::enable_url_table`] that additionally routes
+/// directory-shaped external formats through the single-object table
+/// path. Mirrors `enable_url_table`'s wiring: it wraps the current catalog
+/// list in a [`DynamicFileCatalog`] backed by a [`SedonaUrlTableFactory`],
+/// then points the factory's session store at the (unchanged) session
+/// state so it can resolve registered file formats at query time.
+///
+/// **Experimental.**
+pub fn enable_sedona_url_table(ctx: SessionContext) -> SessionContext {
+ let factory = Arc::new(SedonaUrlTableFactory::new());
+ let current_catalog_list = ctx.state().catalog_list().clone();
+ let catalog_list = Arc::new(DynamicFileCatalog::new(
+ current_catalog_list,
+ Arc::clone(&factory) as Arc<dyn UrlTableFactory>,
+ ));
+ ctx.register_catalog_list(catalog_list);
+ factory.session_store().with_state(ctx.state_weak_ref());
+ ctx
+}
+
+/// [`UrlTableFactory`] that pre-routes directory-shaped external formats to
+/// the single-object table path and delegates everything else to
+/// DataFusion's default [`DynamicListTableFactory`].
+///
+/// **Experimental.**
+#[derive(Debug)]
+pub struct SedonaUrlTableFactory {
+ /// DataFusion's default resolver, used for every URL that does not
+ /// resolve to a registered directory-shaped format. Owns the
+ /// [`SessionStore`] that both it and our routing logic read the live
+ /// [`SessionState`] from.
+ inner: DynamicListTableFactory,
+}
+
+impl SedonaUrlTableFactory {
+ /// Create a factory with a fresh [`SessionStore`]. Wire the store to a
+ /// session with [`SessionStore::with_state`] (done for you by
+ /// [`enable_sedona_url_table`]) before resolving any URL.
+ pub fn new() -> Self {
+ Self {
+ inner: DynamicListTableFactory::new(SessionStore::new()),
+ }
+ }
+
+ /// The [`SessionStore`] shared by the routing logic and the delegated
+ /// [`DynamicListTableFactory`].
+ pub fn session_store(&self) -> &SessionStore {
+ self.inner.session_store()
+ }
+
+ /// Resolve the current [`SessionState`] from the session store, or
+ /// `None` if the session has gone away (in which case the caller falls
+ /// back to the default resolver, which surfaces the canonical error).
+ fn session_state(&self) -> Option<SessionState> {
+ self.session_store()
+ .get_session()
+ .upgrade()
+ .and_then(|session| {
+ session
+ .read()
+ .as_any()
+ .downcast_ref::<SessionState>()
+ .cloned()
+ })
+ }
Review Comment:
It seems like this is handled here, but note that anywhere the `Session` is
downcast into a `SessionState` could fail (usually the culprit is the
FFI_Session or a testing session).
##########
rust/sedona-datasource/src/url_table.rs:
##########
@@ -0,0 +1,218 @@
+// 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.
+
+//! URL-as-table resolution for directory-shaped external formats.
+//!
+//! **Experimental**: the shape of this hook may change.
+//!
+//! DataFusion's [`enable_url_table`](SessionContext::enable_url_table)
+//! installs a resolver (`DynamicListTableFactory`) that always builds a
+//! `ListingTable` for a bare `FROM '<url>'`: it lists the files under the
+//! URL prefix, takes the first object it finds, and picks a file format by
+//! *that object's* extension. For a directory-shaped format like Zarr —
+//! where the "table" is the `.zarr` directory itself, not the files within
+//! it — this lists the directory contents and tries to parse an inner chunk
+//! (e.g. `zarr.json` or a raw binary chunk) as the wrong format, which fails.
+//!
+//! [`enable_sedona_url_table`] installs [`SedonaUrlTableFactory`] instead. It
+//! matches the URL's extension against the session's registered
+//! [`ExternalFormatSpec`](crate::spec::ExternalFormatSpec)s: when the match
+//! is a directory-shaped format
+//!
([`list_single_object`](crate::spec::ExternalFormatSpec::list_single_object)
+//! `== true`), it builds a
+//! [`SingleObjectExternalTable`](crate::provider::SingleObjectExternalTable)
+//! (via [`external_table`]) that passes the URL through untouched. Everything
+//! else — the file-shaped formats DataFusion already handles (GeoParquet,
+//! CSV, ...) — delegates to DataFusion's default resolver unchanged.
+
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use datafusion::{
+ catalog::TableProvider,
+ datasource::{dynamic_file::DynamicListTableFactory,
listing::ListingTableUrl},
+ execution::SessionState,
+ prelude::SessionContext,
+};
+use datafusion_catalog::{DynamicFileCatalog, UrlTableFactory};
+use datafusion_common::Result;
+use datafusion_session::SessionStore;
+
+use crate::{format::ExternalFormatFactory, provider::external_table};
+
+/// Install SedonaDB's URL-as-table resolver on `ctx`.
+///
+/// Drop-in replacement for DataFusion's
+/// [`SessionContext::enable_url_table`] that additionally routes
+/// directory-shaped external formats through the single-object table
+/// path. Mirrors `enable_url_table`'s wiring: it wraps the current catalog
+/// list in a [`DynamicFileCatalog`] backed by a [`SedonaUrlTableFactory`],
+/// then points the factory's session store at the (unchanged) session
+/// state so it can resolve registered file formats at query time.
+///
+/// **Experimental.**
+pub fn enable_sedona_url_table(ctx: SessionContext) -> SessionContext {
Review Comment:
This file should probably go in the `sedona` crate, where we have other
catalog/url stuff. I haven't been perfect about this but I would like to remove
the `datafusion` dependency here (a lot of upstream work has been done to make
datafusion-datasource *not* depend on `datafusion` to reduce the size of the
binaries and compilation time for extension formats).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
