feat: add Arraylake source control

lumen/ai/controls/__init__.py

@@ -4,11 +4,11 @@
 from .explorer import TableExplorer
 from .ingest import (
     METADATA_EXTENSIONS, METADATA_FILENAME_PATTERNS, TABLE_EXTENSIONS,
-    BaseSourceControls, CatalogSourceControls, CodeSourceControls,
-    DownloadConfig, DownloadSourceControls, FileSourceControls,
-    OpenAPISourceControls, ParametricSourceControls, RESTAPISourceControls,
-    SourceResult, UploadedFileRow, UploadSourceControls, URLSourceControls,
-    download_file,
+    ArraylakeSourceControls, BaseSourceControls, CatalogSourceControls,
+    CodeSourceControls, DownloadConfig, DownloadSourceControls,
+    FileSourceControls, OpenAPISourceControls, ParametricSourceControls,
+    RESTAPISourceControls, SourceResult, UploadedFileRow, UploadSourceControls,
+    URLSourceControls, download_file,
 )
 from .revision import AnnotationControls, RetryControls, RevisionControls
 
@@ -17,6 +17,7 @@
     "METADATA_FILENAME_PATTERNS",
     "TABLE_EXTENSIONS",
     "AnnotationControls",
+    "ArraylakeSourceControls",
     "BaseSourceControls",
     "CatalogSourceControls",
     "CodeSourceControls",

lumen/ai/controls/ingest/__init__.py

@@ -1,3 +1,7 @@
+try:
+    from .arraylake import ArraylakeSourceControls
+except ImportError:
+    ArraylakeSourceControls = None
 from .base import BaseSourceControls
 from .catalog import CatalogSourceControls
 from .code import CodeSourceControls
@@ -20,6 +24,7 @@
     "METADATA_EXTENSIONS",
     "METADATA_FILENAME_PATTERNS",
     "TABLE_EXTENSIONS",
+    "ArraylakeSourceControls",
     "BaseSourceControls",
     "CatalogSourceControls",
     "CodeSourceControls",

lumen/ai/controls/ingest/arraylake.py

@@ -0,0 +1,132 @@
+from __future__ import annotations
+
+import param
+
+from ....sources.xarray_sql import XArraySQLSource
+from ...translate import params_to_callable
+from .parametric import ParametricSourceControls
+from .result import SourceResult
+from .utils import open_arraylake_dataset
+
+
+class ArraylakeSourceControls(ParametricSourceControls):
+    """
+    Parametric controls that open an Arraylake repository as a
+    SQL-queryable xarray source.
+
+    The user supplies a repo, branch, and optional Zarr group; the control
+    opens the Icechunk-backed store with xarray and registers the result as
+    an ``XArraySQLSource`` (one SQL table per gridded data variable), mirroring
+    how uploaded CSVs become a ``DuckDBSource``.
+
+    Example
+    -------
+    ::
+
+        from lumen.ai.controls.ingest import ArraylakeSourceControls
+        ui = ExplorerUI(source_controls=[ArraylakeSourceControls])
+
+    Requires the optional dependencies (Python >=3.12)::
+
+        pip install lumen[arraylake]
+    """
+
+    repo = param.String(default="", doc="""
+        Arraylake repository, e.g. 'earthmover-public/goes-16'.""")
+
+    branch = param.String(default="main", doc="""
+        Branch or ref to read from.""")
+
+    group = param.String(default="", doc="""
+        Optional Zarr group within the repository.""")
+
+    variables = param.List(default=None, allow_None=True, precedence=-1, doc="""
+        Subset of data variables to expose as tables. When omitted, all
+        gridded (non-scalar) variables are exposed.""")
+
+    label = '<span class="material-icons" style="vertical-align: middle;">cloud_queue</span> Arraylake'
+
+    def as_tools(
+        self, query: str | None = None, top_k: int = 5,
+    ) -> list[tuple[str, callable]]:
+        """Expose the repo/branch/group params as a single typed agent tool.
+
+        Mirrors ``URLSourceControls.as_tools``: subclass-pattern controls
+        declare class-level params rather than registering actions, so we
+        synthesize one callable whose signature mirrors those params for
+        ``FunctionTool`` to build a schema from.
+        """
+        if self._cached_tools is not None:
+            return self._cached_tools
+
+        query_names = self._get_query_param_names()
+        if not query_names:
+            self._cached_tools = []
+            return self._cached_tools
+
+        action_name = self.__class__.__name__.removesuffix("Controls")
+
+        async def _tool_callable(**kwargs) -> SourceResult:
+            return await self.load_action(action_name, **kwargs)
+
+        query_params = {name: self.param[name] for name in query_names}
+        doc = (self.__doc__ or f"Load data from {action_name}.").strip()
+        params_to_callable(
+            _tool_callable, query_params,
+            name=action_name.lower() or "load_data",
+            doc=doc,
+        )
+
+        self._cached_tools = [(action_name, _tool_callable)]
+        return self._cached_tools
+
+    async def _fetch_data(self, action_name: str, **params) -> SourceResult:
+        repo = (params.get("repo") or self.repo).strip()
+        if not repo:
+            return SourceResult.empty(
+                "Provide an Arraylake repo, e.g. 'earthmover-public/goes-16'."
+            )
+        branch = (params.get("branch") or self.branch).strip() or "main"
+        group = (params.get("group") or self.group).strip() or None
+
+        self.progress(f"Opening Arraylake repo {repo!r}")
+        try:
+            ds = open_arraylake_dataset(repo, branch=branch, group=group)
+        except ImportError:
+            return SourceResult.empty(
+                "Arraylake support requires `pip install lumen[arraylake]` "
+                "(Python >=3.12)."
+            )
+        except Exception as e:
+            return SourceResult.empty(f"Could not open Arraylake repo {repo!r}: {e}")
+
+        # Drop 0-dim scalar metadata variables: they have no chunkable
+        # dimension and cannot become a partitioned table.
+        variables = self.variables or [
+            name for name, var in ds.data_vars.items() if var.ndim > 0
+        ]
+        if not variables:
+            return SourceResult.empty(
+                f"No gridded data variables to load from {repo!r}."
+            )
+
+        source_id = f"{self.source_name_prefix}{self._count:06d}"
+        self._count += 1
+        try:
+            # Virtual-zarr stores (e.g. GOES) can chunk variables
+            # inconsistently along a shared dimension; unify so the dataset
+            # registers as a single coherent set of tables.
+            ds = ds.unify_chunks()
+            source = XArraySQLSource.from_dataset(
+                ds, variables=variables, name=source_id
+            )
+        except Exception as e:
+            return SourceResult.empty(
+                f"Could not register {repo!r} as a source: {e}"
+            )
+
+        tables = source.get_tables()
+        first_table = tables[0] if tables else None
+        location = "/".join(part for part in (repo, branch, group) if part)
+        message = f"Loaded {len(tables)} tables from Arraylake {location}"
+        return SourceResult.from_source(source, table=first_table, message=message)

