feat: skip per-dim discovery scans when query is an unfiltered scan (#193)

Author: ghostiee-11

tests/test_ds.py

@@ -202,6 +202,110 @@ def test_order_by_direction_sets_dim_order(air_dataset_small):
     np.testing.assert_allclose(out["air_avg"].values, expected.values)
 
 
+def test_unfiltered_scan_skips_dim_discovery(air_dataset_small):
+    """``SELECT * FROM <registered>`` does no per-dim discovery scans.
+
+    The pre-fix :func:`_build_lazy_scan` ran one full source scan per dim
+    via ``inner_df.select(col(d)).distinct().sort(...)`` -- a single-column
+    projection per dim, multiplied across the dim count. With a registered
+    template and an unfiltered plan, :func:`_maybe_template_coords` returns
+    the template's coord arrays directly and the discovery loop is skipped
+    entirely. Verified by the iteration callback's projection list: no read
+    where the projection is a single dim column.
+    """
+    from xarray_sql.reader import read_xarray_table
+
+    reads: list = []
+    table = read_xarray_table(
+        air_dataset_small,
+        chunks={"time": 6},
+        _iteration_callback=lambda block, proj: reads.append(list(proj)),
+    )
+    ctx = XarrayContext()
+    ctx.register_table("air", table)
+    ctx._registered_datasets["air"] = air_dataset_small
+
+    reads.clear()
+    out = ctx.sql('SELECT * FROM "air"').to_dataset()
+    single_dim_reads = [
+        proj for proj in reads if set(proj) <= {"time", "lat", "lon"}
+    ]
+    assert not single_dim_reads, (
+        f"fast path should skip per-dim discovery scans, got "
+        f"{single_dim_reads!r}"
+    )
+    # The fast path must not break correctness on first access.
+    out.compute()  # safety: the Dataset still resolves
+
+
+def test_filtered_scan_still_uses_dim_discovery(air_dataset_small):
+    """A ``WHERE`` clause forbids the template-coord fast path.
+
+    With a filter the result's coord extent is a subset of the source, so
+    reusing template coords would silently widen the output. The dispatch
+    in :func:`_maybe_template_coords` falls back to per-dim discovery in
+    that case, which is observable as a positive read count on the source.
+    """
+    from xarray_sql.reader import read_xarray_table
+
+    reads: list = []
+    table = read_xarray_table(
+        air_dataset_small,
+        chunks={"time": 6},
+        _iteration_callback=lambda block, proj: reads.append(block),
+    )
+    ctx = XarrayContext()
+    ctx.register_table("air", table)
+    ctx._registered_datasets["air"] = air_dataset_small
+
+    reads.clear()
+    ctx.sql('SELECT * FROM "air" WHERE lat > 30').to_dataset()
+    assert reads, "filtered scan must hit the source to discover coord extent"
+
+
+def test_fast_path_uses_scanned_tables_coords_not_user_template(
+    air_dataset_small,
+):
+    """Fast path sources coord values from the scanned table, not ``template=``.
+
+    With multiple registered Datasets, a user may pass ``template=other``
+    for metadata recovery while the query scans a different registered
+    table. The coord values must come from the **scanned** table's
+    registered Dataset; using ``other``'s coords would silently widen the
+    output if their coord ranges differ.
+    """
+    other = air_dataset_small.isel(lat=slice(0, 5))
+    assert other.sizes["lat"] != air_dataset_small.sizes["lat"]
+    ctx = XarrayContext()
+    ctx.from_dataset("air", air_dataset_small, chunks={"time": 24})
+    ctx.from_dataset("other", other, chunks={"time": 24})
+    out = ctx.sql('SELECT * FROM "air"').to_dataset(
+        dims=["time", "lat", "lon"], template=other
+    )
+    # Scanned table is "air", so lat must match air's full lat axis.
+    np.testing.assert_array_equal(
+        out["lat"].values, air_dataset_small["lat"].values
+    )
+
+
+def test_round_trip_preserves_descending_lat_on_lazy_path(air_dataset_small):
+    """Lazy round-trip preserves source dim order (xarray-sql#171).
+
+    NCEP ``air_temperature`` ships descending lat (75.0 -> 15.0). The
+    discovery path's ``.distinct().sort()`` previously flipped lat to
+    ascending on the lazy result. The template-coord fast path returns the
+    source's coord arrays as-is, so the descending order survives.
+    """
+    ds = air_dataset_small
+    assert (np.diff(ds["lat"].values) < 0).all(), (
+        "test relies on a descending-lat fixture"
+    )
+    ctx = XarrayContext()
+    ctx.from_dataset("air", ds, chunks={"time": 24})
+    lazy = ctx.sql('SELECT * FROM "air"').to_dataset()
+    np.testing.assert_array_equal(lazy["lat"].values, ds["lat"].values)
+
+
 def test_chunks_argument_controls_partitioning(synthetic_dataset):
     """``chunks`` controls eager-vs-chunked and inherits the source grid.
 

xarray_sql/ds.py

@@ -432,33 +432,115 @@ def _materialize(
     return xr.Dataset(data_vars=data_vars, coords=coords_arg)
 
 
+_PURE_SCAN_NODES = {"Projection", "Sort", "TableScan", "SubqueryAlias"}
+
+
+def _unfiltered_scan_table(inner_df: Any) -> str | None:
+    """Return the scanned table name iff the query is a pure unfiltered scan.
+
+    A pure scan only contains ``Projection``, ``Sort``, ``TableScan``,
+    ``SubqueryAlias`` nodes and exactly one ``TableScan``. Anything else
+    (``Filter``, ``Aggregate``, ``Join``, ``Union``, ``Limit``, multi-table
+    joins, ...) returns ``None`` so the caller falls back to per-dim
+    discovery. The returned name is the registered table the caller can
+    look up to source coord arrays from.
+    """
+    try:
+        lp = inner_df.logical_plan()
+    except Exception:
+        return None
+    table_name: str | None = None
+    stack = [lp]
+    while stack:
+        node = stack.pop()
+        try:
+            variant = node.to_variant()
+        except Exception:
+            return None
+        cls = type(variant).__name__
+        if cls not in _PURE_SCAN_NODES:
+            return None
+        if cls == "TableScan":
+            try:
+                this = variant.table_name()
+            except (AttributeError, TypeError):
+                return None
+            if not isinstance(this, str):
+                return None
+            if table_name is not None and table_name != this:
+                return None  # multi-table scan; not a single source
+            table_name = this
+        stack.extend(node.inputs())
+    return table_name
+
+
+def _maybe_template_coords(
+    templates: dict[str, xr.Dataset] | None,
+    dimension_columns: list[str],
+    inner_df: Any,
+) -> dict[str, np.ndarray] | None:
+    """Use the scanned table's registered coord arrays directly when safe.
+
+    Returns coord arrays sourced from the registered Dataset for the
+    scanned table iff the query is an unfiltered scan over that single
+    table and the registered Dataset carries all requested dims. Returns
+    ``None`` otherwise so the caller falls back to per-dim discovery.
+    Skipping discovery avoids one full plan execution per dim and
+    preserves the source's coordinate order (xarray-sql#171).
+
+    Coord values come from the **scanned** registered Dataset, not from
+    any user-supplied ``template=`` (which is for metadata recovery
+    only). That keeps the fast path correct when a user with multiple
+    registered Datasets passes a metadata template that differs from
+    the query's source.
+    """
+    if not templates:
+        return None
+    table = _unfiltered_scan_table(inner_df)
+    if table is None or table not in templates:
+        return None
+    source = templates[table]
+    if not all(d in source.coords for d in dimension_columns):
+        return None
+    return {d: np.asarray(source.coords[d].values) for d in dimension_columns}
+
+
 def _build_lazy_scan(
     inner_df: Any,
     dimension_columns: list[str],
     field_names: list[str],
     field_types: dict[str, Any],
+    templates: dict[str, xr.Dataset] | None = None,
 ) -> xr.Dataset:
     """Build a lazy Dataset whose data vars are :class:`SQLBackendArray`.
 
     Used when output chunking is requested: each data variable stays lazy and,
     once wrapped by ``Dataset.chunk``, every chunk reads its coordinate range via
-    a pushdown filter on first access. Coordinates are discovered per dim via
+    a pushdown filter on first access. Coordinates come either from the
+    scanned table's registered Dataset (fast path, for unfiltered scans -- see
+    :func:`_maybe_template_coords`) or from per-dim
     ``inner_df.select(col(d)).distinct().sort(...).execute_stream()``; the table
     provider projects to that single coordinate column and skips data variables,
     so discovery reads coordinate values only (no data-variable I/O).
     """
-    coord_arrays: dict[str, np.ndarray] = {}
-    for d in dimension_columns:
-        dim_only = (
-            inner_df.select(col(f'"{d}"')).distinct().sort(col(f'"{d}"').sort())
-        )
-        chunks = [b.to_pyarrow() for b in dim_only.execute_stream()]
-        if not chunks:
-            coord_arrays[d] = np.asarray([])
-            continue
-        coord_arrays[d] = np.concatenate(
-            [c.column(0).to_numpy(zero_copy_only=False) for c in chunks]
-        )
+    coord_arrays = _maybe_template_coords(
+        templates, dimension_columns, inner_df
+    )
+    if coord_arrays is None:
+        coord_arrays = {}
+        for d in dimension_columns:
+            dim_only = (
+                inner_df.select(col(f'"{d}"'))
+                .distinct()
+                .sort(col(f'"{d}"').sort())
+            )
+            chunks = [b.to_pyarrow() for b in dim_only.execute_stream()]
+            if not chunks:
+                coord_arrays[d] = np.asarray([])
+                continue
+            coord_arrays[d] = np.concatenate(
+                [c.column(0).to_numpy(zero_copy_only=False) for c in chunks]
+            )
     shape = tuple(len(coord_arrays[d]) for d in dimension_columns)
 
     data_vars: dict[str, xr.Variable] = {}
@@ -550,6 +632,7 @@ def _result_to_xarray(
     sparsity: Sparsity,
     fill_value: Any,
     chunks: Mapping[str, int] | str | None,
+    templates: dict[str, xr.Dataset] | None = None,
 ) -> xr.Dataset:
     """Reconstruct an ``xr.Dataset`` from a SQL result.
 
@@ -583,7 +666,11 @@ def _result_to_xarray(
         ds = _materialize(inner_df, dimension_columns, field_names, field_types)
     else:
         ds = _build_lazy_scan(
-            inner_df, dimension_columns, field_names, field_types
+            inner_df,
+            dimension_columns,
+            field_names,
+            field_types,
+            templates=templates,
         )
 
     if sparsity == "template":
@@ -730,6 +817,7 @@ def to_dataset(
             sparsity=sparsity,
             fill_value=fill_value,
             chunks=resolved_chunks,
+            templates=self._templates,
         )
 
     # ------------------------------------------------------------------