From 03d0c7566d75f34f886c6fa098fc2c0ca19acd97 Mon Sep 17 00:00:00 2001
From: ghostiee-11 <168410465+ghostiee-11@users.noreply.github.com>
Date: Fri, 24 Jul 2026 00:49:11 +0530
Subject: [PATCH] fix: cftime coord dropped by a projection crashes batch
 reading

iter_record_batches preloaded every dim's coord values and called
schema.field(name) for cftime coords, but a projection (e.g. GROUP BY on a
non-time dim) drops columns from the batch schema, so the lookup raised
KeyError. Skip dims absent from the projected schema; they are never read in
the batch loop, which iterates only the schema's fields.
---
 tests/test_df.py | 24 ++++++++++++++++++++++++
 xarray_sql/df.py |  7 +++++++
 2 files changed, 31 insertions(+)

diff --git a/tests/test_df.py b/tests/test_df.py
index 5ba5b3f..aec174a 100644
--- a/tests/test_df.py
+++ b/tests/test_df.py
@@ -181,6 +181,30 @@ def test_iter_record_batches_matches_dataset_to_record_batch(air_small):
     pd.testing.assert_frame_equal(actual_df, expected_df)
 
 
+def test_iter_record_batches_projection_drops_cftime_dim():
+    """A projection that drops a cftime dim (e.g. time under GROUP BY level)
+    must not call schema.field() for it. The dim is absent from the projected
+    schema, and cftime coords take the convert_for_field(schema.field(name))
+    path, so an unguarded lookup raised KeyError during batch reading."""
+    cftime = pytest.importorskip("cftime")
+    times = np.array(
+        [cftime.DatetimeGregorian(2020, m, 1) for m in (1, 2, 3)], dtype=object
+    )
+    ds = xr.Dataset(
+        {"air": (["time", "lat"], np.arange(3 * 2, dtype=float).reshape(3, 2))},
+        coords={"time": times, "lat": [0.0, 1.0]},
+    )
+    full = _parse_schema(ds)
+    projected = pa.schema(
+        [full.field("lat"), full.field("air")]
+    )  # time dropped
+    table = pa.Table.from_batches(
+        list(iter_record_batches(ds, projected, batch_size=16)), projected
+    )
+    assert table.schema.names == ["lat", "air"]
+    assert table.num_rows == 6
+
+
 def test_iter_record_batches_default_batch_size():
     """A single-batch partition (rows <= DEFAULT_BATCH_SIZE) yields exactly one batch."""
     ds = xr.tutorial.open_dataset("air_temperature").isel(time=slice(0, 2))
diff --git a/xarray_sql/df.py b/xarray_sql/df.py
index 88aa06f..cd5e432 100644
--- a/xarray_sql/df.py
+++ b/xarray_sql/df.py
@@ -339,7 +339,14 @@ def iter_record_batches(
     # Preload small 1-D coordinate arrays (negligible memory).
     # Convert cftime objects to numeric values matching the schema type.
     coord_values = {}
+    schema_names = set(schema.names)
     for name in dim_names:
+        # A dim the projection dropped (e.g. time under GROUP BY level) is never
+        # read in the batch loop below, which only iterates the schema's fields.
+        # Skip it so schema.field(name) is not called for a projected-away name
+        # (it raises for cftime coords, which take the convert_for_field path).
+        if name not in schema_names:
+            continue
         vals = ds.coords[name].values
         if cft.is_cftime(vals):
             coord_values[name] = cft.convert_for_field(vals, schema.field(name))