lumen/ai/controls/ingest/utils.py

@@ -202,6 +202,40 @@ def read_json_to_dataframe(content: str | bytes) -> pd.DataFrame:
     raise ValueError(f"Unsupported JSON root type: {type(data).__name__}")
 
 
+def open_arraylake_dataset(
+    repo: str,
+    branch: str = "main",
+    group: str | None = None,
+    snapshot_cache_nodes: int = 10_000,
+):
+    """Open an Arraylake repository as an xarray Dataset.
+
+    Requires the optional ``arraylake`` and ``icechunk`` dependencies
+    (``pip install lumen[arraylake]``, Python >=3.12). A larger Icechunk
+    snapshot cache avoids repeatedly re-reading metadata while opening large
+    stores (on GOES this cut the open from ~200s to ~16s).
+
+    Raises
+    ------
+    ImportError
+        If ``arraylake``/``xarray`` are not installed.
+    """
+    import arraylake as al
+    import xarray as xr
+
+    repo_config = None
+    try:
+        import icechunk
+        repo_config = icechunk.RepositoryConfig(
+            caching=icechunk.CachingConfig(num_snapshot_nodes=snapshot_cache_nodes)
+        )
+    except Exception:
+        pass
+
+    store = al.Client().get_repo(repo, config=repo_config).readonly_session(branch).store
+    return xr.open_zarr(store, group=group)
+
+
 @dataclass
 class FileReadResult:
     """Result of parsing a file into DataFrames."""

lumen/ai/ui.py

@@ -49,8 +49,8 @@
 )
 from .context import TContext
 from .controls import (
-    BaseSourceControls, DownloadSourceControls, FileSourceControls,
-    SourceCatalog, TableExplorer, UploadSourceControls,
+    ArraylakeSourceControls, BaseSourceControls, DownloadSourceControls,
+    FileSourceControls, SourceCatalog, TableExplorer, UploadSourceControls,
 )
 from .controls.ingest.constants import XARRAY_EXTENSIONS
 from .coordinator import Coordinator, Plan, Planner
@@ -451,7 +451,7 @@ class UI(Viewer):
     page_config = param.Dict(default={}, doc="""
         Configuration for the panel-material-ui Page component the UI is rendered into.""")
 
