From d50a5e2050348525c9513eea37f007a5d705f493 Mon Sep 17 00:00:00 2001
From: Stefan Krawczyk <stefan@dagworks.io>
Date: Sat, 8 Feb 2025 23:52:59 -0800
Subject: [PATCH 01/13] WIP fix dask changes

---
 examples/dask/hello_world/data_loaders.py | 6 ++----
 hamilton/plugins/h_dask.py                | 3 ++-
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/examples/dask/hello_world/data_loaders.py b/examples/dask/hello_world/data_loaders.py
index 5f0377149..38c4300d1 100644
--- a/examples/dask/hello_world/data_loaders.py
+++ b/examples/dask/hello_world/data_loaders.py
@@ -9,9 +9,7 @@ def spend(spend_location: str, spend_partitions: int) -> dataframe.Series:
     :param spend_partitions: number of partitions to segment the data into
     :return:
     """
-    return dataframe.from_pandas(
-        pd.Series([10, 10, 20, 40, 40, 50]), name="spend", npartitions=spend_partitions
-    )
+    return dataframe.from_pandas(pd.Series([10, 10, 20, 40, 40, 50]), npartitions=spend_partitions)
 
 
 def signups(signups_location: str, signups_partitions: int) -> dataframe.Series:
@@ -22,5 +20,5 @@ def signups(signups_location: str, signups_partitions: int) -> dataframe.Series:
     :return:
     """
     return dataframe.from_pandas(
-        pd.Series([1, 10, 50, 100, 200, 400]), name="signups", npartitions=signups_partitions
+        pd.Series([1, 10, 50, 100, 200, 400]), npartitions=signups_partitions
     )
diff --git a/hamilton/plugins/h_dask.py b/hamilton/plugins/h_dask.py
index 6a15c8b91..7966197ae 100644
--- a/hamilton/plugins/h_dask.py
+++ b/hamilton/plugins/h_dask.py
@@ -3,6 +3,7 @@
 
 import dask.array
 import dask.dataframe
+import dask_expr
 import numpy as np
 import pandas as pd
 from dask import compute
@@ -227,7 +228,7 @@ def get_output_name(output_name: str, column_name: str) -> str:
             elif isinstance(v, (list, tuple)):
                 massaged_outputs[k] = dask.dataframe.from_array(dask.array.from_array(v))
                 columns_expected.append(k)
