Simplifying unravel logic

perf_tests/compute_air.py

@@ -8,6 +8,6 @@
   chunks = {'time': 240}
   air = air.chunk(chunks)
 
-  df = qr.to_dd(air).compute()
+  df = qr.read_xarray(air).compute()
 
   print(len(df))

perf_tests/groupby_air.py

@@ -13,7 +13,7 @@
     time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10)
   ).chunk(chunks)
 
-  df = qr.to_dd(air_small)
+  df = qr.read_xarray(air_small)
 
   c = Context()
   c.create_table('air', df)

perf_tests/groupby_air_full.py

@@ -10,7 +10,7 @@
   chunks = {'time': 240}
   air = air.chunk(chunks)
 
-  df = qr.to_dd(air)
+  df = qr.read_xarray(air)
 
   c = Context()
   c.create_table('air', df)

perf_tests/open_era5.py

@@ -8,6 +8,6 @@
   'gs://gcp-public-data-arco-era5/ar/1959-2022-full_37-1h-0p25deg-chunk-1.zarr-v2',
   chunks={'time': 240, 'level': 1}
 )
-era5_wind_df = qr.to_dd(era5_ds[['u_component_of_wind', 'v_component_of_wind']])
+era5_wind_df = qr.read_xarray(era5_ds[['u_component_of_wind', 'v_component_of_wind']])
 
 print(era5_wind_df.columns)

perf_tests/sanity.py

@@ -11,6 +11,6 @@
     time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10)
   ).chunk(chunks)
 
-  df = qr.to_dd(air_small).compute()
+  df = qr.read_xarray(air_small).compute()
 
   print(len(df))

qarray/__init__.py

@@ -1,3 +1,2 @@
-from .core import unravel
-from .df import to_pd, to_dd
+from .df import read_xarray
 

qarray/core.py

@@ -11,6 +11,7 @@ def get_columns(ds: xr.Dataset) -> t.List[str]:
   return list(ds.dims.keys()) + list(ds.data_vars.keys())
 
 
+# Deprecated
 def unravel(ds: xr.Dataset) -> t.Iterator[Row]:
   dim_keys, dim_vals = zip(*ds.dims.items())
 
@@ -22,6 +23,7 @@ def unravel(ds: xr.Dataset) -> t.Iterator[Row]:
     yield row
 
 
+# Deprecated
 def unbounded_unravel(ds: xr.Dataset) -> np.ndarray:
   """Unravel with unbounded memory (as a NumPy Array)."""
   dim_keys, dim_vals = zip(*ds.dims.items())

qarray/df.py

@@ -11,7 +11,7 @@
 from . import core
 
 Block = t.Dict[str, slice]
-Chunks = t.Dict[str, int]
+Chunks = t.Optional[t.Dict[str, int]]
 
 # Turn on Dask-Expr
 dask.config.set({'dataframe.query-planning-warning': False})