-    source_controls = param.List(default=[UploadSourceControls, DownloadSourceControls], doc="""
+    source_controls = param.List(default=[UploadSourceControls, DownloadSourceControls, ArraylakeSourceControls], doc="""
         List of SourceControls types to manage datasets.""")
 
     filedropper_kwargs = param.Dict(default={}, doc="""Keyword arguments to pass to FileDropper in UploadControls.
@@ -1320,6 +1320,10 @@ async def _do_sync():
         self._source_controls = []
         control_tabs = []
         for control in self.source_controls:
+            if control is None:
+                # Optional controls (e.g. Arraylake) are None when their
+                # dependencies are unavailable; skip them.
+                continue
             if isinstance(control, BaseSourceControls):
                 # Already instantiated — adopt it, just wire up context/catalog
                 control_inst = control

lumen/tests/ai/test_controls/test_arraylake.py

@@ -0,0 +1,116 @@
+import inspect
+import sys
+import types
+
+from unittest.mock import MagicMock
+
+import numpy as np
+import pytest
+
+try:
+    import xarray as xr
+    import xarray_sql  # noqa
+except ModuleNotFoundError:
+    pytest.skip(
+        "xarray / xarray-sql not installed, skipping Arraylake control tests.",
+        allow_module_level=True,
+    )
+
+from lumen.ai.controls import ArraylakeSourceControls
+from lumen.ai.controls.ingest.result import SourceResult
+from lumen.sources.xarray_sql import XArraySQLSource
+
+
+def _goes_like():
+    """A GOES-like dataset: (y, x) bands plus a 0-dim scalar metadata var."""
+    return xr.Dataset(
+        {
+            "CMI_C13": (("y", "x"), np.random.rand(4, 5)),
+            "DQF_C13": (("y", "x"), np.random.randint(0, 4, (4, 5))),
+            "goes_imager_projection": ((), 0),
+        },
+        coords={"y": np.arange(4), "x": np.arange(5)},
+    )
+
+
+@pytest.fixture
+def arraylake_controls(context, source_catalog):
+    return ArraylakeSourceControls(context=context, source_catalog=source_catalog)
+
+
+@pytest.fixture
+def fake_arraylake(monkeypatch):
+    """Inject a stub ``arraylake`` module so the lazy import succeeds."""
+    module = types.ModuleType("arraylake")
+    module.Client = MagicMock()
+    monkeypatch.setitem(sys.modules, "arraylake", module)
+    return module
+
+
+@pytest.mark.asyncio
+class TestArraylakeSourceControls:
+
+    async def test_fetch_data_builds_xarray_source(
+        self, arraylake_controls, fake_arraylake, monkeypatch
+    ):
+        monkeypatch.setattr(xr, "open_zarr", lambda *a, **k: _goes_like())
+        result = await arraylake_controls._fetch_data(
+            "ArraylakeSource",
+            repo="earthmover-public/goes-16",
+            branch="main",
+            group="",
+        )
+        assert isinstance(result, SourceResult)
+        assert len(result.sources) == 1
+        source = result.sources[0]
+        assert isinstance(source, XArraySQLSource)
+        assert set(source.get_tables()) == {"CMI_C13", "DQF_C13"}
+        assert result.table == "CMI_C13"
+
+    async def test_scalar_variable_dropped(
+        self, arraylake_controls, fake_arraylake, monkeypatch
+    ):
+        monkeypatch.setattr(xr, "open_zarr", lambda *a, **k: _goes_like())
+        result = await arraylake_controls._fetch_data("ArraylakeSource", repo="o/r")
+        assert "goes_imager_projection" not in result.sources[0].get_tables()
+
+    async def test_explicit_variables_override(
+        self, arraylake_controls, fake_arraylake, monkeypatch
+    ):
+        monkeypatch.setattr(xr, "open_zarr", lambda *a, **k: _goes_like())
+        arraylake_controls.variables = ["CMI_C13"]
+        result = await arraylake_controls._fetch_data("ArraylakeSource", repo="o/r")
+        assert result.sources[0].get_tables() == ["CMI_C13"]
+
+    async def test_missing_repo_returns_empty(self, arraylake_controls):
+        result = await arraylake_controls._fetch_data("ArraylakeSource", repo="")
+        assert result.sources == []
+        assert "repo" in result.message.lower()
+
+    async def test_missing_arraylake_returns_empty(
+        self, arraylake_controls, monkeypatch
+    ):
+        # Setting the module to None makes ``import arraylake`` raise ImportError.
+        monkeypatch.setitem(sys.modules, "arraylake", None)
+        result = await arraylake_controls._fetch_data("ArraylakeSource", repo="o/r")
+        assert result.sources == []
+        assert "lumen[arraylake]" in result.message
+
+    async def test_open_error_returns_empty(
+        self, arraylake_controls, fake_arraylake, monkeypatch
+    ):
+        def boom(*a, **k):
+            raise RuntimeError("connection refused")
+
+        monkeypatch.setattr(xr, "open_zarr", boom)
+        result = await arraylake_controls._fetch_data("ArraylakeSource", repo="o/r")
+        assert result.sources == []
+        assert "connection refused" in result.message
+
+
+def test_as_tools_exposes_typed_action(arraylake_controls):
+    tools = arraylake_controls.as_tools()
+    assert len(tools) == 1
+    name, func = tools[0]
+    assert name == "ArraylakeSource"
+    assert list(inspect.signature(func).parameters) == ["repo", "branch", "group"]

pyproject.toml

@@ -83,6 +83,7 @@ chromadb = ['chromadb']
 bigquery = ['google-cloud-bigquery', 'sqlalchemy-bigquery']
 xarray = ['xarray', 'xarray-sql', 'netCDF4', 'zarr']
 xarray-grib = ['lumen[xarray]', 'cfgrib', 'eccodes']
+arraylake = ['lumen[xarray]', 'arraylake', 'icechunk']
 snowflake = ['snowflake-connector-python', 'cryptography']
 ae5 = ['ae5-tools']
 mcp = ['fastmcp']