Datafusion filter

dask_planner/src/sql/logical.rs

@@ -212,7 +212,7 @@ impl PyLogicalPlan {
             .iter()
             .map(|f| RelDataTypeField::from(f, schema.as_ref()))
             .collect::<Result<Vec<_>>>()
-            .map_err(py_type_err)?;
+            .map_err(|e| py_type_err(e))?;
         Ok(RelDataType::new(false, rel_fields))
     }
 }

dask_sql/physical/rex/core/call.py

@@ -2,7 +2,7 @@
 import logging
 import operator
 import re
-from functools import reduce
+from functools import partial, reduce
 from typing import TYPE_CHECKING, Any, Callable, Union
 
 import dask.array as da
@@ -36,6 +36,17 @@
 SeriesOrScalar = Union[dd.Series, Any]
 
 
+def as_timelike(op):
+    if isinstance(op, np.int64):
+        return np.timedelta64(op, "D")
+    elif isinstance(op, str):
+        return np.datetime64(op)
+    elif pd.api.types.is_datetime64_dtype(op):
+        return op
+    else:
+        raise ValueError(f"Don't know how to make {type(op)} timelike")
+
+
 class Operation:
     """Helper wrapper around a function, which is used as operator"""
 
@@ -115,10 +126,39 @@ def __init__(self, operation: Callable, unary_operation: Callable = None):
 
     def reduce(self, *operands, **kwargs):
         if len(operands) > 1:
-            enriched_with_kwargs = lambda kwargs: (
-                lambda x, y: self.operation(x, y, **kwargs)
-            )
-            return reduce(enriched_with_kwargs(kwargs), operands)
+            if any(map(pd.api.types.is_datetime64_dtype, operands)):
+                operands = tuple(map(as_timelike, operands))
+            return reduce(partial(self.operation, **kwargs), operands)
+        # if len(operands) > 1:
+        #     # Doing math against dates requires a Timedelta
+        #     # Find all the operands that are of a datetime64 type
+        #     date_operands = [
+        #         idx for idx in {0, 1} if pd.api.types.is_datetime64_dtype(operands[idx])
+        #     ]
+
+        #     # If there are datetime64 operands we need to make sure that the other operands
+        #     # in the list are Timedelta for the operation to work.
+        #     if date_operands:
+        #         # Operands is a Set, since we are altering it must convert to a List
+        #         operands = list(operands)
+
+        #         # Knowing there are datetime types in the operands check for incompatable other types
+        #         # If found, convert them to Timedelta
+        #         for idx, operand in enumerate(operands):
+        #             # Default to `Day`/`D` since that is what PostgreSQL does
+        #             # TODO: Would be nice to have the scalar conversion pushed further
+        #             # up in the processing chain here for clauses like `DATE '2022-06-16'`
+        #             # for example so at the point the operand reaches here it is already
+        #             # a date type
+        #             if isinstance(operand, np.int64):
+        #                 operands[idx] = np.timedelta64(operand, "D")
+        #             elif isinstance(operand, str):
+        #                 operands[idx] = np.datetime64(operand)
+
+        #     enriched_with_kwargs = lambda kwargs: (
+        #         lambda x, y: self.operation(x, y, **kwargs)
+        #     )
+        #     return reduce(enriched_with_kwargs(kwargs), operands)
         else:
             return self.unary_operation(*operands, **kwargs)
 
@@ -169,11 +209,10 @@ def div(self, lhs, rhs):
         # We do not need to truncate in this case
         # So far, I did not spot any other occurrence
         # of this function.
-        if isinstance(result, datetime.timedelta):
+        if isinstance(result, (datetime.timedelta, np.timedelta64)):
             return result
         else:  # pragma: no cover
-            result = da.trunc(result)
-            return result
+            return da.trunc(result).astype(np.int64)
 
 
 class CaseOperation(Operation):

tests/integration/test_filter.py

@@ -74,7 +74,7 @@ def test_string_filter(c, string_table):
         "datetime_table",
         pytest.param(
             "gpu_datetime_table",
-            marks=(pytest.mark.gpu, pytest.mark.skip(reason="WIP DataFusion")),
+            marks=(pytest.mark.gpu),
         ),
     ],
 )
@@ -100,7 +100,7 @@ def test_filter_cast_date(c, input_table, request):
         "datetime_table",
         pytest.param(
             "gpu_datetime_table",
-            marks=(pytest.mark.gpu, pytest.mark.skip(reason="WIP DataFusion")),
+            marks=(pytest.mark.gpu),
         ),
     ],
 )