Add STACSource to query STAC catalogs#1867
Conversation
Points Lumen at a SpatioTemporal Asset Catalog (STAC) API. Every cataloged collection is exposed as a selectable, queryable table via get_tables() without fetching any data. When a collection is queried it is resolved to an xarray dataset through xpystac and delegated to an XArraySQLSource, so a climate scientist goes from "download a multi-GB NetCDF, then write YAML" to pointing Lumen at a catalog. - lumen/sources/stac.py: STACSource(BaseSQLSource). Lazy catalog; resolves a collection to a delegate XArraySQLSource and forwards get/get_schema/execute to it (no duplicated SQL/schema logic); accessibility-ranked asset selection with item fallback; dataset normalization (unify_chunks, drop CF bounds dims and 0-dim grid-mapping vars); datacube cube:variables exposed as columns. - lumen/util.py: check_stac_available() mirroring check_xarray_available. - lumen/sources/xarray_sql.py: fix the global chunk spec being applied to variables of heterogeneous dimensionality; make get_schema build from a bounded sample with an O(1) metadata row count instead of COUNT(*)/DISTINCT/MIN/MAX over the full (possibly multi-TB) grid. - pyproject.toml / pixi.toml: optional stac extra; planetary-computer split into a separate stac-planetary extra. - Tests for construction, resolution, schema/metadata, asset selection, signing, and the optional-import guard. Verified against the live Planetary Computer STAC API (135 collections discovered; real daymet-daily-hi data returned). Full suite passes with no regressions.
71f9a02 to
ce2cdae
Compare
|
Thanks! This looks good. Couple questions:
|
|
Firstly thanks for the review, According to your questions :
Heads up: I have a follow-up wiring PR nearly ready, but it's a different thing. It adds |
Either way if you want lumen-stac, Happy to break it out |
| def normalize_table(self, table: str) -> str: | ||
| return table.strip('"').strip("'").strip("`") |
There was a problem hiding this comment.
Is there a reason why we override here
There was a problem hiding this comment.
I just mirrored XArraySQLSource.normalize_table _split uses the bare collection_id as a cache key, so a quoted "daymet-daily-hi:precipitation" would split to ('"daymet-daily-hi', 'precipitation"')
Will drop this override anyways!!
|
|
||
| To use this source, install the optional STAC dependencies:: | ||
|
|
||
| pip install lumen[stac] |
There was a problem hiding this comment.
| pip install lumen[stac] | |
| pip install 'lumen[stac]' |
| except ImportError as e: | ||
| raise ImportError( | ||
| "STACSource requires the 'pystac-client', 'xpystac' and 'xarray' " | ||
| "packages. Install them with: pip install lumen[stac]" |
There was a problem hiding this comment.
| "packages. Install them with: pip install lumen[stac]" | |
| "packages. Install them with: pip install 'lumen[stac]'" |
What about planetary-computer? We also need to add docs |
| "No STAC collection has been resolved yet; query a collection " | ||
| "table first so it can be opened before running raw SQL." |
There was a problem hiding this comment.
How do you query a collection table first?
I think this should be more actionable
|
In addition to the above comments,
I'd love to see this alongside this PR first before committing to merging this. Also, is adlfs needed as dependency? I tried pulling in daymet dataset and it errored. |
Docs: yep, will add |
|
Yep, both of those are bugs:
|
will tuck On adlfs: yes, daymet's COGs live on Azure Blob (abfs://), so adlfs is genuinely required. Will add it to lumen[stac-planetary] since most PC datasets are on Azure. |
Maintainer review on holoviz#1867 surfaced six items; this commit addresses all of them on the same branch so each is reviewable as one logical change. * Quote 'lumen[stac]' in the install hints at lines 12 and 47 so zsh does not eat the brackets when a user copy-pastes (inline suggestion applied verbatim). * Add adlfs to lumen[stac-planetary] so daymet and other Azure-hosted Planetary Computer collections actually resolve without a separate install step. * execute(sql) no longer requires a private _resolve() call first. Variables declared in a collection's STAC datacube extension are pre-indexed during get_tables(), and execute() parses the SQL via sqlglot (already a base dep) to auto-resolve the owning collection before dispatching. The "no collection resolved" raise becomes an actionable message naming the unknown table and listing the known collections and variables, pointing at the public get_schema/get recovery path. * get(table) now caps rows at default_limit (100k by default) by injecting SQLLimit into sql_transforms before the delegate call. default_limit=None disables the safety net; an explicit caller sql_transforms=[...] passes through unchanged. * STACCatalogControls lives at lumen.ai.controls.ingest.stac and is the first concrete subclass of CatalogSourceControls. _load_catalog enumerates collections via pystac_client; _fetch_entry returns a STACSource scoped to the picked collection through SourceResult. Inherits as_tools() so SourceAgent gets the natural-language search path for free. * Add a STAC section to docs/configuration/sources.md (install, URL form, collection:variable syntax, PC + non-PC examples, the default_limit knob, and a one-liner on STACCatalogControls). Tests: 42 (was 34 + 8 new); ruff and isort clean on the touched files.
Two follow-ups surfaced by end-to-end testing on 'lumen-ai serve --catalog https://planetarycomputer.microsoft.com/api/stac/v1': * execute(sql) now auto-rewrites bare collection-id references in the SQL to the delegate's first variable. The LLM naturally writes SELECT * FROM "daymet-daily-hi" but the underlying XArraySQLSource exposes variable tables (prcp, tmax, ...), so the rewrite lets callers query a collection by id from the chat or notebook path without learning the xarray variable naming. Wired through the same sqlglot walk that already powers the auto-resolve. * Add a 'tables' param (None default) so consumers like SQLAgent._execute_query that read source.tables directly do not AttributeError on STACSource (BaseSQLSource itself does not declare the param; STAC discovers tables dynamically via get_tables()). Mirrors XArraySQLSource's 'tables' shape. All 39 existing STAC unit tests stay green.
| source.get_tables() # list of collection IDs | ||
|
|
||
| ui = lmai.ExplorerUI(data=source) | ||
| ui.servable() |
There was a problem hiding this comment.
This doesn't really match the table I description.
get_tables doesn't get shown because it's not the last item
There was a problem hiding this comment.
Fixed: split into its own block with print(source.get_tables()) so the listing actually shows in a REPL / notebook.
| pip install 'lumen[stac]' | ||
| ``` | ||
|
|
||
| For Planetary Computer collections (Azure-hosted), also install the signing |
There was a problem hiding this comment.
What's the signing extra? Is it just stac-planetary?
There was a problem hiding this comment.
Renamed: now says "install lumen[stac-planetary] to pull in planetary-computer and adlfs" instead of vaguely calling it "the signing extra".
| table syntax: | ||
|
|
||
| ``` py | ||
| df = source.get('daymet-daily-hi:prcp') |
There was a problem hiding this comment.
Where do you get a list of the variables?
There was a problem hiding this comment.
Added a source.get_schema(collection) step before the collection:variable example so users can see what is queryable on a given collection.
| label = "STAC catalog" | ||
|
|
||
| display_columns = { | ||
| "id": {"title": "Collection", "width": "25%"}, | ||
| "title": {"title": "Title", "width": "30%"}, | ||
| "description": {"title": "Description", "width": "35%"}, | ||
| "license": {"title": "License", "width": "10%"}, | ||
| } | ||
|
|
||
| search_columns = ["id", "title", "description", "keywords"] | ||
|
|
||
| filter_columns = { | ||
| "license": { | ||
| "type": "input", "func": "like", "placeholder": "License...", | ||
| }, | ||
| } | ||
|
|
||
| detail_columns = ["description", "keywords"] |
There was a problem hiding this comment.
Should these be private? Or is that inherited from the parent class?
There was a problem hiding this comment.
Inherited. CatalogSourceControls declares display_columns, search_columns, filter_columns, and detail_columns as public param.Dict / param.List (lumen/ai/controls/ingest/catalog.py:42-55), so I'm overriding the defaults under the same public names. Keeping them public to match the base class.
| "id": collection.id, | ||
| "title": collection.title or collection.id, | ||
| # Trim long descriptions so the Tabulator row stays readable; | ||
| # the full description is available via the detail panel. | ||
| "description": (collection.description or "")[:300], | ||
| "license": collection.license or "", | ||
| "keywords": ", ".join(collection.keywords or []), |
There was a problem hiding this comment.
Let's follow PEP-8 ideally
| display_columns = { | ||
| "id": {"title": "Collection", "width": "25%"}, | ||
| "title": {"title": "Title", "width": "30%"}, | ||
| "description": {"title": "Description", "width": "35%"}, | ||
| "license": {"title": "License", "width": "10%"}, |
| # Asset media-type / role hints that indicate an xarray-openable asset. | ||
| _xarray_media_tokens: ClassVar[tuple[str, ...]] = ( | ||
| "zarr", "netcdf", "x-hdf", "kerchunk", "reference", | ||
| ) | ||
| _xarray_roles: ClassVar[frozenset[str]] = frozenset( | ||
| {"data", "zarr", "references", "index"} | ||
| ) |
There was a problem hiding this comment.
| ``` py | ||
| source = STACSource(url='https://earth-search.aws.element84.com/v1') | ||
| ``` |
There was a problem hiding this comment.
I couldn't get it working
from lumen.sources.stac import STACSource
import lumen.ai as lmai
source = STACSource(url='https://earth-search.aws.element84.com/v1')
source.get("naip")
---------------------------------------------------------------------------
AttributeError Traceback (most recent call last)
File ~/repos/lumen/lumen/sources/stac.py:250, in STACSource._resolve(self, collection_id)
249 signed = signer(obj) if signer is not None else obj
--> 250 ds = self._open_dataset(signed)
251 except Exception as e:
File ~/repos/lumen/lumen/sources/stac.py:236, in STACSource._open_dataset(self, obj)
235 kwargs["patch_url"] = self.patch_url
--> 236 return xr.open_dataset(
237 obj, engine="stac", chunks=self.chunks, **kwargs
238 )
File ~/miniconda3/envs/lumen/lib/python3.12/site-packages/xarray/backends/api.py:613, in open_dataset(filename_or_obj, engine, chunks, cache, decode_cf, mask_and_scale, decode_times, decode_timedelta, use_cftime, concat_characters, decode_coords, drop_variables, create_default_indexes, inline_array, chunked_array_type, from_array_kwargs, backend_kwargs, **kwargs)
607 backend_ds = backend.open_dataset(
608 filename_or_obj,
609 drop_variables=drop_variables,
610 **decoders,
611 **kwargs,
612 )
--> 613 ds = _dataset_from_backend_dataset(
614 backend_ds,
615 filename_or_obj,
616 engine,
617 chunks,
618 cache,
619 overwrite_encoded_chunks,
620 inline_array,
621 chunked_array_type,
622 from_array_kwargs,
623 drop_variables=drop_variables,
624 create_default_indexes=create_default_indexes,
625 **decoders,
626 **kwargs,
627 )
628 return ds
File ~/miniconda3/envs/lumen/lib/python3.12/site-packages/xarray/backends/api.py:300, in _dataset_from_backend_dataset(backend_ds, filename_or_obj, engine, chunks, cache, overwrite_encoded_chunks, inline_array, chunked_array_type, from_array_kwargs, create_default_indexes, **extra_tokens)
296 raise ValueError(
297 f"chunks must be an int, dict, 'auto', or None. Instead found {chunks}."
298 )
--> 300 _protect_dataset_variables_inplace(backend_ds, cache)
302 if create_default_indexes:
File ~/miniconda3/envs/lumen/lib/python3.12/site-packages/xarray/backends/api.py:111, in _protect_dataset_variables_inplace(dataset, cache)
110 def _protect_dataset_variables_inplace(dataset: Dataset, cache: bool) -> None:
--> 111 for name, variable in dataset.variables.items():
112 if name not in dataset._indexes:
113 # no need to protect IndexVariable objects
AttributeError: 'NoneType' object has no attribute 'variables'
The above exception was the direct cause of the following exception:
ValueError Traceback (most recent call last)
Cell In[4], line 5
1 from lumen.sources.stac import STACSource
2 import lumen.ai as lmai
3
4 source = STACSource(url='https://earth-search.aws.element84.com/v1')
----> 5 source.get("naip")
File ~/repos/lumen/lumen/sources/stac.py:326, in STACSource.get(self, table, **query)
325 def get(self, table: str, **query):
--> 326 delegate, variable = self._delegate_table(table)
327 if self.default_limit is not None and "sql_transforms" not in query:
328 query["sql_transforms"] = [SQLLimit(limit=self.default_limit)]
File ~/repos/lumen/lumen/sources/stac.py:296, in STACSource._delegate_table(self, table)
294 def _delegate_table(self, table: str) -> tuple[XArraySQLSource, str]:
295 collection_id, variable = self._split(table)
--> 296 delegate = self._resolve(collection_id)
297 if variable is None:
298 variable = delegate.get_tables()[0]
File ~/repos/lumen/lumen/sources/stac.py:259, in STACSource._resolve(self, collection_id)
257 last_error = last_error or ValueError("no data variables resolved")
258 if dataset is None:
--> 259 raise ValueError(
260 f"STAC collection {collection_id!r} could not be opened as an "
261 f"xarray dataset via xpystac ({type(last_error).__name__}: "
262 f"{last_error}). Its assets may need a reader (e.g. "
263 "stackstac/odc-stac for COGs) or storage credentials/extra "
264 "filesystem packages (e.g. adlfs) not configured here."
265 ) from last_error
266 # CF "bounds" variables (a small vertices dim like nv/bnds) are not
267 # queryable tables and break tabular chunk handling; drop them.
268 bounds_dims = [
269 d for d in dataset.dims if d in ("nv", "bnds", "bounds", "vertices")
270 ]
ValueError: STAC collection 'naip' could not be opened as an xarray dataset via xpystac (AttributeError: 'NoneType' object has no attribute 'variables'). Its assets may need a reader (e.g. stackstac/odc-stac for COGs) or storage credentials/extra filesystem packages (e.g. adlfs) not configured here.
There was a problem hiding this comment.
Fixed. _is_xarray_asset used to trust the data role even when media_type said image/tiff; profile=cloud-optimized, so COG assets reached xpystac and blew up. Tightened it so media_type is authoritative, and skipped the first-item fallback when the item has no xarray-openable asset. Your snippet now returns:
ValueError: STAC collection 'naip' has no xarray-openable asset (saw media_types: ['application/xml', 'image/tiff; application=geotiff; profile=cloud-optimized']); for COG or other image collections use stackstac or odc-stac instead.
NAIP is COG-only so xpystac can't open it regardless. For an end-to-end STACSource example use a Zarr-backed PC collection like daymet-daily-hi. Native COG support (stackstac fallback) is a planned follow-up.
ahuang11
left a comment
There was a problem hiding this comment.
I think needs cleanup, clearer instructions, and fixing tests.
* Fix the test_optional_imports.py regression the earlier 'pip install lumen[stac]' quote suggestion broke: assert on the 'lumen[<extra>]' substring instead of the full 'pip install ...' string so individual sources can shell-quote without breaking the contract. * PEP-8 sweep on lumen/ai/controls/ingest/stac.py (drop the column-aligned dict syntax in display_columns / filter_columns / the _load_catalog rows). * Fix the NAIP / COG-only collection bug from the review. STACSource(url='https://earth-search.aws.element84.com/v1').get('naip') used to raise AttributeError from inside xarray's backend because _is_xarray_asset trusted role hints ('data') even when media_type explicitly said 'image/tiff; profile=cloud-optimized', so COG assets reached xpystac and exploded. - Tighten _is_xarray_asset so a declared media_type is authoritative (role hints are a fallback only when no media_type is present). - Short-circuit _openable_candidates when collection-level assets already satisfy the request, skipping the per-item STAC API search on healthy collections. - Skip the first-item fallback for items that do not themselves carry an xarray-openable asset. - Reword the raise to name the actual observed asset media_types (no more guessing 'COG'); point at stackstac / odc-stac as the right tool for COG / image collections. - Parametrise the two negative-path tests in TestAssetSelection into a single test covering both 'no items' and 'cog-only item' cases. Verified end-to-end against the live Earth Search snippet: get('naip') now returns the actionable ValueError instead of AttributeError. * Rewrite the STAC docs section: - name 'lumen[stac-planetary]' explicitly instead of vaguely calling it 'the signing extra'; - print(source.get_tables()) in its own block so the listing is actually visible in a REPL / notebook; - show source.get_schema(collection) as the way to discover variables before reaching for 'collection:variable'.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #1867 +/- ##
==========================================
+ Coverage 70.65% 71.15% +0.50%
==========================================
Files 195 199 +4
Lines 33331 34266 +935
==========================================
+ Hits 23549 24382 +833
- Misses 9782 9884 +102 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
* chunks defaults to None instead of 'auto'. Measured ~65% faster first xr.open_dataset (no dask graph build) on PC daymet without changing any returned data. Callers who want dask-backed lazy IO on huge unbounded queries can set chunks='auto' or a per-dim dict on the source. * Sort the param block alphabetically (chunks, collections, default_limit, open_kwargs, patch_url, sql_expr, tables, url) per the holoviz cleanup checklist (https://holoviz-dev.github.io/holoviz-skills/contributing-to-holoviz/cleanup/). Doc strings, defaults, and behaviour unchanged. pixi run lint clean (ruff, isort, type-annotations-not-comments, docstring-first, etc.). 40 STAC unit tests stay green.

Description
Adds
STACSourceso Lumen can point at a SpatioTemporal Asset Catalog (STAC) API and query any cataloged collection as a table. Builds onXArraySQLSource(#1741): a collection is resolved to xarray viaxpystacand the resolved dataset is delegated to anXArraySQLSource, so there is no duplicated SQL/schema logic, it just inherits the parent.This turns the climate-scientist onboarding flow from "download a multi-GB NetCDF, then write a YAML spec" into pointing Lumen at a catalog:
get_tables()/get_metadata()stay lazy (no data fetched for discovery). Optional dependency viapip install lumen[stac];pip install lumen[stac-planetary]adds Planetary Computer asset signing.It also fixes two
XArraySQLSourcescalability issues this surfaced: a global chunk spec applied to variables of differing dimensionality, andget_schemarunningCOUNT(*)/DISTINCT/MIN/MAXover the full grid (now a bounded sample plus an O(1) metadata row count).After
Run live against the public Microsoft Planetary Computer STAC API (135 collections discovered; real
daymet-daily-hiprecipitation returned). Executed notebook with real outputs: https://gist.github.com/ghostiee-11/ee42a767740e2720e339b5cf76962654How Has This Been Tested?
34 tests covering construction, lazy listing, resolution and delegation, asset selection/signing, datacube metadata, schema, caching, and the optional-import guard, plus xarray_sql coverage for the chunk/schema fixes.
AI Disclosure
I planned the architecture and integration strategy. AI tools helped scaffold the implementation and tests based on my design decisions. I reviewed and tested all code manually.