Skip to content

Commit ade544a

Browse files
committed
Rearchitect lazy backend on DataFusion DataFrame API + Arrow streaming
Replaces the SQL-string-subquery + pandas implementation with native DataFusion DataFrame operations (df.filter, df.select, execute_stream) and Arrow RecordBatch -> numpy scatter throughout. Addresses every concern from the round-2 review on #167 plus the round-1 cleanup items. Production data path: - Slab access (SQLBackendArray._raw_getitem): Was: build "SELECT cols FROM (base_query) AS _xql_base WHERE ..." string, _raw_sql(...).to_pandas(), set_index().to_xarray(), .sel for reorder. Now: filter the wrapper's live DataFrame with typed Expr objects (col(d) == literal(v), Or-chained for IN-equivalent), select the needed columns, execute_stream, scatter Arrow RecordBatches into a preallocated numpy buffer via np.searchsorted-based per-dim position lookup. No SQL strings, no pandas. - Coord-array discovery (_lazy_to_xarray): Was: regex-based "bare SELECT *" detection chose between template coords and per-dim "SELECT DISTINCT d FROM (base) ORDER BY d" string SQL. Now: per-dim inner_df.select(col(d)).distinct().sort(...) chains executed via execute_stream into Arrow Arrays. Zero regex, zero string SQL. - Eager / aggregation path (_eager_to_xarray): Was: _normalize_to_pandas -> set_index().to_xarray(). Now: _to_arrow_table consumes XarrayDataFrame and DataFusion DataFrame inputs via execute_stream; per-dim coord discovery via pyarrow.compute; duplicate detection by Python-set scan over dim tuples (pyarrow.compute.unique lacks a struct kernel in 23.0.0); scatter via the same Arrow -> numpy helper. Deleted (no longer needed): - _FROM_OR_JOIN_RE, _UNFILTERED_SELECT_STAR_RE, _extract_from_tables. - _sql_literal, _build_cond, _raw_sql. - SQLBackendArray._fallback_materialize_and_index. - import re from ds.py. Renamed and cleaned up (round-1 review feedback): - dim_cols parameter -> dimension_columns (xarray-native naming). - Test module docstring scrubbed of "Phase 1/2/3" and #58 references. - ds.py top-line is now a declarative sentence for autodoc. - _RegistryView docstring rewritten to say "the registered Datasets" rather than coin "registration." - SparseExtent now has an explanatory comment above the Literal type. - _apply_template docstring explains why dtype-bound encoding keys are stripped (SQL can change a column's dtype; reattaching the source packing would corrupt a later to_netcdf write). - _apply_template now also copies dim-coordinate attrs (standard_name, long_name, units, ...) so the round-trip is identical down to coord metadata. Behavior changes (intentional): - _resolve_template no longer parses SQL. Auto-resolution works only when exactly one Dataset is registered on the context; multi-registered users pass template_table= explicitly. - The bare-SELECT-* fast path is gone; coord discovery is uniform. Test consolidation: - 8 narrow round-trip / metadata tests collapsed into a single parametrized property test test_round_trip_identity over air_dataset_small, weather_dataset, and synthetic_dataset, using xr.testing.assert_identical (the equality system the maintainer pointed at). It exercises values, dims, coord values, dim-coord dtype, dim-coord attrs, non-dim coords, and dataset attrs in one check per fixture. - Kept the distinct cases that the property test does not cover: encoding dtype-bound-key strip, aggregation alias gets no attrs, template_table explicit override, mutual-exclusion error. - test_extract_from_tables (8 parametrized) deleted with the regex. - test_bare_select_star_skips_distinct_scans deleted with the removed optimization. - test_lazy_isel_int_pushes_down_equality and test_full_dim_slice_omits_filter_for_full_dims rewritten to patch SQLBackendArray._raw_getitem / DataFrame.filter (the new entry points) instead of the gone _raw_sql. - super(type(ctx), ctx) replaced with datafusion.SessionContext.sql. - repr(result) == result._inner private-attr assertion dropped. Memory (RSS, fresh process, air_temperature ~31 MB): - Lazy construction: ~+2 MB (unchanged). - Lazy slab access (isel(time=0).values): ~+2.3 MB (was ~+17 MB on the previous subquery+pandas path; ~7x reduction). Net diff on this branch: +492 / -543 = -51 lines despite adding the Arrow scatter helper, the Arrow-native eager path, and the new parametrized property test. Refs #58.
1 parent f50ba86 commit ade544a

4 files changed

Lines changed: 492 additions & 543 deletions

File tree

README.md

Lines changed: 16 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -69,11 +69,13 @@ SQL queries against them.
6969
## Round-tripping back to Xarray
7070

7171
`ctx.sql(...)` returns an `XarrayDataFrame` that exposes `.to_pandas()`
72-
(unchanged) and a new `.to_dataset()` for converting the result back to an
73-
`xr.Dataset`. The reverse path is **lazy by default**: the returned Dataset
74-
is backed by a `BackendArray` that translates xarray indexing into SQL
75-
`WHERE` clauses pushed into a sub-query of the original SQL. Data only
76-
materializes for the slice actually accessed.
72+
(unchanged) and a new `.to_dataset()` for converting the result back into
73+
an `xr.Dataset`. The reverse path is **lazy by default**: the returned
74+
Dataset is backed by an `xarray.backends.BackendArray` that translates
75+
xarray indexers into DataFusion `filter` expressions and consumes the
76+
filtered DataFrame via `execute_stream`. Arrow `RecordBatch` es scatter
77+
directly into a preallocated numpy buffer with no pandas hop, so only
78+
the slab actually accessed is materialized.
7779

7880
```python
7981
out = ctx.sql('SELECT * FROM "air"').to_dataset()
@@ -86,16 +88,18 @@ out = ctx.sql('SELECT * FROM "air"').to_dataset()
8688
# Data variables:
8789
# air (time, lat, lon) float32 ...
8890

89-
# Slicing pushes down into SQL; only the requested slab is materialized.
91+
# Slicing pushes down into DataFusion; only the requested slab is
92+
# materialized.
9093
slab = out["air"].isel(time=0).values
9194

9295
# For full eager materialization, call .compute().
9396
eager = out.compute()
9497
```
9598

96-
`dim_cols` defaults to the dims of the registered Dataset referenced by the
97-
SQL `FROM` clause. Variable attrs, dataset attrs, non-dimension coords, and
98-
dim-coordinate dtypes are recovered from the registered Dataset
99+
`dimension_columns` defaults to the dims of the single registered Dataset
100+
on the context (or the one named via `template_table=` when several are
101+
registered). Variable attrs, dataset attrs, non-dimension coordinates,
102+
and dim-coordinate dtype are recovered from the registered Dataset
99103
automatically.
100104

101105
For filtered queries that return only part of the original extent, pass
@@ -111,8 +115,9 @@ out = ctx.sql(
111115

112116
Aggregation queries (e.g. `AVG(air) AS air_avg ... GROUP BY lat, lon`)
113117
materialize once because their output does not align with the source dim
114-
structure. Pass `dim_cols=[...]` explicitly when an aggregation drops a
115-
dim.
118+
structure; the aggregation path is also Arrow-native (no pandas
119+
intermediates). Pass `dimension_columns=[...]` explicitly when an
120+
aggregation drops a dim.
116121

117122
## Why build this?
118123

0 commit comments

Comments
 (0)