From a63aa3d979eff963ad89a2e0b7e07ecfaca1827b Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sat, 4 Oct 2025 15:42:52 +0200 Subject: [PATCH 1/5] read_xarray now uses iterables all the way through. When working with small datasets, we materialized blocks via a `list()` call. This is slow on large datasets. So, we now work only with the iterable of the blocks. --- xarray_sql/df.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/xarray_sql/df.py b/xarray_sql/df.py index 8b700add..dc0e7eb7 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -167,13 +167,12 @@ def read_xarray(ds: xr.Dataset, chunks: Chunks = None) -> pa.RecordBatchReader: da.dims == fst for da in ds.values() ), "All dimensions must be equal. Please filter data_vars in the Dataset." - blocks = list(block_slices(ds, chunks)) + blocks = block_slices(ds, chunks) def pivot_block(b: Block): return pivot(ds.isel(b)) - schema = pa.Schema.from_pandas(pivot_block(blocks[0])) - last_schema = pa.Schema.from_pandas(pivot_block(blocks[-1])) - assert schema == last_schema, "Schemas must be consistent across blocks!" - - return from_map_batched(pivot_block, blocks, schema=schema) + head, *tail = blocks + schema = pa.Schema.from_pandas(pivot_block(head)) + joined = itertools.chain([head], tail) + return from_map_batched(pivot_block, joined, schema=schema) From a3039a61b93167dd7b0b4efe368e424d7154b1b8 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sat, 4 Oct 2025 16:16:48 +0200 Subject: [PATCH 2/5] Lighter-weight way to parse the schema. --- xarray_sql/df.py | 23 +++++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/xarray_sql/df.py b/xarray_sql/df.py index dc0e7eb7..94e868f7 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -86,7 +86,10 @@ def from_map_batched( args = () def map_batches(): - for items in zip(*iterables): + print('entering map_batched') + for i, items in enumerate(zip(*iterables)): + if i % 100 == 0: + print(f'item {i} in map_batches') df = func(*items, *args, **kwargs) yield pa.RecordBatch.from_pandas(df, schema=schema) @@ -149,6 +152,22 @@ def pivot(ds: xr.Dataset) -> pd.DataFrame: """Converts an xarray Dataset to a pandas DataFrame.""" return ds.to_dataframe().reset_index() +def _parse_schema(ds) -> pa.Schema: + """Extracts a `pa.Schema` from the Dataset, treating dims and data_vars as columns.""" + columns = [] + + for coord_name, coord_var in ds.coords.items(): + # Only include dimension coordinates + if coord_name in ds.dims: + pa_type = pa.from_numpy_dtype(coord_var.dtype) + columns.append(pa.field(coord_name, pa_type)) + + for var_name, var in ds.data_vars.items(): + pa_type = pa.from_numpy_dtype(var.dtype) + columns.append(pa.field(var_name, pa_type)) + + return pa.schema(columns) + def read_xarray(ds: xr.Dataset, chunks: Chunks = None) -> pa.RecordBatchReader: """Pivots an Xarray Dataset into a PyArrow Table, partitioned by chunks. @@ -173,6 +192,6 @@ def pivot_block(b: Block): return pivot(ds.isel(b)) head, *tail = blocks - schema = pa.Schema.from_pandas(pivot_block(head)) + schema = _parse_schema(ds) joined = itertools.chain([head], tail) return from_map_batched(pivot_block, joined, schema=schema) From a0a0b89fc9a5adb3b8666f102f10be91fad35a89 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sat, 4 Oct 2025 16:18:06 +0200 Subject: [PATCH 3/5] Better order of read_xarray function. --- xarray_sql/df.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/xarray_sql/df.py b/xarray_sql/df.py index 94e868f7..233dd190 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -181,17 +181,14 @@ def read_xarray(ds: xr.Dataset, chunks: Chunks = None) -> pa.RecordBatchReader: Returns: A PyArrow Table, which is a table representation of the input Dataset. """ + def pivot_block(b: Block): + return pivot(ds.isel(b)) + fst = next(iter(ds.values())).dims assert all( da.dims == fst for da in ds.values() ), "All dimensions must be equal. Please filter data_vars in the Dataset." - blocks = block_slices(ds, chunks) - - def pivot_block(b: Block): - return pivot(ds.isel(b)) - - head, *tail = blocks schema = _parse_schema(ds) - joined = itertools.chain([head], tail) - return from_map_batched(pivot_block, joined, schema=schema) + blocks = block_slices(ds, chunks) + return from_map_batched(pivot_block, blocks, schema=schema) From c1a2ae6cf475d8876441f0ae2e3495b40faad6d0 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sat, 4 Oct 2025 16:30:23 +0200 Subject: [PATCH 4/5] rm print debugging. --- xarray_sql/df.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/xarray_sql/df.py b/xarray_sql/df.py index 233dd190..148c96a3 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -86,10 +86,7 @@ def from_map_batched( args = () def map_batches(): - print('entering map_batched') - for i, items in enumerate(zip(*iterables)): - if i % 100 == 0: - print(f'item {i} in map_batches') + for items in zip(*iterables): df = func(*items, *args, **kwargs) yield pa.RecordBatch.from_pandas(df, schema=schema) From 0a803d9f3e44c357ddd636b9c10c4a98188ecaa8 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sat, 4 Oct 2025 16:32:37 +0200 Subject: [PATCH 5/5] Reformat --- xarray_sql/df.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/xarray_sql/df.py b/xarray_sql/df.py index 148c96a3..644dfee5 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -149,6 +149,7 @@ def pivot(ds: xr.Dataset) -> pd.DataFrame: """Converts an xarray Dataset to a pandas DataFrame.""" return ds.to_dataframe().reset_index() + def _parse_schema(ds) -> pa.Schema: """Extracts a `pa.Schema` from the Dataset, treating dims and data_vars as columns.""" columns = [] @@ -158,7 +159,7 @@ def _parse_schema(ds) -> pa.Schema: if coord_name in ds.dims: pa_type = pa.from_numpy_dtype(coord_var.dtype) columns.append(pa.field(coord_name, pa_type)) - + for var_name, var in ds.data_vars.items(): pa_type = pa.from_numpy_dtype(var.dtype) columns.append(pa.field(var_name, pa_type)) @@ -178,6 +179,7 @@ def read_xarray(ds: xr.Dataset, chunks: Chunks = None) -> pa.RecordBatchReader: Returns: A PyArrow Table, which is a table representation of the input Dataset. """ + def pivot_block(b: Block): return pivot(ds.isel(b))