Lumen + Xarray Integration: XArraySQLSource, Transforms, and AI Support

[ghostiee-11]

Overview

I'm working on integrating xarray support into Lumen, building on the proof-of-concept in #1434 and the request in #1508. This issue outlines my plan to deliver this in 3 incremental PRs, each independently mergeable and useful.

Motivation

Lumen currently handles tabular data (CSV, Parquet, SQL) via DuckDB. Researchers working with N-dimensional scientific data (climate, satellite, microscopy) stored in NetCDF/Zarr/HDF5 formats cannot use Lumen AI to explore their data. The xarray-sql library (Apache DataFusion backend) makes it possible to expose xarray datasets as SQL-queryable tables, bridging this gap.

Plan: 3 PRs

PR 1: Add XArraySQLSource (core source class)

Status: Ready for review

• New lumen/sources/xarray_sql.py - XArraySQLSource(BaseSQLSource)
• Uses xarray-sql (Apache DataFusion) for SQL queries over xarray datasets
• Each data variable becomes a separate SQL table with coordinate columns
• DataFusion uses postgres-compatible dialect for sqlglot translation
• Supports NetCDF, Zarr, HDF5, GRIB with auto-engine detection
• Rich metadata from xarray attributes (units, dimensions, coordinates)
• Dimension info API for UI controls and AI context
• Async execute/get for non-blocking AI agents
• Shared DataFusion context via create_sql_expr_source
• Full to_spec/from_spec serialization
• Optional dependency: pip install lumen[xarray]
• 50 tests passing

PR 2: Add xarray scientific transforms

• DimensionSlice - slice along any dimension by range or values
• SpatialBBox - filter to lat/lon bounding box
• DimensionAggregate - reduce dimensions via groupby (mean, sum, etc.)
• TimeResample - resample time series at different frequencies
• Anomaly - compute deviations from climatological mean
• RollingWindow - moving average/sum/std smoothing

These work on any DataFrame with coordinate columns, useful even without XArraySQLSource.

PR 3: Wire xarray into Lumen AI

• Add .nc/.zarr/.h5 to TABLE_EXTENSIONS in ai/controls/base.py
• Add xarray file reader in _add_table()
• Resolve .nc/.zarr paths in _resolve_data() for lumen-ai serve data.nc CLI
• Upload flow: user drags NetCDF into Lumen AI chat, data becomes queryable

What to expect after all 3 PRs

A scientist can:

1. Run lumen-ai serve climate_data.nc from the command line
2. Or drag-and-drop a .nc file into Lumen AI chat
3. Ask natural language questions like "What's the average temperature by latitude?"
4. The SQL agent generates valid DataFusion SQL against the xarray data
5. Scientific transforms enable dimension slicing, spatial filtering, time resampling, and anomaly detection

Further work

After these 3 PRs establish the foundation, I plan to move forward with:

• Remote Zarr support via fsspec (S3, GCS, Azure)
• xarray DataTree support (hierarchical multi-resolution data)
• anndata compatibility (as @droumis suggested in Add xarray-sql support #1434)
• Multi-file datasets via xr.open_mfdataset()
• Coordinate-aware AI prompt engineering (custom Jinja2 templates for DataFusion)
• Performance testing at scale
• Comprehensive documentation and example notebooks

Related

• Add xarray-sql support #1508 - Add xarray-sql support
• Add xarray-sql support #1434 - Proof of concept by @ahuang11

[ghostiee-11]

We can close this, as what's proposed in this is done and merged

[github-actions]

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