Turning on Dask-Expr (#28)

This PR enables dask-expr with copy on write semantics in Dask Dataframes. Fixes #14.
alxmrs · Mar 3, 2024 · 35f910f · 35f910f
1 parent 1de22ff
commit 35f910f
Show file tree

Hide file tree

Showing 9 changed files with 3,327 additions and 3 deletions.
diff --git a/perf_tests/compute_air.py-2024-03-03T07:42:36+05:30.svg b/perf_tests/compute_air.py-2024-03-03T07:42:36+05:30.svg
diff --git a/perf_tests/compute_air.py-2024-03-03T07:50:18+05:30.svg b/perf_tests/compute_air.py-2024-03-03T07:50:18+05:30.svg
diff --git a/perf_tests/groupby_air.py-2024-03-03T07:43:22+05:30.svg b/perf_tests/groupby_air.py-2024-03-03T07:43:22+05:30.svg
diff --git a/perf_tests/groupby_air.py-2024-03-03T07:50:06+05:30.svg b/perf_tests/groupby_air.py-2024-03-03T07:50:06+05:30.svg
diff --git a/perf_tests/groupby_air_full.py-2024-03-03T07:44:17+05:30.svg b/perf_tests/groupby_air_full.py-2024-03-03T07:44:17+05:30.svg
diff --git a/perf_tests/groupby_air_full.py-2024-03-03T07:49:53+05:30.svg b/perf_tests/groupby_air_full.py-2024-03-03T07:49:53+05:30.svg
diff --git a/perf_tests/sanity.py-2024-03-03T07:45:47+05:30.svg b/perf_tests/sanity.py-2024-03-03T07:45:47+05:30.svg
diff --git a/perf_tests/sanity.py-2024-03-03T07:49:44+05:30.svg b/perf_tests/sanity.py-2024-03-03T07:49:44+05:30.svg
diff --git a/qarray/df.py b/qarray/df.py
@@ -1,6 +1,7 @@
 import itertools
 import typing as t
 
+import dask
 import dask.dataframe as dd
 import numpy as np
 import pandas as pd
@@ -12,6 +13,12 @@
 Block = t.Dict[str, slice]
 Chunks = t.Dict[str, int]
 
+# Turn on Dask-Expr
+dask.config.set({'dataframe.query-planning-warning': False})
+dask.config.set({"dataframe.query-planning": True})
+# Turn on Copy-On-Write (needs Pandas 2.0).
+pd.options.mode.copy_on_write = True
+
 
 # Borrowed from Xarray
 def _get_chunk_slicer(dim: t.Hashable, chunk_index: t.Mapping,
@@ -123,6 +130,3 @@ def f(b: Block) -> pd.DataFrame:
     divisions=divisions,
     token=token,
   )
-
-# TODO(alxmrs): Try dask expressions dataframe:
-#  https://github.com/dask-contrib/dask-expr