diff --git a/lumen/ai/controls.py b/lumen/ai/controls.py index cbd66fb30..164965c2d 100644 --- a/lumen/ai/controls.py +++ b/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 diff --git a/lumen/ai/ui.py b/lumen/ai/ui.py index ccfddaf8d..4ec8c932f 100644 --- a/lumen/ai/ui.py +++ b/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): diff --git a/lumen/transforms/sql.py b/lumen/transforms/sql.py index 0e7b7359a..8c94c0834 100644 --- a/lumen/transforms/sql.py +++ b/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) diff --git a/pixi.toml b/pixi.toml index 1965eca22..788e3adef 100644 --- a/pixi.toml +++ b/pixi.toml @@ -32,6 +32,7 @@ panel = ">=1.7.5" param = ">=2.2.1" pip = "*" sqlglot = "*" +sqlglotrs = "*" [feature.py310.dependencies] python = "3.10.*" diff --git a/pyproject.toml b/pyproject.toml index bd721a316..89b17a520 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,7 @@ dependencies = [ "panel >=1.7.5", "param >=2.2.1", "sqlglot", + "sqlglotrs", ] [project.urls]