Dictionary-encode coordinate columns

A dense grid repeats every coordinate value across the whole partition (a
chunk of shape (time, lat, lon) carries each latitude time×lon times). The
reader materialized coordinate columns as dense, fully-repeated Arrow arrays,
so every GROUP BY / JOIN on a coordinate re-hashed a hugely redundant column
and the pivot moved far more bytes than the data itself.

Encode coordinate columns as Arrow dictionaries: the distinct coordinate
values are the dictionary and the strided per-row indices we already compute
are the dictionary indices — no broadcast of repeated values. The index type
is sized to the dimension length (int8/int16/int32), so a 6-step time chunk
uses 1 byte/row and a 721/1440-point lat/lon uses 2. On an ERA5-shaped chunk
the coordinate columns shrink ~4.8x (and equality GROUP BY / JOIN keys become
small integers).

- df.py: _parse_schema declares dimension coordinates as dictionary(index,
  value) with the value type/metadata preserved; iter_record_batches and
  dataset_to_record_batch emit DictionaryArrays.
- src/lib.rs: keep partition pruning and exact statistics working on the new
  encoding — DataFusion coerces a coordinate filter to either a Dictionary
  literal (timestamp) or Cast(col AS value_type) (float), so compare_to_scalar
  unwraps Dictionary scalars, the pruning matchers strip decode casts, and
  bound_to_scalar / total_byte_size unwrap the dictionary value type.
- tests: pin the encoding contract and the group-by round-trip; update schema
  and memory-characterization expectations to the (smaller) encoded columns.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019VuSeCio99NcME5eubcN3N

Author: claude

src/lib.rs

