Add XArraySQLSource for N-dimensional scientific data

[ghostiee-11]

Description

Adds XArraySQLSource so Lumen can query xarray datasets (NetCDF, Zarr, HDF5, GRIB) with SQL via Apache DataFusion. Builds on the proof-of-concept from #1434 and addresses #1508.

Each data variable in the dataset becomes a SQL table with coordinate columns. For example, a NetCDF file with an air variable on (time, lat, lon) dimensions gives you a table you can query like:

source = XArraySQLSource(uri="air_temperature.nc")
source.execute("SELECT lat, AVG(air) FROM air GROUP BY lat")

Inherits from BaseSQLSource to reuse schema inference, caching, SQL transforms, and create_sql_expr_source. Optional dependency via pip install lumen[xarray].

After

Tested with NOAA air temperature dataset (3.8M rows):

How Has This Been Tested?

58 tests covering construction (from dataset, file, zarr), SQL execution, schema/metadata, table normalization, serialization roundtrip, async, and resource cleanup.

pytest lumen/tests/sources/test_xarray_sql.py -v

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.

[Copilot]

Pull request overview

Adds a new optional xarray-backed SQL Source implementation to Lumen, enabling SQL queries (via xarray-sql/DataFusion) over N-dimensional scientific datasets and exposing each data variable as a SQL table.

Changes:

• Added XArraySQLSource to query xarray datasets with DataFusion SQL, plus schema/metadata/dimension-info helpers and async query APIs.
• Added a new optional dependency group (lumen[xarray]) for installing xarray + xarray-sql related packages.
• Added a dedicated test suite for XArraySQLSource and wired an optional import into lumen.sources.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 13 comments.

File	Description
pyproject.toml	Adds xarray optional dependency group for the new source.
lumen/sources/__init__.py	Registers XArraySQLSource via optional import (graceful when deps missing).
lumen/sources/xarray_sql.py	Implements XArraySQLSource (DataFusion execution, schema/metadata, async APIs).
lumen/tests/sources/test_xarray_sql.py	Adds tests for construction, querying, schema/metadata, serialization, async.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[codecov]

Codecov Report

❌ Patch coverage is 95.76108% with 22 lines in your changes missing coverage. Please review.
✅ Project coverage is 69.75%. Comparing base (88a9a82) to head (71040e0).
⚠️ Report is 9 commits behind head on main.

Files with missing lines	Patch %	Lines
lumen/tests/sources/test_xarray_sql.py	96.59%	12 Missing ⚠️
lumen/sources/xarray_sql.py	94.01%	10 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1741      +/-   ##
==========================================
+ Coverage   68.75%   69.75%   +1.00%     
==========================================
  Files         171      173       +2     
  Lines       28979    29770     +791     
==========================================
+ Hits        19924    20767     +843     
+ Misses       9055     9003      -52     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[ahuang11]

Didn't read through it completely; I think there's a lot of duplication instead of just properly inheriting from parent.

[ghostiee-11]

Thanks for the review @ahuang11! You were absolutely right -- there was way too much duplication. I went through BaseSQLSource and DuckDBSource carefully and refactored to properly inherit from the parent.

What I removed (334 lines net deleted):

• _sql_literal() / _quote_identifier() -- these reimplemented what SQLFilter, SQLDistinct, and SQLMinMax already do via sqlglot. Deleted entirely.
• get_sql_expr() -- added a sql_expr = param.String(default='SELECT * FROM {table}') param (same pattern as DuckDBSource), so the parent's get_sql_expr() with SQLSelectFrom works directly.
• get() with manual WHERE building -- replaced with the DuckDBSource pattern: pop __dask, build conditions list, pass to SQLFilter(conditions=conditions, read=self.dialect). Had to add slice-to-tuple conversion since SQLFilter expects tuples for BETWEEN, and dunder key filtering since Lumen passes internal keys like __limit__ through query params.
• execute_async() / get_async() -- parent's asyncio.to_thread(self.execute, ...) is identical. Deleted.
• to_spec() / from_spec() -- parent's generic param serialization handles this. Deleted.
• estimate_size() -- not part of the BaseSQLSource API, no other source has it. Removed.

