Wire XArraySQLSource into Lumen AI for natural language queries on scientific data by ghostiee-11 · Pull Request #1791 · holoviz/lumen

ghostiee-11 · 2026-03-24T09:55:00Z

I've been building a climate data analysis tool with xarray datasets and wanted to use Lumen AI to query them with natural language. Since XArraySQLSource landed in #1741, this wires it into the AI layer so lumen-ai serve data.nc just works, including drag-and-drop upload of .nc and .h5 files.

Depends on: #1741 (XArraySQLSource)

E2E validated with 9 queries + 2 plots on an 11MB climate dataset (screenshots in comments).

Adds a new optional source that enables SQL queries over xarray datasets via xarray-sql/DataFusion. Each data variable is exposed as a SQL table, inheriting from BaseSQLSource to avoid duplication. Includes documentation, 58 tests, and pyproject.toml extras for xarray dependencies.

… metadata - _resolve_chunks: use dataset.chunks directly instead of getattr - get_schema: expand comment explaining TABLESAMPLE/DataFusion chain - normalize_table: add step-by-step comments for clarity - get(): document __limit__ filter usage - _get_table_metadata: replace math.prod with var.size, add auxiliary coord descriptions for non-linear grids (e.g. RASM 2D lat/lon) - Add cfgrib/eccodes deps for GRIB file support - Add tests: engine detection, GRIB I/O, 2D auxiliary coord metadata

- Remove dialect='postgres', inherit 'any' from BaseSQLSource - Add from_dataset() classmethod mirroring DuckDBSource.from_df() - Trim xarray deps: remove cfgrib/eccodes - Simplify normalize_table, _resolve_chunks, get_schema comment - Update dialect test, add from_dataset test (58 pass)

…rib group - Extract _format_aux_coords() helper to reduce indentation in _get_table_metadata - Change .h5/.hdf5 engine map from h5netcdf to netcdf4 (fewer deps, same functionality) - Add xarray-grib optional dep group with cfgrib + eccodes - Update engine detection tests to match

Add xarray file support to Lumen AI so users can query NetCDF, Zarr, and HDF5 datasets via natural language. CLI: detect xarray file extensions in _resolve_data, create XArraySQLSource Upload: register upload handlers for .nc/.h5 drag-and-drop with temp file cleanup via atexit Controls: add xarray extensions to TABLE_EXTENSIONS SQL agent: fix crash when source.tables is a list instead of dict Source: set dialect=None for sqlglot compat, auto-assign integer coords for bare dimensions (fixes RASM-type datasets) Deps: remove h5netcdf, add xarray-grib optional group 8 new tests, 123 total passing. E2E validated with 9 queries + 2 plots on 11MB climate dataset (NetCDF) and Zarr ocean data.