@@ -31,10 +31,7 @@ def _get_chunk_slicer(dim: t.Hashable, chunk_index: t.Mapping,
 
 
 # Adapted from Xarray `map_blocks` implementation.
-def block_slices(
-    ds: xr.Dataset,
-    chunks: t.Optional[Chunks] = None
-) -> t.Iterator[Block]:
+def block_slices(ds: xr.Dataset, chunks: Chunks = None) -> t.Iterator[Block]:
   """Compute block slices for a chunked Dataset."""
   if chunks is not None:
     for_chunking = ds.copy(data=None, deep=False).chunk(chunks)
@@ -63,32 +60,17 @@ def block_slices(
   yield from blocks
 
 
-def explode(
-    ds: xr.Dataset,
-    chunks: t.Optional[Chunks] = None
-) -> t.Iterator[xr.Dataset]:
+def explode(ds: xr.Dataset, chunks: Chunks = None) -> t.Iterator[xr.Dataset]:
   """Explodes a dataset into its chunks."""
   yield from (ds.isel(b) for b in block_slices(ds, chunks=chunks))
 
 
-def to_pd(ds: xr.Dataset, bounded=True) -> pd.DataFrame:
-  columns = core.get_columns(ds)
-  if bounded:
-    df = pd.DataFrame(core.unravel(ds), columns=columns)
-    for c in columns:
-      df[c] = df[c].astype(ds[c].dtype)
-    return df
-  else:
-    data = core.unbounded_unravel(ds)
-    return pd.DataFrame.from_records(data)
-
-
 def _block_len(block: Block) -> int:
   return np.prod([v.stop - v.start for v in block.values()])
 
 
-def to_dd(ds: xr.Dataset, chunks: t.Optional[Chunks] = None) -> dd.DataFrame:
-  """Unravel a Dataset into a Dataframe, partitioned by chunks.
+def read_xarray(ds: xr.Dataset, chunks: Chunks = None) -> dd.DataFrame:
+  """Pivots an Xarray Dataset into a Dask Dataframe, partitioned by chunks.
 
   Args:
     ds: An Xarray Dataset. All `data_vars` mush share the same dimensions.
@@ -104,8 +86,8 @@ def to_dd(ds: xr.Dataset, chunks: t.Optional[Chunks] = None) -> dd.DataFrame:
   block_lengths = [_block_len(b) for b in blocks]
   divisions = tuple(np.cumsum([0] + block_lengths))  # 0 ==> start partition.
 
-  def f(b: Block) -> pd.DataFrame:
-    return to_pd(ds.isel(b), bounded=False)
+  def pivot(b: Block) -> pd.DataFrame:
+    return ds.isel(b).to_dataframe().reset_index()
 
   # Token is needed to prevent Dask from spending too many cycles calculating
   # it's own token from the constituent parts.
@@ -124,7 +106,7 @@ def f(b: Block) -> pd.DataFrame:
   }
 
   return from_map(
-    f,
+    pivot,
     blocks,
     meta=meta,
     divisions=divisions,

qarray/df_integration_test.py

@@ -2,7 +2,7 @@
 
 import xarray as xr
 
-from .df import to_dd
+from . import read_xarray
 
 
 class Era5TestCast(unittest.TestCase):
@@ -11,7 +11,7 @@ def test_open_era5(self):
       'gs://gcp-public-data-arco-era5/ar/1959-2022-full_37-1h-0p25deg-chunk-1.zarr-v2',
       chunks={'time': 240, 'level': 1}
     )
-    era5_wind_df = to_dd(era5_ds[['u_component_of_wind', 'v_component_of_wind']])
+    era5_wind_df = read_xarray(era5_ds[['u_component_of_wind', 'v_component_of_wind']])
 
     self.assertEqual(list(era5_wind_df.columns), [
       'time', 'level', 'latitude', 'longitude',

qarray/df_test.py

@@ -5,7 +5,7 @@
 import numpy as np
 import xarray as xr
 
-from .df import explode, to_dd, block_slices
+from .df import explode, read_xarray, block_slices
 
 
 class DaskTestCase(unittest.TestCase):
@@ -53,30 +53,30 @@ def test_data_equal__one__last(self):
 class DaskDataframeTest(DaskTestCase):
 
   def test_sanity(self):
-    df = to_dd(self.air_small).compute()
+    df = read_xarray(self.air_small).compute()
     self.assertIsNotNone(df)
     self.assertEqual(len(df), np.prod(list(self.air_small.dims.values())))
 
   def test_columns(self):
-    df = to_dd(self.air_small).compute()
+    df = read_xarray(self.air_small).compute()
     cols = list(df.columns)
     self.assertEqual(cols, ['lat', 'time', 'lon', 'air'])
 
   def test_dtypes(self):
-    df: dd.DataFrame = to_dd(self.air_small).compute()
+    df: dd.DataFrame = read_xarray(self.air_small).compute()
     types = list(df.dtypes)
     self.assertEqual([self.air_small[c].dtype for c in df.columns], types)
 
   def test_partitions_dont_match_dataset_chunks(self):
     standard_blocks = list(block_slices(self.air_small))
-    default: dd.DataFrame = to_dd(self.air_small)
-    chunked: dd.DataFrame = to_dd(self.air_small, dict(time=5))
+    default: dd.DataFrame = read_xarray(self.air_small)
+    chunked: dd.DataFrame = read_xarray(self.air_small, dict(time=5))
 
     self.assertEqual(default.npartitions, len(standard_blocks))
     self.assertNotEqual(chunked.npartitions, len(standard_blocks))
 
   def test_chunk_perf(self):
-    df = to_dd(self.air, chunks=dict(time=6)).compute()
+    df = read_xarray(self.air, chunks=dict(time=6)).compute()
     self.assertIsNotNone(df)
     self.assertEqual(len(df), np.prod(list(self.air.dims.values())))
 

qarray/sql_test.py

@@ -2,13 +2,13 @@
 
 from dask_sql import Context
 
-from .df import to_dd
+from . import read_xarray
 from .df_test import DaskTestCase
 
 
 class SqlTestCase(DaskTestCase):
   def test_sanity(self):
-    df = to_dd(self.air_small)
+    df = read_xarray(self.air_small)
 
     c = Context()
     c.create_table('air', df)
@@ -20,7 +20,7 @@ def test_sanity(self):
     self.assertEqual(len(result), 100)
 
   def test_agg_small(self):
-    df = to_dd(self.air_small)
+    df = read_xarray(self.air_small)
 
     c = Context()
     c.create_table('air', df)
@@ -41,7 +41,7 @@ def test_agg_small(self):
     self.assertEqual(len(result), expected)
 
   def slow_test_agg_regular(self):
-    df = to_dd(self.air)
+    df = read_xarray(self.air)
 
     c = Context()
     c.create_table('air', df)