Wire XArraySQLSource into Lumen AI for natural language queries on scientific data

docs/configuration/sources.md

@@ -29,6 +29,7 @@ ui.servable()
 |--------|---------|
 | Files | CSV, Parquet, JSON (local or URL) |
 | DuckDB | Local SQL queries on files |
+| xarray | N-dimensional scientific data (NetCDF, Zarr, HDF5) |
 | Snowflake | Cloud data warehouse |
 | BigQuery | Google's data warehouse |
 | PostgreSQL | PostgreSQL via SQLAlchemy |
@@ -173,6 +174,45 @@ source = DuckDBSource(
 
 1. Required for HTTP/S3 access
 
+### xarray for scientific data
+
+Query N-dimensional scientific datasets (NetCDF, Zarr, HDF5) with SQL via Apache DataFusion:
+
+``` bash title="Install dependencies"
+pip install lumen[xarray]
+```
+
+``` py title="Query a NetCDF file"
+from lumen.sources.xarray_sql import XArraySQLSource
+import lumen.ai as lmai
+
+source = XArraySQLSource(uri='air_temperature.nc')
+
+ui = lmai.ExplorerUI(data=source)
+ui.servable()
+```
+
+Each data variable in the dataset becomes a separate SQL table with coordinate columns:
+
+``` py title="Direct SQL queries"
+source = XArraySQLSource(uri='air_temperature.nc')
+source.get_tables()  # ['air']
+
+df = source.execute(
+    "SELECT lat, lon, AVG(air) as avg_temp "
+    "FROM air GROUP BY lat, lon"
+)
+```
+
+**Select specific variables:**
+
+``` py hl_lines="3"
+source = XArraySQLSource(
+    uri='climate_data.nc',
+    variables=['temperature', 'pressure']
+)
+```
+
 ### Multiple sources
 
 ``` py title="Mix sources"

lumen/ai/controls/base.py

@@ -22,7 +22,7 @@
 from ..utils import log_debug
 from .progress import Progress
 
-TABLE_EXTENSIONS = ("csv", "parquet", "parq", "json", "xlsx", "geojson", "wkt", "zip")
+TABLE_EXTENSIONS = ("csv", "parquet", "parq", "json", "xlsx", "geojson", "wkt", "zip", "nc", "nc4", "h5", "hdf5")
 
 METADATA_EXTENSIONS = ("md", "txt", "yaml", "yml", "json", "pdf", "docx", "doc", "pptx", "ppt")
 METADATA_FILENAME_PATTERNS = ("_metadata", "metadata_", "readme", "schema")

lumen/ai/ui.py

@@ -28,6 +28,9 @@
 )
 from panel_splitjs import HSplit, MultiSplit, VSplit
 
+XARRAY_FILE_EXTENSIONS = ('.nc', '.nc4', '.netcdf', '.h5', '.hdf5', '.he5', '.zarr', '.grib', '.grib2', '.grb', '.grb2')
+XARRAY_UPLOAD_EXTENSIONS = ('nc', 'nc4', 'h5', 'hdf5')
+
 from lumen.ai.agents.deck_gl import DeckGLAgent
 
 from ..pipeline import Pipeline
@@ -502,6 +505,40 @@ def __init__(
         self._configure_session()
         self._render_page()
 
+    @staticmethod
+    def _get_xarray_upload_handlers():
+        """Build upload handlers for xarray file formats."""
+        import atexit
+        import os
+
+        _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):
+            try:
+                from ..sources.xarray_sql import XArraySQLSource
+            except ImportError:
+                return None
+            import tempfile
+            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)
+
+        return {ext: _handle_xarray_upload for ext in XARRAY_UPLOAD_EXTENSIONS}
+
     @classmethod
     def _resolve_data(
         cls, data: DataT | list[DataT] | dict[DataT] | None
@@ -591,6 +628,19 @@ def _resolve_data(
                         sources.append(source)
                     continue
 
+                # Handle xarray files (NetCDF, Zarr, HDF5, GRIB)
+                if src.endswith(XARRAY_FILE_EXTENSIONS):
+                    try:
+                        from ..sources.xarray_sql import XArraySQLSource
+                    except ImportError as e:
+                        raise ImportError(
+                            "xarray and xarray-sql are required to load xarray files. "
+                            "Install them with: pip install lumen[xarray]"
+                        ) from e
+                    source = XArraySQLSource(uri=str(Path(src).absolute()))
+                    sources.append(source)
+                    continue
+
                 if src.startswith('http'):
                     remote = True
                 if src.endswith(('.parq', '.parquet', '.csv', '.json', '.tsv', '.jsonl', '.ndjson')):
@@ -1216,14 +1266,18 @@ async def _do_sync():
             param.parameterized.async_executor(_do_sync)
         self._source_catalog.param.watch(_schedule_visibility_change, 'visibility_changed')
 
+        # Register xarray upload handlers (if xarray-sql is available)
+        xarray_handlers = self._get_xarray_upload_handlers()
+        merged_upload_handlers = {**xarray_handlers, **self.upload_handlers}
+
         # Initialize source controls
         self._source_controls = []
         control_tabs = []
         for control in self.source_controls:
             control_kwargs = {
                 'context': self.context,
                 'source_catalog': self._source_catalog,
-                'upload_handlers': self.upload_handlers
+                'upload_handlers': merged_upload_handlers
             }
             if control is UploadControls and self.filedropper_kwargs:
                 control_kwargs['filedropper_kwargs'] = self.filedropper_kwargs

lumen/command/ai.py

@@ -241,10 +241,15 @@ def __init__(
                 continue
 
             path = Path(table_path)
-            if not path.is_file():
-                if path.exists():
-                    raise ValueError(f"Table path is not a file: {table_path}")
+            if not path.exists():
                 raise FileNotFoundError(f"Table file not found: {table_path}")
+            if path.is_dir() and not table_path.endswith('.zarr'):
+                raise ValueError(
+                    f"Table path is a directory: {table_path}. "
+                    "Only .zarr directories are supported."
+                )
+            if not path.is_file() and not path.is_dir():
+                raise ValueError(f"Table path is not a file: {table_path}")
 
         source = self._build_source_code(
             tables=tables,

lumen/sources/__init__.py

@@ -1 +1,6 @@
 from .base import *
+
+try:
+    from .xarray_sql import XArraySQLSource  # noqa: F401
+except ImportError:
+    pass