diff --git a/tests/test_ds.py b/tests/test_ds.py index b674239..aa3deb2 100644 --- a/tests/test_ds.py +++ b/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 `` 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. diff --git a/xarray_sql/ds.py b/xarray_sql/ds.py index c33f996..5dfdf42 100644 --- a/xarray_sql/ds.py +++ b/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, ) # ------------------------------------------------------------------