unit test for reproducing bug

Eventual-Inc · Oct 7, 2022 · 4ef5fd1 · 4ef5fd1
1 parent 482dc16
commit 4ef5fd1
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/tests/dataframe_cookbook/test_explodes.py b/tests/dataframe_cookbook/test_explodes.py
@@ -34,9 +34,14 @@ def add_one(o: MyObj):
 
     assert df.schema()["explode"].daft_type == ExpressionType.python_object()
 
+    df = df.select(col("explode_plus1"))
+
+    assert df.schema()["explode_plus1"].daft_type == ExpressionType.python_object()
+
     pd_df = pd.DataFrame(data)
     pd_df = pd_df.explode("explode")
     pd_df["explode_plus1"] = pd_df["explode"].apply(add_one)
+    pd_df = pd_df[["explode_plus1"]]
 
     df.collect()
     daft_pd_df = pd.DataFrame(df._result.to_pydict())