The one override I kept: get_schema()

DataFusion does not support TABLESAMPLE SQL syntax, which BaseSQLSource.get_schema uses via SQLSample when shuffle=True. The override uses SQLLimit instead -- but otherwise follows the parent's logic exactly, using SQLCount, SQLDistinct, and SQLMinMax transforms instead of raw SQL strings.

Test changes: removed test classes for deleted methods (TestEstimateSize, TestSQLLiteral, test_get_sql_expr_dict_input). All 52 remaining tests pass.

Let me know if there's anything else to clean up!

[ghostiee-11]

Hey @ahuang11, pushed a follow-up fix (f812387) addressing the duplication feedback:

• get_schema(): replaced 70-line copy-paste with 3-line super() delegation (forces shuffle=False since DataFusion lacks TABLESAMPLE)
• execute(): removed broken repr() param substitution
• get(): simplified to match DuckDBSource pattern
• Removed unused constants and logging

-101 lines, all 58 tests pass locally.

I also noticed a few things I could harden further (like wiring param_values support into execute(), adding file validation in __init__, and capping chunk sizes to prevent OOM on large datasets). Would you prefer I address those in this PR or keep this focused and do them as follow-ups?

[ahuang11]

Thanks I'll review when time allows.

In the future, there's no need to tag me or other maintainers. We'll review as soon as we can. For reference: https://holoviz.org/contribute.html#using-ai-readme

[ghostiee-11]

Thanks for the heads up, noted! I am really sorry, won't tag in the future. And thanks for the link.

[ahuang11]

Thanks for rebuilding my draft PR; interestingly I lost xarray.py somewhere.

I think what you have is in a decent state. I'd like to see how this works on a variety of datasets and edge cases, including large datasets that are go out of memory, xqlsystems/xarray-sql#93 (comment), non-linear grids, like RASM tutorial dataset, all the other file formats you listed (grib, h5), and perhaps some prints of the table metadata.

[ghostiee-11]

Yeah sure, I'll put together a demo notebook covering these cases:

1. Large/out-of-memory datasets - test with chunked loading to verify DataFusion handles lazy evaluation without loading everything into memory
2. Non-linear grids - RASM tutorial dataset (xr.tutorial.open_dataset("rasm")) which has 2D lat/lon coordinates
3. Other file formats - HDF5 via h5netcdf, GRIB via cfgrib (ERA5 sample), and Zarr
4. Table metadata output - print _get_table_metadata() results for each dataset to show what the AI discovery layer sees
5. The xarray-sql#93 edge case - test how the source handles datasets that trigger that issue

I'll add the results as a comment here once ready.

[ghostiee-11]

Hey!! You can check the notebook here, I have compiled all the 5 cases mentioned here :

https://colab.research.google.com/drive/1jvvJS7J4bYMK09SUpCdIyVtKjK9pWekF

[ahuang11]

I think this PR is really close; left a couple more comments; will do some testing with it and see if it's missing anything.

[ghostiee-11]

I think this PR is really close; left a couple more comments; will do some testing with it and see if it's missing anything.

Ofcc, surely will love to refactor if anything breaks...

[ahuang11]

Just a minor comment about the h5 extensions; otherwise, I believe this will go in v1.2.0

[ghostiee-11]

Just a minor comment about the h5 extensions; otherwise, I believe this will go in v1.2.0

Fixed this!! Thanks for the constant review and suggestions!!

[ahuang11]

Thanks!

[ghostiee-11]

Thanks a lot for the merge!!

[github-actions]

This pull request has been automatically locked since there has not been any recent activity after it was closed. Please open a new issue for related bugs.