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))