@@ -93,6 +93,13 @@ impl ScalarBound {
     /// Compare this bound with a DataFusion ScalarValue.
     /// Returns None if types are incompatible.
     fn compare_to_scalar(&self, scalar: &ScalarValue) -> Option<std::cmp::Ordering> {
+        // Coordinate columns are dictionary-encoded, so DataFusion coerces the
+        // comparison literal to a `Dictionary` scalar (the column stays a bare
+        // Column). Unwrap it to the underlying value and compare against that,
+        // so partition pruning keeps working on dictionary coordinates.
+        if let ScalarValue::Dictionary(_, value) = scalar {
+            return self.compare_to_scalar(value);
+        }
         match (self, scalar) {
             // Integer comparisons
             (ScalarBound::Int64(a), ScalarValue::Int64(Some(b))) => Some(a.cmp(b)),
@@ -265,8 +272,11 @@ impl PrunableStreamingTable {
         right: &Expr,
         meta: &PartitionMetadata,
     ) -> bool {
-        // Try to extract column and literal from either side
-        let (col_name, scalar, flipped) = match (left, right) {
+        // Try to extract column and literal from either side. Coordinate
+        // columns are dictionary-encoded, so a filter may arrive as
+        // `Cast(col AS value_type) op literal`; strip the (lossless) decode cast
+        // to reach the column.
+        let (col_name, scalar, flipped) = match (strip_cast(left), strip_cast(right)) {
             (Expr::Column(c), Expr::Literal(s, _)) => (c.name.clone(), s, false),
             (Expr::Literal(s, _), Expr::Column(c)) => (c.name.clone(), s, true),
             _ => return false, // Not a simple column-literal comparison
@@ -348,8 +358,8 @@ impl PrunableStreamingTable {
             return false;
         }
 
-        // Extract column name
-        let col_name = match between.expr.as_ref() {
+        // Extract column name (peeling any dictionary-decode cast)
+        let col_name = match strip_cast(between.expr.as_ref()) {
             Expr::Column(c) => c.name.clone(),
             _ => return false,
         };
@@ -387,8 +397,8 @@ impl PrunableStreamingTable {
             return false;
         }
 
-        // Extract column name
-        let col_name = match in_list.expr.as_ref() {
+        // Extract column name (peeling any dictionary-decode cast)
+        let col_name = match strip_cast(in_list.expr.as_ref()) {
             Expr::Column(c) => c.name.clone(),
             _ => return false,
         };
@@ -440,13 +450,33 @@ impl PrunableStreamingTable {
 
     /// Check if an expression references a dimension column.
     fn expr_references_dimension(&self, expr: &Expr) -> bool {
-        match expr {
+        match strip_cast(expr) {
             Expr::Column(c) => self.dimension_columns.contains(&c.name),
             _ => false,
         }
     }
 }
 
+/// Peel `Cast`/`TryCast` wrappers to reach the inner expression.
+///
+/// Dictionary-encoded coordinate columns are decoded with a lossless
+/// `CAST(col AS value_type)` before a comparison, so a filter on a coordinate
+/// can arrive as `Cast(Column) op literal`. Stripping the cast lets pruning see
+/// the underlying column; the cast targets the dictionary's own value type, so
+/// it never changes the compared value. If a cast were ever lossy, the bound
+/// comparison would simply fail to match and the partition is conservatively
+/// kept — pruning never becomes wrong, only (at worst) less effective.
+fn strip_cast(expr: &Expr) -> &Expr {
+    let mut current = expr;
+    loop {
+        match current {
+            Expr::Cast(cast) => current = &cast.expr,
+            Expr::TryCast(cast) => current = &cast.expr,
+            _ => return current,
+        }
+    }
+}
+
 /// Extension trait for partition streams that support column projection.
 ///
 /// Implemented by `PyArrowStreamPartition` so that `PrunableStreamingTable`
@@ -702,6 +732,12 @@ fn fold_bound(a: &ScalarBound, b: &ScalarBound, keep_min: bool) -> Option<Scalar
 /// unit we can't scale exactly), in which case the column is left without
 /// min/max rather than risk a wrong value.
 fn bound_to_scalar(bound: &ScalarBound, dtype: &DataType) -> Option<ScalarValue> {
+    // Coordinate columns are dictionary-encoded; their min/max is the min/max of
+    // the underlying values, so report a value-type scalar (what DataFusion uses
+    // for dictionary-column statistics).
+    if let DataType::Dictionary(_, value_type) = dtype {
+        return bound_to_scalar(bound, value_type);
+    }
     match (bound, dtype) {
         (ScalarBound::Int64(v), DataType::Int64) => Some(ScalarValue::Int64(Some(*v))),
         (ScalarBound::Int64(v), DataType::Int32) => {
@@ -736,6 +772,19 @@ fn bound_to_scalar(bound: &ScalarBound, dtype: &DataType) -> Option<ScalarValue>
     }
 }
 
+/// Fixed byte width of one value of `dtype`, or `None` if it is variable-width.
+///
+/// Dictionary-encoded coordinate columns are sized by their *value* type: the
+/// statistic estimates the logical data volume (a conservative upper bound on
+/// the encoded column's real footprint), and keeps the reported size stable
+/// whether or not a column is dictionary-encoded.
+fn fixed_width(dtype: &DataType) -> Option<usize> {
+    match dtype {
+        DataType::Dictionary(_, value_type) => fixed_width(value_type),
+        other => other.primitive_width(),
+    }
+}
+
 /// Exact in-memory byte size of `num_rows` rows of `schema`, or `Absent` if any
 /// column is variable-width (e.g. Utf8) and cannot be sized from the row count
 /// alone. Our data model is dense fixed-width grids, so this is normally exact.
@@ -745,7 +794,7 @@ fn total_byte_size(schema: &Schema, num_rows: &Precision<usize>) -> Precision<us
     };
     let mut row_width = 0usize;
     for field in schema.fields() {
-        match field.data_type().primitive_width() {
+        match fixed_width(field.data_type()) {
             Some(w) => row_width += w,
             None => return Precision::Absent,
         }

tests/test_cft.py

@@ -153,18 +153,23 @@ class TestParseSchemaIntegration:
     def test_noleap_produces_timestamp_us(self, rasm_ds):
         schema = _parse_schema(rasm_ds[["Tair"]])
         time_field = schema.field("time")
-        assert time_field.type == pa.timestamp("us")
+        # Coordinate columns are dictionary-encoded; the cftime encoding lives
+        # in the dictionary's value type (and the field metadata is preserved).
+        assert pa.types.is_dictionary(time_field.type)
+        assert time_field.type.value_type == pa.timestamp("us")
         assert time_field.metadata[b"xarray:calendar"] == b"noleap"
 
     def test_360day_produces_int64(self, ds_360day):
         schema = _parse_schema(ds_360day)
         time_field = schema.field("time")
-        assert time_field.type == pa.int64()
+        assert pa.types.is_dictionary(time_field.type)
+        assert time_field.type.value_type == pa.int64()
         assert time_field.metadata[b"xarray:calendar"] == b"360_day"
 
     def test_datetime64_unchanged(self):
         ds = xr.tutorial.open_dataset("air_temperature")
         schema = _parse_schema(ds)
         time_field = schema.field("time")
-        assert pa.types.is_timestamp(time_field.type)
+        assert pa.types.is_dictionary(time_field.type)
+        assert pa.types.is_timestamp(time_field.type.value_type)
         assert time_field.metadata is None  # no xarray: metadata for native

tests/test_df.py

@@ -198,25 +198,28 @@ def test_dataset_to_record_batch_matches_pivot(air_small):
     we sort by the coordinate columns before comparing.
     """
     schema = _parse_schema(air_small)
+    schema_cols = [f.name for f in schema]
     dim_cols = [f.name for f in schema if f.name in air_small.dims]
     blocks = list(
         block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4})
     )
 
     for block in blocks:
         ds_block = air_small.isel(block)
-        actual_df = (
-            dataset_to_record_batch(ds_block, schema)
-            .to_pandas()
-            .sort_values(dim_cols)
-            .reset_index(drop=True)
-        )
         expected_df = (
-            pa.RecordBatch.from_pandas(pivot(ds_block), schema=schema)
-            .to_pandas()
+            pivot(ds_block)[schema_cols]
             .sort_values(dim_cols)
             .reset_index(drop=True)
         )
+        actual_df = dataset_to_record_batch(ds_block, schema).to_pandas()[
+            schema_cols
+        ]
+        # Coordinate columns are dictionary-encoded, so they come back from
+        # Arrow as pandas Categorical; decode them to the plain dtype pivot uses
+        # before comparing values.
+        for col in dim_cols:
+            actual_df[col] = actual_df[col].astype(expected_df[col].dtype)
+        actual_df = actual_df.sort_values(dim_cols).reset_index(drop=True)
 
         pd.testing.assert_frame_equal(actual_df, expected_df, check_like=False)
 
@@ -403,19 +406,21 @@ def test_read_xarray_loads_one_chunk_at_a_time(large_ds):
             peaks.append(cur_peak)
 
         for size in sizes:
-            # Observed range: 1.59–1.83× on macOS, up to ~2.7× on Linux
-            # (glibc + Arrow allocate more intermediate buffers).
-            # iter_record_batches holds data-variable arrays (≈1× chunk) while
-            # yielding sub-batches, plus the current Arrow batch (≈0.65× chunk).
-            assert chunk_size * 1.3 < size, f"size {size} unexpectedly low"
+            # iter_record_batches holds the data-variable arrays (≈1× chunk)
+            # while yielding sub-batches, plus the current Arrow batch. The
+            # batch's coordinate columns are dictionary-encoded (only the
+            # distinct values plus small int32 indices), so they add far less
+            # than a broadcast column would — steady state sits just above 1×.
+            assert chunk_size * 1.1 < size, f"size {size} unexpectedly low"
             assert chunk_size * 3.5 > size, f"size {size} unexpectedly high"
 
         for peak in peaks:
-            # Observed range: 1.84–3.28× on macOS, up to ~4.15× on Linux
-            # (glibc + Arrow hold more intermediate buffers at peak).
-            # Peak includes data arrays + Arrow batch + temporary coordinate index
-            # arrays; the first batch of each chunk is highest (Dask compute overhead).
-            assert chunk_size * 1.5 < peak, f"peak {peak} unexpectedly low"
+            # Peak includes the data arrays + the current Arrow batch +
+            # temporary coordinate index arrays; the first batch of each chunk
+            # is highest (Dask compute overhead). Dictionary-encoded coordinate
+            # columns keep the batch (and so the peak) smaller than a broadcast
+            # column would.
+            assert chunk_size * 1.3 < peak, f"peak {peak} unexpectedly low"
             assert chunk_size * 5.0 > peak, f"peak {peak} unexpectedly high"
 
         assert max(peaks) < large_ds.nbytes

tests/test_dict_coords.py

@@ -0,0 +1,74 @@
+"""Coordinate columns are dictionary-encoded.
+
+A dense grid repeats every coordinate value across the whole partition (a chunk
+of shape ``(time, lat, lon)`` carries each latitude ``time × lon`` times).
+Encoding coordinate columns as Arrow dictionaries keeps only the distinct values
+plus small integer indices, which shrinks the bytes the engine moves and lets
+``GROUP BY`` / equality ``JOIN`` on coordinates compare integer keys. These
+tests pin that the coordinates are dictionary-encoded end to end and that the
+values still round-trip correctly.
+"""
+
+import numpy as np
+import pyarrow as pa
+import xarray as xr
+
+from xarray_sql import XarrayContext
+from xarray_sql.df import _parse_schema, block_slices, iter_record_batches
+
+
+def _grid() -> xr.Dataset:
+    return xr.Dataset(
+        {"v": (("time", "lat", "lon"), np.random.rand(6, 4, 5))},
+        coords={
+            "time": xr.date_range("2020-01-01", periods=6, freq="D"),
+            "lat": np.array([-90.0, -30.0, 30.0, 90.0]),
+            "lon": np.array([0.0, 72.0, 144.0, 216.0, 288.0]),
+        },
+    )
+
+
+def test_coordinate_fields_are_dictionary_encoded():
+    """Dimension coordinates are dictionary-typed; data variables are not."""
+    schema = _parse_schema(_grid())
+    for dim in ("time", "lat", "lon"):
+        assert pa.types.is_dictionary(schema.field(dim).type), dim
+    assert not pa.types.is_dictionary(schema.field("v").type)
+
+
+def test_iter_record_batches_emits_dictionary_coords():
+    """A coordinate column arrives as a DictionaryArray with correct values."""
+    ds = _grid()
+    schema = _parse_schema(ds)
+    block = next(block_slices(ds, chunks={"time": 6}))
+    batch = next(iter_record_batches(ds.isel(block), schema, batch_size=1024))
+
+    lat = batch.column(batch.schema.names.index("lat"))
+    assert isinstance(lat, pa.DictionaryArray)
+    # The dictionary holds only the distinct latitudes; decoding reproduces the
+    # per-row values in row-major order.
+    assert lat.dictionary.to_pylist() == [-90.0, -30.0, 30.0, 90.0]
+    decoded = lat.to_numpy(zero_copy_only=False)
+    assert decoded.shape == (batch.num_rows,)
+    # First lon×... rows share lat=-90 (lat varies slower than lon in C order).
+    assert decoded[0] == -90.0
+
+    # Data variables stay plain (not dictionary-encoded).
+    v = batch.column(batch.schema.names.index("v"))
+    assert not isinstance(v, pa.DictionaryArray)
+
+
+def test_groupby_coordinate_roundtrips_through_dictionary():
+    """GROUP BY on a dictionary coordinate returns the same numbers as xarray."""
+    ds = _grid()
+    ctx = XarrayContext()
+    ctx.from_dataset("grid", ds, chunks={"time": 3})
+
+    got = ctx.sql(
+        "SELECT lat, AVG(v) AS m FROM grid GROUP BY lat ORDER BY lat"
+    ).to_dataset(dims=["lat"])
+    ref = ds["v"].mean(["time", "lon"])
+
+    xr.testing.assert_allclose(
+        got.m, ref.reindex_like(got.m), rtol=1e-6, atol=1e-9
+    )