Commit ade544a
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
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
69 | 69 | | |
70 | 70 | | |
71 | 71 | | |
72 | | - | |
73 | | - | |
74 | | - | |
75 | | - | |
76 | | - | |
| 72 | + | |
| 73 | + | |
| 74 | + | |
| 75 | + | |
| 76 | + | |
| 77 | + | |
| 78 | + | |
77 | 79 | | |
78 | 80 | | |
79 | 81 | | |
| |||
86 | 88 | | |
87 | 89 | | |
88 | 90 | | |
89 | | - | |
| 91 | + | |
| 92 | + | |
90 | 93 | | |
91 | 94 | | |
92 | 95 | | |
93 | 96 | | |
94 | 97 | | |
95 | 98 | | |
96 | | - | |
97 | | - | |
98 | | - | |
| 99 | + | |
| 100 | + | |
| 101 | + | |
| 102 | + | |
99 | 103 | | |
100 | 104 | | |
101 | 105 | | |
| |||
111 | 115 | | |
112 | 116 | | |
113 | 117 | | |
114 | | - | |
115 | | - | |
| 118 | + | |
| 119 | + | |
| 120 | + | |
116 | 121 | | |
117 | 122 | | |
118 | 123 | | |
| |||
0 commit comments