Add xarray-sql support

lumen/ai/controls.py

@@ -1,7 +1,9 @@
 import asyncio
 import io
+import os
 import pathlib
 import re
+import tempfile
 import zipfile
 
 from urllib.parse import urlparse
@@ -22,7 +24,13 @@
 from ..util import detect_file_encoding
 from .memory import _Memory, memory
 
-TABLE_EXTENSIONS = ("csv", "parquet", "parq", "json", "xlsx", "geojson", "wkt", "zip")
+try:
+    from ..sources.xarray_sql import XArraySource
+    XARRAY_AVAILABLE = True
+except ImportError:
+    XARRAY_AVAILABLE = False
+
+TABLE_EXTENSIONS = ("csv", "parquet", "parq", "json", "xlsx", "geojson", "wkt", "zip", "nc", "zarr")
 
 # Download configuration constants
 class DownloadConfig:
@@ -594,6 +602,51 @@ def _add_table(
             conn.execute(init[0])
             cols = ', '.join(f'"{c}"' for c in df.columns if c != 'geometry')
             conversion = f'CREATE TEMP TABLE {table} AS SELECT {cols}, ST_GeomFromWKB(geometry) as geometry FROM {table}_temp'
+        elif extension.endswith(('nc', 'zarr')):
+            # Handle NetCDF and Zarr files with XArraySource
+            if not XARRAY_AVAILABLE:
+                self._error_placeholder.object += f"\nCannot process {table_controls.filename}.{extension}: xarray-sql is not installed."
+                self._error_placeholder.visible = True
+                return 0
+
+            # Create temporary file with appropriate extension
+            with tempfile.NamedTemporaryFile(delete=False, suffix=f'.{extension}') as tmp:
+                file.seek(0)
+                tmp.write(file.read())
+                tmp_path = tmp.name
+
+            try:
+                # Create XArraySource for this file
+                xarray_source = XArraySource(
+                    tables={table: tmp_path},
+                    chunks={'time': 24},  # Default chunking
+                    name=f'XArraySource_{table}'
+                )
+
+                # Add to memory sources
+                if 'sources' in self._memory:
+                    # Check if source already exists
+                    existing_sources = self._memory['sources']
+                    # Remove any existing source with the same table name
+                    self._memory['sources'] = [
+                        s for s in existing_sources
+                        if not (isinstance(s, XArraySource) and table in s.get_tables())
+                    ]
+                    self._memory['sources'].append(xarray_source)
+                else:
+                    self._memory['sources'] = [xarray_source]
+
+                self._memory['source'] = xarray_source
+                self._memory['table'] = table
+                self._last_table = table
+                return 1
+            except Exception as e:
+                self._error_placeholder.object += f"\nCould not load {table_controls.filename}.{extension}: {e}"
+                self._error_placeholder.visible = True
+                # Clean up temp file on error
+                if os.path.exists(tmp_path):
+                    os.unlink(tmp_path)
+                return 0
         else:
             self._error_placeholder.object += f"\nCould not convert {table_controls.filename}.{extension}."
             self._error_placeholder.visible = True

lumen/ai/ui.py

@@ -31,6 +31,12 @@
 from ..pipeline import Pipeline
 from ..sources import Source
 from ..sources.duckdb import DuckDBSource
+
+try:
+    from ..sources.xarray_sql import XArraySource
+    XARRAY_AVAILABLE = True
+except ImportError:
+    XARRAY_AVAILABLE = False
 from ..transforms.sql import SQLLimit
 from ..util import log
 from .agents import (
@@ -502,7 +508,7 @@ def _resolve_data(self, data: DataT | list[DataT] | dict[DataT] | None):
         elif not isinstance(data, list):
             data = [data]
         sources = []
-        mirrors, tables = {}, {}
+        mirrors, tables, xarray_tables = {}, {}, {}
         remote = False
         for src in data:
             if isinstance(src, Source):
@@ -518,12 +524,18 @@ def _resolve_data(self, data: DataT | list[DataT] | dict[DataT] | None):
                 if src.startswith('http'):
                     remote = True
                 if src.endswith(('.parq', '.parquet', '.csv', '.json', '.tsv', '.jsonl', '.ndjson')):
-                    table = src
+                    tables[name] = src
+                elif src.endswith(('.nc', '.zarr')):
+                    if not XARRAY_AVAILABLE:
+                        raise ValueError(
+                            f"Cannot load {src}: xarray-sql is not installed. "
+                            "Install it with: pip install xarray-sql"
+                        )
+                    xarray_tables[name.split(".")[0]] = src
                 else:
                     raise ValueError(
                         f"Could not determine how to load {src} file."
                     )
-                tables[name] = table
         if tables or mirrors:
             initializers = ["INSTALL httpfs;", "LOAD httpfs;"] if remote else []
             source = DuckDBSource(
@@ -534,6 +546,16 @@ def _resolve_data(self, data: DataT | list[DataT] | dict[DataT] | None):
                 uri=':memory:'
             )
             sources.append(source)
+
+        if xarray_tables:
+            # Create XArraySource for NetCDF and Zarr files
+            xarray_source = XArraySource(
+                tables=xarray_tables,
+                chunks={'time': 24},  # Default chunking, can be customized
+                name=f'{PROVIDED_SOURCE_NAME}_xarray'
+            )
+            sources.append(xarray_source)
+
         memory["sources"] = sources
 
     def show(self, **kwargs):

lumen/transforms/sql.py

@@ -723,7 +723,7 @@ def apply(self, sql_in: str) -> str:
         expression = self.parse_sql(sql_in)
         dialect = self.write or self.read
 
-        if dialect in ('sqlserver', 'duckdb') or (dialect in ('postgres', 'redshift', 'snowflake', 'bigquery')):
+        if dialect in ('sqlserver', 'duckdb') or (dialect in ('redshift', 'snowflake', 'bigquery')):
             return self._apply_tablesample_dialect(expression)
         elif dialect in ('mysql', 'mariadb'):
             return self._apply_mysql_dialect(expression)

pixi.toml

@@ -32,6 +32,7 @@ panel = ">=1.7.5"
 param = ">=2.2.1"
 pip = "*"
 sqlglot = "*"
+sqlglotrs = "*"
 
 [feature.py310.dependencies]
 python = "3.10.*"

pyproject.toml

@@ -38,6 +38,7 @@ dependencies = [
     "panel >=1.7.5",
     "param >=2.2.1",
     "sqlglot",
+    "sqlglotrs",
 ]
 
 [project.urls]