fix: infer surviving dims in to_dataset for aggregations (closes #189)

tests/test_ds.py

@@ -139,6 +139,42 @@ def test_aggregation_drops_dim(air_dataset_small):
     np.testing.assert_allclose(actual, expected)
 
 
+def test_aggregation_infers_dims(air_dataset_small):
+    """#189: to_dataset() infers the surviving GROUP BY dims without dims=."""
+    ctx = XarrayContext()
+    ctx.from_dataset("air", air_dataset_small)
+
+    # GROUP BY lat, lon: time is aggregated away, lat/lon survive.
+    out = ctx.sql(
+        "SELECT lat, lon, AVG(air) AS air_avg FROM air GROUP BY lat, lon"
+    ).to_dataset()
+    assert set(out.dims) == {"lat", "lon"}
+    assert "air_avg" in out.data_vars
+    expected = (
+        air_dataset_small.compute()
+        .sortby(["lat", "lon"])
+        .mean(dim="time")["air"]
+        .values
+    )
+    np.testing.assert_allclose(
+        out.sortby(["lat", "lon"])["air_avg"].values, expected
+    )
+
+    # The reporter's exact case: GROUP BY the time coordinate.
+    single = ctx.sql(
+        'SELECT "time", AVG("air") AS air FROM "air" GROUP BY "time"'
+    ).to_dataset()
+    assert set(single.dims) == {"time"}
+    assert "air" in single.data_vars
+    expected_t = (
+        air_dataset_small.compute()
+        .sortby("time")
+        .mean(dim=["lat", "lon"])["air"]
+        .values
+    )
+    np.testing.assert_allclose(single.sortby("time")["air"].values, expected_t)
+
+
 def test_barrier_query_scans_source_once(air_dataset_small):
     """A barrier plan (aggregation) executes the source exactly once.
 
@@ -383,14 +419,12 @@ def test_to_dataset_multi_registered_requires_explicit_template(
     assert set(out.dims) == {"time", "lat", "lon"}
 
 
-def test_to_dataset_infer_fails_when_no_template_fits(air_dataset_small):
-    """If no registered Dataset's dims fit the result -> clear error."""
+def test_to_dataset_infer_fails_when_no_dim_survives(air_dataset_small):
+    """A global aggregation leaves no registered dim in the result -> clear error."""
     ctx = XarrayContext()
     ctx.from_dataset("air", air_dataset_small)
     with pytest.raises(ValueError, match="dims cannot be inferred"):
-        ctx.sql(
-            "SELECT lat, lon, AVG(air) AS air_avg FROM air GROUP BY lat, lon"
-        ).to_dataset()
+        ctx.sql("SELECT AVG(air) AS air_avg FROM air").to_dataset()
 
 
 def test_template_accepts_name_or_dataset(air_dataset_small):

xarray_sql/ds.py

@@ -757,10 +757,12 @@ def to_dataset(
 
         Args:
             dims: Result columns to use as Dataset dimensions. When
-                ``None``, defaults to the dims of the registered Dataset
-                referenced by the SQL ``FROM`` clause (if exactly one
-                matches), or any single registered Dataset whose dims are
-                all present in the result columns.
+                ``None``, defaults to a registered Dataset's dimensions that
+                survive into the result columns, so an aggregation that drops
+                dims (e.g. ``GROUP BY time`` over a ``(time, lat, lon)`` grid)
+                round-trips on the remaining dim. Raises when no dimension
+                survives, or when several registered Datasets imply different
+                dims (pass ``dims`` explicitly then).
             template: Source to recover metadata (attrs, encoding, non-dim
                 coordinates, dim-coord dtype) from. Either an ``xr.Dataset``
                 used directly, or the name of a registered table (e.g.
@@ -879,33 +881,37 @@ def _infer_dimension_columns(
     ) -> list[str]:
         """Pick a default ``dimension_columns`` from the registry, or raise.
 
-        Uses the data variable's dim order (via :func:`_ds_var_dims`) so
-        the round-trip preserves the original axis order.
+        A registered Dataset's dims that survive into the result columns
+        become the dimensions, so aggregations that drop dims (e.g.
+        ``GROUP BY time`` over a ``(time, lat, lon)`` grid) round-trip on the
+        surviving dim(s). Uses the data variable's dim order (via
+        :func:`_ds_var_dims`) so the original axis order is preserved.
         """
         result_cols = set(self._result_columns())
-        if (
-            preferred_template is not None
-            and set(preferred_template.dims) <= result_cols
-        ):
-            return _ds_var_dims(preferred_template)
+
+        def surviving(template: xr.Dataset) -> list[str]:
+            # Template dims still present in the result, in var axis order.
+            return [d for d in _ds_var_dims(template) if d in result_cols]
+
+        if preferred_template is not None:
+            preferred = surviving(preferred_template)
+            if preferred:
+                return preferred
         if not self._templates:
             raise ValueError(
                 "dims cannot be inferred (no registered "
                 "Dataset on this result); pass dims=[...] "
                 "explicitly."
             )
-        candidates = [
-            _ds_var_dims(t)
-            for t in self._templates.values()
-            if set(t.dims) <= result_cols
-        ]
+        candidates = {tuple(surviving(t)) for t in self._templates.values()}
+        candidates.discard(())  # templates with no surviving dim
         if len(candidates) == 1:
-            return candidates[0]
+            return list(next(iter(candidates)))
         if not candidates:
             raise ValueError(
-                "dims cannot be inferred: no registered "
-                "Dataset has all of its dims present in the result "
-                "columns. Pass dims=[...] explicitly."
+                "dims cannot be inferred: no registered Dataset "
+                "dimension survives in the result columns. Pass "
+                "dims=[...] explicitly."
             )
         raise ValueError(
             "dims cannot be inferred unambiguously: multiple "