Datafusion filter

dask_sql/physical/rex/core/call.py

@@ -115,6 +115,25 @@ def __init__(self, operation: Callable, unary_operation: Callable = None):
 
     def reduce(self, *operands, **kwargs):
         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
+                    if isinstance(operand, (np.int64, str)):
+                        operands[idx] = np.timedelta64(operand, "D")
+
             enriched_with_kwargs = lambda kwargs: (
                 lambda x, y: self.operation(x, y, **kwargs)
             )
@@ -169,11 +188,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),
         ),
     ],
 )