- Remove .h5/.hdf5/.he5 from XARRAY_ENGINES (xarray doesn't handle pure HDF5 well, h5py is better for those) - Drop h5netcdf from xarray deps in pyproject.toml and pixi.toml - Add separate xarray-grib extra with cfgrib and eccodes - Update engine detection test

ghostiee-11 · 2026-03-24T10:09:56Z

E2E Validation Screenshots

Tested on an 11MB synthetic climate dataset (365 days x 37 lats x 72 lons, 3 variables: temperature, precipitation, wind_speed).

Q1: Data points per variable (UNION ALL + COUNT)

Q2: Wind stats southern hemisphere (MIN/MAX/STDDEV + WHERE lat < 0)

Q3: Tropical precipitation by longitude (SUM + BETWEEN + GROUP BY)

Q4: Quarterly average temperature (GROUP BY time + ORDER BY)

Q5: 5 coldest locations (MIN + GROUP BY lat,lon + ORDER ASC)

Q6: Average temperature by latitude (AVG + GROUP BY + ORDER BY)

Q7: Monthly precipitation + temperature JOIN (JOIN + date_trunc + WHERE BETWEEN)

Q8: Line chart - temperature by latitude (SQL + Vega-Lite)

Q9: Bar chart - seasonal precipitation (CASE/WHEN + Vega-Lite)

codecov · 2026-03-24T10:58:04Z

Codecov Report

❌ Patch coverage is 69.69697% with 40 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.02%. Comparing base (45f06f7) to head (dfa4689).
⚠️ Report is 4 commits behind head on main.

Files with missing lines	Patch %	Lines
lumen/tests/ai/test_ui.py	77.77%	18 Missing ⚠️
lumen/ai/ui.py	57.14%	15 Missing ⚠️
lumen/command/ai.py	0.00%	5 Missing ⚠️
lumen/util.py	71.42%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1791      +/-   ##
==========================================
+ Coverage   69.63%   70.02%   +0.38%     
==========================================
  Files         190      193       +3     
  Lines       31495    32445     +950     
==========================================
+ Hits        21933    22721     +788     
- Misses       9562     9724     +162

☔ 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 · 2026-03-24T17:35:39Z

+        import numpy as np
+        import xarray as xr


Ensure imports are up top (in the future, please do a quick look through your PRs before marking ready for review; it can get a bit tiring to be repeating the same comment).

Yeah!! i am sorry, i understand will keep this in mind.. Thankyou

Still here~

Remove top-level import xarray/xarray_sql so installing xarray does not cause it to load by default. Use _check_xarray_available() helper and inline imports in methods that need them. Remove XArraySQLSource from sources/__init__.py -- other optional sources (Snowflake, BigQuery) are not imported there either.

ahuang11 · 2026-03-26T18:36:16Z

Merge conflicts here now.

ghostiee-11 · 2026-03-26T18:40:22Z

Merge conflicts here now.

Yup due to #1741 merge will fix!!

Resolve conflicts from holoviz#1741 merge by adopting main's canonical XArraySQLSource implementation (class-level _xarray_engines, classmethod _detect_engine, dialect="any").

ahuang11 · 2026-03-26T22:43:53Z

+            try:
+                from ..sources.xarray_sql import XArraySQLSource
+            except ImportError:
+                return None


We should make this a util, rather than try/except Import everywhere.

if not _check_xarray_available(): raise ImportError( "xarray and xarray-sql are required for XArraySQLSource. " "Install them with: pip install lumen[xarray]" )

ahuang11 · 2026-03-26T22:44:27Z

+        """Test that .nc creates separate source while .csv goes to DuckDB."""
+        pytest.importorskip("xarray_sql")
+        import numpy as np
+        import xarray as xr


Please take some time to look at diffs on GitHub and check for imports inlined :'(

Yess yess for sure.. I will keep this in mind.. this pr is outdated the day you gave me feedback about the imports on top thing!!

ahuang11 · 2026-03-26T22:45:58Z

+                pytest.skip("NetCDF write backend unavailable")
+            result = UI._resolve_data(str(nc_path))
+            assert len(result) == 1
+            from lumen.sources.xarray_sql import XArraySQLSource


This import is even more insane, not even at the top of the function.

lemme fix all

ahuang11 · 2026-03-26T22:48:04Z

+XARRAY_FILE_EXTENSIONS = ('.nc', '.nc4', '.netcdf', '.h5', '.hdf5', '.he5', '.zarr', '.grib', '.grib2', '.grb', '.grb2')
+XARRAY_UPLOAD_EXTENSIONS = ('nc', 'nc4', 'h5', 'hdf5')
+


Why are constants located between imports?

ahuang11

Take some time to look at GitHub diff! I'm finding myself repeating the same comments and requests.

ghostiee-11 · 2026-03-26T22:52:29Z

Take some time to look at GitHub diff! I'm finding myself repeating the same comments and requests.

Yes, yes, I understand. I will keep this in mind.

…, gate upload handlers

ahuang11 · 2026-03-27T18:04:42Z

+def _check_xarray_available():
+    """Check if xarray and xarray-sql are installed."""
+    try:
+        import xarray  # noqa
+        import xarray_sql  # noqa
+        return True
+    except ImportError:
+        return False


I suspect more suited to be in util. XArraySQLSource can also use it I think...

ahuang11 · 2026-03-27T18:06:26Z

+        if not _check_xarray_available():
+            return {}
+
+        _temp_files: list[str] = []
+
+        def _cleanup_temp_files():
+            for path in _temp_files:
+                try:
+                    os.unlink(path)
+                except OSError:
+                    pass
+
+        atexit.register(_cleanup_temp_files)
+
+        def _handle_xarray_upload(context, file_obj, alias, filename):
+            from ..sources.xarray_sql import XArraySQLSource
+            ext = filename.rsplit('.', 1)[-1] if '.' in filename else 'nc'
+            tmp = tempfile.NamedTemporaryFile(delete=False, suffix=f'.{ext}')
+            file_obj.seek(0)
+            tmp.write(file_obj.read())
+            tmp.flush()
+            tmp.close()
+            _temp_files.append(tmp.name)
+            return XArraySQLSource(uri=tmp.name, name=alias)


Not a big fan of nested functions not at the top, e.g.

def func1() # no other code here please def nested() .... other_code

ahuang11 · 2026-03-27T18:09:01Z

Wanted to bring this up again in your description
https://holoviz.org/contribute.html#using-ai-readme

Do NOT list out what changed, unless you want to bring attention to specific changes

- Move check_xarray_available to lumen/util.py (shared by ui.py and xarray_sql.py) - Remove duplicate ImportError check in _resolve_data (XArraySQLSource handles it) - Move constants after imports, reorder nested functions in upload handler - Move inline import io to top of test file, add skipif decorators

ghostiee-11 · 2026-04-04T19:39:29Z

Hey Andrew, addressed all your feedback:

Moved _check_xarray_available to lumen/util.py as a shared utility. Both xarray_sql.py and ui.py now import from there, no more duplication.
Removed the duplicate ImportError check in _resolve_data. XArraySQLSource.__init__ already handles that.
Moved constants after imports (no longer wedged between import blocks).
Reordered nested functions in _get_xarray_upload_handlers so definitions come first, side-effects after.
Moved import io to the top of the test file and added skipif decorators for xarray-dependent tests.

Sorry again for the repeated import issues. I'm being more careful reviewing the GitHub diff before pushing now.

- Resolve controls/base.py -> controls/ingest/ refactor - Define XARRAY_UPLOAD_EXTENSIONS in controls/ingest/constants.py and splat into TABLE_EXTENSIONS (per review) - Derive XARRAY_FILE_EXTENSIONS from XARRAY_UPLOAD_EXTENSIONS so the xarray extension list has a single source of truth - Move XArraySQLSource import to top of lumen/ai/ui.py; drop the inline imports in _get_xarray_upload_handlers and _resolve_data - Reorder _get_xarray_upload_handlers so nested defs come before body - Move MagicMock import to top of tests/ai/test_ui.py

ahuang11 · 2026-04-17T16:17:16Z

+XARRAY_FILE_EXTENSIONS = tuple(f'.{ext}' for ext in XARRAY_UPLOAD_EXTENSIONS) + (
+    '.netcdf', '.he5', '.zarr', '.grib', '.grib2', '.grb', '.grb2',


I'm still confused as to why we have upload extensions vs file extensions. Can you explain?

You're right, there's no real reason for them to be separate. XArraySQLSource accepts all these formats via a path, and the upload handler just writes to a temp file before calling it, so there's no format supported by one path and not the other.

will fix this!!

The previous split (XARRAY_UPLOAD_EXTENSIONS=4, XARRAY_FILE_EXTENSIONS=11) had no functional reason — XArraySQLSource accepts all 11 formats via URI, and the upload handler writes to a temp file before calling it, so there is no format supported by one path and not the other. - controls/ingest/constants.py: single XARRAY_EXTENSIONS tuple (11 formats, bare names), splatted into TABLE_EXTENSIONS - ui.py: import XARRAY_EXTENSIONS, derive XARRAY_DOTTED_EXTENSIONS for the endswith check in _resolve_data

… directly

ahuang11 · 2026-04-17T23:07:56Z

Thanks! This is definitely useful for opening use cases with xarray. In the future, we'll need to handle plotting these gridded datasets better. Really appreciate the work thus far! #1821

ghostiee-11 · 2026-04-18T04:40:51Z

Thanks a lot, Andrew!! Will look into this, and escalate it properly for every type of use cases

ghostiee-11 · 2026-04-18T04:41:22Z

Thanks for the quick review and merge!!

After holoviz#1791 landed, loading a NetCDF file and asking the AI to plot it produces a scatter or line chart rather than a 2D heatmap, because: - hvPlotBaseView had no way to pass a value column (C) to hvplot - The AI agents had no awareness the pipeline was backed by gridded xarray data Changes: - Add `C` param (value column) to hvPlotBaseView; wire into df.hvplot() - Add `quadmesh` to the kind selector - New lumen/ai/gridded.py: detect_gridded() returns coordinate/variable metadata when pipeline.source is XArraySQLSource, else None - BaseViewAgent._generate_yaml_spec: call detect_gridded() once, pass result as `gridded` template variable - BaseViewAgent/main.jinja2 + hvPlotAgent/main.jinja2: conditional block guides LLM to use kind='heatmap' with x/y/C for gridded data - hvPlotAgent._extract_spec: drop by/groupby when kind=heatmap Addresses holoviz#1821 (1/3)

ghostiee-11 added 9 commits March 16, 2026 23:36

Merge origin/main, resolve pyproject.toml conflict

5a91c7c

Fix isort: add blank line between stdlib imports in test file

f2a3162

Merge origin/main, resolve test_ui.py conflict

fd44709

Skip xarray tests when netCDF4 write backend unavailable in CI

fcb6316

ghostiee-11 force-pushed the feat/xarray-ai-integration branch from e5f921a to fcb6316 Compare March 24, 2026 10:37

ghostiee-11 marked this pull request as ready for review March 24, 2026 11:00

ahuang11 reviewed Mar 24, 2026

View reviewed changes

Comment thread lumen/ai/ui.py Outdated

Move imports to top of file per review

612e589

ahuang11 reviewed Mar 24, 2026

View reviewed changes

ghostiee-11 added 2 commits March 24, 2026 23:52

Merge PR1 lazy imports, inline numpy/xarray per review

6788a9d

ahuang11 marked this pull request as draft March 26, 2026 18:36

Merge main into feat/xarray-ai-integration

56fc483

Resolve conflicts from holoviz#1741 merge by adopting main's canonical XArraySQLSource implementation (class-level _xarray_engines, classmethod _detect_engine, dialect="any").

ghostiee-11 marked this pull request as ready for review March 26, 2026 20:29

ahuang11 reviewed Mar 26, 2026

View reviewed changes

Comment thread lumen/ai/ui.py

ahuang11 reviewed Mar 26, 2026

View reviewed changes

ahuang11 requested changes Mar 26, 2026

View reviewed changes

ghostiee-11 added 2 commits March 27, 2026 04:38

Address review: move imports to top, extract xarray availability util…

4136d12

…, gate upload handlers

Fix isort: add blank line between import groups

c07277e

ghostiee-11 marked this pull request as ready for review March 26, 2026 23:29

ahuang11 reviewed Mar 27, 2026

View reviewed changes

ahuang11 requested changes Mar 27, 2026

View reviewed changes

ahuang11 reviewed Mar 27, 2026

View reviewed changes

Comment thread lumen/ai/ui.py Outdated

ahuang11 marked this pull request as draft March 27, 2026 18:10

ghostiee-11 marked this pull request as ready for review April 4, 2026 19:53

ahuang11 requested changes Apr 6, 2026

View reviewed changes

Comment thread lumen/ai/ui.py Outdated

Comment thread lumen/ai/ui.py Outdated

ahuang11 reviewed Apr 7, 2026

View reviewed changes

Comment thread lumen/ai/controls/base.py Outdated

ahuang11 reviewed Apr 17, 2026

View reviewed changes

ghostiee-11 marked this pull request as draft April 17, 2026 17:24

Update test_upload_handler_returns_dict for unified extensions

3c85c61

ghostiee-11 marked this pull request as ready for review April 17, 2026 17:42

ahuang11 reviewed Apr 17, 2026

View reviewed changes

Comment thread lumen/ai/ui.py Outdated

Drop XARRAY_DOTTED_EXTENSIONS; check suffix against XARRAY_EXTENSIONS…

dfa4689

… directly

ahuang11 enabled auto-merge (squash) April 17, 2026 23:08

ahuang11 approved these changes Apr 17, 2026

View reviewed changes

ahuang11 merged commit 674e698 into holoviz:main Apr 17, 2026
12 checks passed

		XARRAY_FILE_EXTENSIONS = ('.nc', '.nc4', '.netcdf', '.h5', '.hdf5', '.he5', '.zarr', '.grib', '.grib2', '.grb', '.grb2')
		XARRAY_UPLOAD_EXTENSIONS = ('nc', 'nc4', 'h5', 'hdf5')

		XARRAY_FILE_EXTENSIONS = tuple(f'.{ext}' for ext in XARRAY_UPLOAD_EXTENSIONS) + (
		'.netcdf', '.he5', '.zarr', '.grib', '.grib2', '.grb', '.grb2',

Uh oh!

Uh oh!

Conversation

ghostiee-11 commented Mar 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ghostiee-11 commented Mar 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

E2E Validation Screenshots

Uh oh!

codecov Bot commented Mar 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

ahuang11 Mar 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ahuang11 commented Mar 26, 2026

Uh oh!

ghostiee-11 commented Mar 26, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ahuang11 left a comment

Choose a reason for hiding this comment

Uh oh!

ghostiee-11 commented Mar 26, 2026

Uh oh!

ahuang11 Mar 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ahuang11 Mar 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ahuang11 commented Mar 27, 2026

Uh oh!

ghostiee-11 commented Apr 4, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ahuang11 commented Apr 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

ghostiee-11 commented Apr 18, 2026

Uh oh!

ghostiee-11 commented Apr 18, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

ghostiee-11 commented Mar 24, 2026 •

edited

Loading

ghostiee-11 commented Mar 24, 2026 •

edited

Loading

codecov Bot commented Mar 24, 2026 •

edited

Loading

ahuang11 Mar 24, 2026 •

edited

Loading

ahuang11 Mar 27, 2026 •

edited

Loading

ahuang11 Mar 27, 2026 •

edited

Loading

ahuang11 commented Apr 17, 2026 •

edited

Loading