james-willis commented on code in PR #1125: URL: https://github.com/apache/sedona-db/pull/1125#discussion_r3726820772
########## 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: yes, get_session will return None in this case. ########## 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: done -- 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]