-            elif isinstance(v, (dask.dataframe.core.Scalar,)):
+            elif isinstance(v, (dask.dataframe.core.Scalar, dask_expr.Scalar)):
                 scalar = v.compute()
                 if length == 0:
                     massaged_outputs[k] = dask.dataframe.from_pandas(

From 4b9e2cdc97022e732d3d89e91413db1af4ccef4b Mon Sep 17 00:00:00 2001
From: Stefan Krawczyk <stefan@dagworks.io>
Date: Mon, 10 Feb 2025 10:31:39 -0800
Subject: [PATCH 02/13] Remove now incorrect config from dask test

---
 plugin_tests/h_dask/conftest.py | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/plugin_tests/h_dask/conftest.py b/plugin_tests/h_dask/conftest.py
index f667cff2a..bc5ef5b5a 100644
--- a/plugin_tests/h_dask/conftest.py
+++ b/plugin_tests/h_dask/conftest.py
@@ -1,9 +1,4 @@
-import dask
-
 from hamilton import telemetry
 
 # disable telemetry for all tests!
 telemetry.disable_telemetry()
-
-# required until we fix the DataFrameResultBuilder to work with dask-expr
-dask.config.set({"dataframe.query-planning": False})

From 50dab62874baef82be7086591ce465ed25622ba0 Mon Sep 17 00:00:00 2001
From: Stefan Krawczyk <stefan@dagworks.io>
Date: Mon, 10 Feb 2025 10:39:06 -0800
Subject: [PATCH 03/13] WIP

---
 hamilton/plugins/h_dask.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/hamilton/plugins/h_dask.py b/hamilton/plugins/h_dask.py
index 7966197ae..dbfbd7567 100644
--- a/hamilton/plugins/h_dask.py
+++ b/hamilton/plugins/h_dask.py
@@ -3,7 +3,6 @@
 
 import dask.array
 import dask.dataframe
-import dask_expr
 import numpy as np
 import pandas as pd
 from dask import compute
@@ -16,6 +15,11 @@
 
 logger = logging.getLogger(__name__)
 
+try:
+    from dask.dataframe.core import Scalar as dask_scalar
+except ImportError:
+    from dask_expr import Scalar as dask_scalar
+
 
 class DaskGraphAdapter(base.HamiltonGraphAdapter):
     """Class representing what's required to make Hamilton run on Dask.
@@ -228,7 +232,7 @@ def get_output_name(output_name: str, column_name: str) -> str:
             elif isinstance(v, (list, tuple)):
                 massaged_outputs[k] = dask.dataframe.from_array(dask.array.from_array(v))
                 columns_expected.append(k)
-            elif isinstance(v, (dask.dataframe.core.Scalar, dask_expr.Scalar)):
+            elif isinstance(v, (dask_scalar,)):
                 scalar = v.compute()
                 if length == 0:
                     massaged_outputs[k] = dask.dataframe.from_pandas(

From a3a787a76df64977035b5e5b2ba58bb007e9f0d7 Mon Sep 17 00:00:00 2001
From: Stefan Krawczyk <stefan@dagworks.io>
Date: Mon, 10 Feb 2025 10:54:15 -0800
Subject: [PATCH 04/13] Making sure installed deps are updated for dask

---
 pyproject.toml | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 477283834..5e83118a5 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -112,8 +112,8 @@ sdk = ["sf-hamilton-sdk"]
 slack = ["slack-sdk"]
 test = [
   "connectorx",
-  "dask",
-  "dask-expr; python_version >= '3.9'",
+  "dask[complete]",
+#  "dask-expr; python_version >= '3.9'",
   "datasets", # huggingface datasets
   "diskcache",
   "dlt",
@@ -129,7 +129,7 @@ test = [
   "mlflow",
   "networkx",
   "openpyxl", # for excel data loader
-  "pandera",
+  "pandera[dask]",
   "plotly",
   "polars",
   "pyarrow",

From 32bd2a562a54c7de41e71df64429bb36c550e519 Mon Sep 17 00:00:00 2001
From: Stefan Krawczyk <stefan@dagworks.io>
Date: Mon, 10 Feb 2025 11:18:24 -0800
Subject: [PATCH 05/13] WIP

---
 hamilton/plugins/h_dask.py | 11 ++++++-----
 pyproject.toml             |  4 ++--
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/hamilton/plugins/h_dask.py b/hamilton/plugins/h_dask.py
index dbfbd7567..0e806688e 100644
--- a/hamilton/plugins/h_dask.py
+++ b/hamilton/plugins/h_dask.py
@@ -1,4 +1,6 @@
 import logging
+
+# import sys
 import typing
 
 import dask.array
@@ -16,9 +18,10 @@
 logger = logging.getLogger(__name__)
 
 try:
-    from dask.dataframe.core import Scalar as dask_scalar
+    from dask.dataframe.dask_expr import Scalar as dask_scalar
 except ImportError:
-    from dask_expr import Scalar as dask_scalar
+    from dask.dataframe.core import Scalar as dask_scalar
+    # maybe import from dask_expr here?
 
 
 class DaskGraphAdapter(base.HamiltonGraphAdapter):
@@ -262,9 +265,7 @@ def get_output_name(output_name: str, column_name: str) -> str:
 
         # assumption is that everything here is a dask series or dataframe
         # we assume that we do column concatenation and that it's an outer join (TBD: make this configurable)
-        _df = dask.dataframe.multi.concat(
-            [o for o in massaged_outputs.values()], axis=1, join="outer"
-        )
+        _df = dask.dataframe.concat([o for o in massaged_outputs.values()], axis=1, join="outer")
         _df.columns = columns_expected
         return _df
 
diff --git a/pyproject.toml b/pyproject.toml
index 5e83118a5..f155dabf3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -51,7 +51,7 @@ docs = [
   "sf-hamilton[dev]",
   "alabaster>=0.7,<0.8,!=0.7.5", # read the docs pins
   "commonmark==0.9.1", # read the docs pins
-  "dask-expr",
+  "dask-expr; python_version == '3.9'",
   "dask[distributed]",
   "ddtrace",
   "diskcache",
@@ -113,7 +113,7 @@ slack = ["slack-sdk"]
 test = [
   "connectorx",
   "dask[complete]",
-#  "dask-expr; python_version >= '3.9'",
+  "dask-expr; python_version == '3.9'",
   "datasets", # huggingface datasets
   "diskcache",
   "dlt",

From 5bdd255c4bedb505f2421b8a458dbc000a866f54 Mon Sep 17 00:00:00 2001
From: Stefan Krawczyk <stefan@dagworks.io>
Date: Mon, 10 Feb 2025 11:23:19 -0800
Subject: [PATCH 06/13] WIP

---
 tests/integrations/pandera/requirements.txt | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/integrations/pandera/requirements.txt b/tests/integrations/pandera/requirements.txt
index 3b1a6b8d2..397d64eb2 100644
--- a/tests/integrations/pandera/requirements.txt
+++ b/tests/integrations/pandera/requirements.txt
@@ -1 +1,2 @@
 # Additional requirements on top of hamilton...pandera
+pandera[dask]

From 970f3626f79511acba8972619495e2fb92b9a634 Mon Sep 17 00:00:00 2001
From: Stefan Krawczyk <stefan@dagworks.io>
Date: Mon, 10 Feb 2025 11:34:44 -0800
Subject: [PATCH 07/13] Add case for dask + py3.9 test

dask expr was rolled into dask, but not for 3.9.
---
 plugin_tests/h_dask/conftest.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/plugin_tests/h_dask/conftest.py b/plugin_tests/h_dask/conftest.py
index bc5ef5b5a..ac6a9df6c 100644
--- a/plugin_tests/h_dask/conftest.py
+++ b/plugin_tests/h_dask/conftest.py
@@ -2,3 +2,11 @@
 
 # disable telemetry for all tests!
 telemetry.disable_telemetry()
+
+# dask_expr got made default, except for python 3.9 and below
+import sys
+
+if sys.version_info < (3, 10):
+    import dask
+
+    dask.config.set({"dataframe.query-planning": False})

From b4487c3bf7fa9a26960e7ff8f4344e7f57a506f2 Mon Sep 17 00:00:00 2001
From: Stefan Krawczyk <stefan@dagworks.io>
Date: Mon, 10 Feb 2025 11:40:12 -0800
Subject: [PATCH 08/13] Fix deps for pandera dask

---
 .ci/test.sh                                 | 10 +++++-----
 tests/integrations/pandera/requirements.txt |  1 +
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/.ci/test.sh b/.ci/test.sh
index 3a349bbf0..015a5a40e 100755
--- a/.ci/test.sh
+++ b/.ci/test.sh
@@ -27,12 +27,12 @@ fi
 
 if [[ ${TASK} == "integrations" ]]; then
     pip install -e '.[pandera]'
-    pip install dask
-    if python -c 'import sys; exit(0) if sys.version_info > (3, 9) else exit(1)'; then
-        echo "python version is 3.9+"
-        pip install dask-expr
+    pip install -r tests/integrations/requirements.txt
+    if python -c 'import sys; exit(0) if sys.version_info[:2] == (3, 9) else exit(1)'; then
+      echo "Python version is 3.9"
+      pip install dask-expr
     else
-        echo "Python version is 3.8 or less"
+      echo "Python version is not 3.9"
     fi
     pytest tests/integrations
     exit 0
diff --git a/tests/integrations/pandera/requirements.txt b/tests/integrations/pandera/requirements.txt
index 397d64eb2..41f8d8968 100644
--- a/tests/integrations/pandera/requirements.txt
+++ b/tests/integrations/pandera/requirements.txt
@@ -1,2 +1,3 @@
 # Additional requirements on top of hamilton...pandera
+dask
 pandera[dask]

From 455914c7abd512096b3d5492dfcabf7cb691e901 Mon Sep 17 00:00:00 2001
From: Stefan Krawczyk <stefan@dagworks.io>
Date: Mon, 10 Feb 2025 11:51:54 -0800
Subject: [PATCH 09/13] WIP

---
 .ci/test.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.ci/test.sh b/.ci/test.sh
index 015a5a40e..2966d2402 100755
--- a/.ci/test.sh
+++ b/.ci/test.sh
@@ -27,7 +27,7 @@ fi
 
 if [[ ${TASK} == "integrations" ]]; then
     pip install -e '.[pandera]'
-    pip install -r tests/integrations/requirements.txt
+    pip install -r tests/integrations/pandera/requirements.txt
     if python -c 'import sys; exit(0) if sys.version_info[:2] == (3, 9) else exit(1)'; then
       echo "Python version is 3.9"
       pip install dask-expr

From cf552c95b9f4c8d66fe49f16bcbbae5a46f50ec0 Mon Sep 17 00:00:00 2001
From: Stefan Krawczyk <stefan@dagworks.io>
Date: Mon, 10 Feb 2025 12:10:55 -0800
Subject: [PATCH 10/13] WIP

---
 .ci/test.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.ci/test.sh b/.ci/test.sh
index 2966d2402..1d2542ec7 100755
--- a/.ci/test.sh
+++ b/.ci/test.sh
@@ -26,7 +26,7 @@ if [[ ${TASK} == "dask" ]]; then
 fi
 
 if [[ ${TASK} == "integrations" ]]; then
-    pip install -e '.[pandera]'
+    pip install -e '.[pandera, test]'
     pip install -r tests/integrations/pandera/requirements.txt
     if python -c 'import sys; exit(0) if sys.version_info[:2] == (3, 9) else exit(1)'; then
       echo "Python version is 3.9"

From 3edb1b00f629572ae5f5f4906cc278584bbd3567 Mon Sep 17 00:00:00 2001
From: Stefan Krawczyk <stefan@dagworks.io>
Date: Mon, 10 Feb 2025 12:55:46 -0800
Subject: [PATCH 11/13] Pandera dask series mark as xfail

Some weird interplay with dependencies. Punting on this for now.
---
 tests/integrations/pandera/requirements.txt             | 1 -
 tests/integrations/pandera/test_pandera_data_quality.py | 3 +++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/integrations/pandera/requirements.txt b/tests/integrations/pandera/requirements.txt
index 41f8d8968..397d64eb2 100644
--- a/tests/integrations/pandera/requirements.txt
+++ b/tests/integrations/pandera/requirements.txt
@@ -1,3 +1,2 @@
 # Additional requirements on top of hamilton...pandera
-dask
 pandera[dask]
diff --git a/tests/integrations/pandera/test_pandera_data_quality.py b/tests/integrations/pandera/test_pandera_data_quality.py
index 7f3231870..fd1be3221 100644
--- a/tests/integrations/pandera/test_pandera_data_quality.py
+++ b/tests/integrations/pandera/test_pandera_data_quality.py
@@ -200,6 +200,9 @@ def foo(fail: bool = False) -> dd.DataFrame:
 
 
 @pytest.mark.skipif(sys.version_info < (3, 9), reason="requires python3.9 or higher")
+@pytest.xfail(
+    reason="some weird import issue leads to key error in pandera, can't recreate outside of the series decorator"
+)
 def test_pandera_decorator_dask_series():
     """Validates that the function can be annotated with a dask series type it'll work appropriately.
     Install dask if this fails.

From 27d0aa54b648f77748075f4ad72b72861cb4e6a7 Mon Sep 17 00:00:00 2001
From: Stefan Krawczyk <stefan@dagworks.io>
Date: Mon, 10 Feb 2025 13:02:17 -0800
Subject: [PATCH 12/13] WIP

---
 tests/integrations/pandera/test_pandera_data_quality.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/integrations/pandera/test_pandera_data_quality.py b/tests/integrations/pandera/test_pandera_data_quality.py
index fd1be3221..0defdfd91 100644
--- a/tests/integrations/pandera/test_pandera_data_quality.py
+++ b/tests/integrations/pandera/test_pandera_data_quality.py
@@ -200,7 +200,7 @@ def foo(fail: bool = False) -> dd.DataFrame:
 
 
 @pytest.mark.skipif(sys.version_info < (3, 9), reason="requires python3.9 or higher")
-@pytest.xfail(
+@pytest.mark.xfail(
     reason="some weird import issue leads to key error in pandera, can't recreate outside of the series decorator"
 )
 def test_pandera_decorator_dask_series():

From 1b02b7677afe10769c916aa4c79fa36133578712 Mon Sep 17 00:00:00 2001
From: Stefan Krawczyk <stefan@dagworks.io>
Date: Mon, 10 Feb 2025 13:51:54 -0800
Subject: [PATCH 13/13] WIP

---
 hamilton/plugins/h_dask.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/hamilton/plugins/h_dask.py b/hamilton/plugins/h_dask.py
index 0e806688e..2bee07ddc 100644
--- a/hamilton/plugins/h_dask.py
+++ b/hamilton/plugins/h_dask.py
@@ -1,6 +1,4 @@
 import logging
-
-# import sys
 import typing
 
 import dask.array
@@ -20,8 +18,8 @@
 try:
     from dask.dataframe.dask_expr import Scalar as dask_scalar
 except ImportError:
+    # this is for older versions of dask
     from dask.dataframe.core import Scalar as dask_scalar
-    # maybe import from dask_expr here?
 
 
 class DaskGraphAdapter(base.HamiltonGraphAdapter):