From 819ab0d056927573dcd0c2dd38c6c728aa8dd88f Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sun, 3 Mar 2024 09:02:48 +0530 Subject: [PATCH 1/4] `unravel_to_pd`: a four line implementation. --- qarray/core.py | 13 +++++++++++++ qarray/df.py | 3 +-- 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/qarray/core.py b/qarray/core.py index 7915440a..47c12f54 100644 --- a/qarray/core.py +++ b/qarray/core.py @@ -2,6 +2,7 @@ import typing as t import numpy as np +import pandas as pd import xarray as xr Row = t.List[t.Any] @@ -44,3 +45,15 @@ def unbounded_unravel(ds: xr.Dataset) -> np.ndarray: out[d] = coords[:, i] return out + + +def unravel_to_pd(ds: xr.Dataset) -> pd.DataFrame: + data = {name: da.values.ravel() for name, da in ds.items()} + + coord_vals = (ds.coords[k].values for k in ds.dims.keys()) + coords = pd.MultiIndex.from_product(coord_vals, names=ds.dims.keys()) + + # TODO(alxmrs): The reset index takes a lot of time! Maybe it's still OK. + return pd.DataFrame(data, index=coords).reset_index() + + diff --git a/qarray/df.py b/qarray/df.py index dec4392d..7efabc7d 100644 --- a/qarray/df.py +++ b/qarray/df.py @@ -79,8 +79,7 @@ def to_pd(ds: xr.Dataset, bounded=True) -> pd.DataFrame: df[c] = df[c].astype(ds[c].dtype) return df else: - data = core.unbounded_unravel(ds) - return pd.DataFrame.from_records(data) + return core.unravel_to_pd(ds) def _block_len(block: Block) -> int: From f580386273e2d20f9f37d08985b3a76519012d0f Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sun, 3 Mar 2024 14:12:51 +0530 Subject: [PATCH 2/4] Just using Xarray Dataset's built-in to_dataframe() method. They are the same implementation. --- qarray/core.py | 12 ------------ qarray/df.py | 17 +++-------------- 2 files changed, 3 insertions(+), 26 deletions(-) diff --git a/qarray/core.py b/qarray/core.py index 47c12f54..e1a3aae7 100644 --- a/qarray/core.py +++ b/qarray/core.py @@ -45,15 +45,3 @@ def unbounded_unravel(ds: xr.Dataset) -> np.ndarray: out[d] = coords[:, i] return out - - -def unravel_to_pd(ds: xr.Dataset) -> pd.DataFrame: - data = {name: da.values.ravel() for name, da in ds.items()} - - coord_vals = (ds.coords[k].values for k in ds.dims.keys()) - coords = pd.MultiIndex.from_product(coord_vals, names=ds.dims.keys()) - - # TODO(alxmrs): The reset index takes a lot of time! Maybe it's still OK. - return pd.DataFrame(data, index=coords).reset_index() - - diff --git a/qarray/df.py b/qarray/df.py index 7efabc7d..49b924ce 100644 --- a/qarray/df.py +++ b/qarray/df.py @@ -71,17 +71,6 @@ def explode( 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: - return core.unravel_to_pd(ds) - - def _block_len(block: Block) -> int: return np.prod([v.stop - v.start for v in block.values()]) @@ -103,8 +92,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. @@ -123,7 +112,7 @@ def f(b: Block) -> pd.DataFrame: } return from_map( - f, + pivot, blocks, meta=meta, divisions=divisions, From 1118b36c6fdfea49513cd062c69f9ad6462354e1 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sun, 3 Mar 2024 14:13:43 +0530 Subject: [PATCH 3/4] Remove import of to_pd. --- qarray/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/qarray/__init__.py b/qarray/__init__.py index acc1e608..875c829c 100644 --- a/qarray/__init__.py +++ b/qarray/__init__.py @@ -1,3 +1,3 @@ from .core import unravel -from .df import to_pd, to_dd +from .df import to_dd From aec8c9008456d0683b4369497f6032dc9874e709 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sun, 3 Mar 2024 14:20:33 +0530 Subject: [PATCH 4/4] Refactor: using xarray's built-in dataframe conversion. This also renames `to_dd` to `read_xarray`. --- perf_tests/compute_air.py | 2 +- perf_tests/groupby_air.py | 2 +- perf_tests/groupby_air_full.py | 2 +- perf_tests/open_era5.py | 2 +- perf_tests/sanity.py | 2 +- qarray/__init__.py | 3 +-- qarray/core.py | 3 ++- qarray/df.py | 16 +++++----------- qarray/df_integration_test.py | 4 ++-- qarray/df_test.py | 14 +++++++------- qarray/sql_test.py | 8 ++++---- 11 files changed, 26 insertions(+), 32 deletions(-) diff --git a/perf_tests/compute_air.py b/perf_tests/compute_air.py index 368d75b7..3a93fe64 100755 --- a/perf_tests/compute_air.py +++ b/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)) \ No newline at end of file diff --git a/perf_tests/groupby_air.py b/perf_tests/groupby_air.py index 694007c1..ab7e69f3 100755 --- a/perf_tests/groupby_air.py +++ b/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) diff --git a/perf_tests/groupby_air_full.py b/perf_tests/groupby_air_full.py index f4523407..28a988b2 100755 --- a/perf_tests/groupby_air_full.py +++ b/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) diff --git a/perf_tests/open_era5.py b/perf_tests/open_era5.py index 69aaf660..7511faf5 100755 --- a/perf_tests/open_era5.py +++ b/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) \ No newline at end of file diff --git a/perf_tests/sanity.py b/perf_tests/sanity.py index 6a5dc6e9..b58242a8 100755 --- a/perf_tests/sanity.py +++ b/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)) \ No newline at end of file diff --git a/qarray/__init__.py b/qarray/__init__.py index 875c829c..823ce77a 100644 --- a/qarray/__init__.py +++ b/qarray/__init__.py @@ -1,3 +1,2 @@ -from .core import unravel -from .df import to_dd +from .df import read_xarray diff --git a/qarray/core.py b/qarray/core.py index e1a3aae7..a55cd83b 100644 --- a/qarray/core.py +++ b/qarray/core.py @@ -2,7 +2,6 @@ import typing as t import numpy as np -import pandas as pd import xarray as xr Row = t.List[t.Any] @@ -12,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()) @@ -23,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()) diff --git a/qarray/df.py b/qarray/df.py index 49b924ce..bd518618 100644 --- a/qarray/df.py +++ b/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,10 +60,7 @@ 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)) @@ -75,8 +69,8 @@ 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. diff --git a/qarray/df_integration_test.py b/qarray/df_integration_test.py index 857e6b91..61ac0943 100644 --- a/qarray/df_integration_test.py +++ b/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', diff --git a/qarray/df_test.py b/qarray/df_test.py index 1c789afb..e4ea7acd 100644 --- a/qarray/df_test.py +++ b/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()))) diff --git a/qarray/sql_test.py b/qarray/sql_test.py index d121d39a..29a5cc4f 100644 --- a/qarray/sql_test.py +++ b/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)