diff --git a/xarray_sql/df.py b/xarray_sql/df.py index 8b700add..644dfee5 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -150,6 +150,23 @@ def pivot(ds: xr.Dataset) -> pd.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. @@ -162,18 +179,15 @@ def read_xarray(ds: xr.Dataset, chunks: Chunks = None) -> pa.RecordBatchReader: Returns: A PyArrow Table, which is a table representation of the input Dataset. """ - 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 = list(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!" + 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." + schema = _parse_schema(ds) + blocks = block_slices(ds, chunks) return from_map_batched(pivot_block, blocks, schema=schema)