Lazy read_xarray with iterables and parsing metadata.

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)