Repository: ibis-project/ibis

Issue:
ENH: is_computable_arg dispatcher
The list of computable args is fixed right now, but a specialized client may want to define their own notion of a computable arg. A dispatcher would solve this.


Discussion:

Resolving patch:
diff --git a/ibis/pandas/core.py b/ibis/pandas/core.py
--- a/ibis/pandas/core.py
+++ b/ibis/pandas/core.py
@@ -115,14 +115,30 @@
 scalar_types = fixed_width_types + temporal_types
 simple_types = scalar_types + (str, type(None))
 
-_VALID_INPUT_TYPES = (
-    ibis.client.Client,
-    ir.Expr,
-    dt.DataType,
-    type(None),
-    win.Window,
-    tuple,
-) + scalar_types
+
+@functools.singledispatch
+def is_computable_input(arg):
+    """All inputs are not computable without a specific override."""
+    return False
+
+
+@is_computable_input.register(ibis.client.Client)
+@is_computable_input.register(ir.Expr)
+@is_computable_input.register(dt.DataType)
+@is_computable_input.register(type(None))
+@is_computable_input.register(win.Window)
+@is_computable_input.register(tuple)
+def is_computable_input_arg(arg):
+    """Return whether `arg` is a valid computable argument."""
+    return True
+
+
+# Register is_computable_input for each scalar type (int, float, date, etc).
+# We use consume here to avoid leaking the iteration variable into the module.
+ibis.util.consume(
+    is_computable_input.register(t)(is_computable_input_arg)
+    for t in scalar_types
+)
 
 
 def execute_with_scope(expr, scope, aggcontext=None, clients=None, **kwargs):
@@ -222,23 +238,6 @@ def execute_until_in_scope(
     )
 
 
-def is_computable_arg(op, arg):
-    """Is `arg` a valid input to an ``execute_node`` rule?
-
-    Parameters
-    ----------
-    arg : object
-        Any Python object
-
-    Returns
-    -------
-    result : bool
-    """
-    return isinstance(
-        op, (ops.ExpressionList, ops.ValueList, ops.WindowOp)
-    ) or isinstance(arg, _VALID_INPUT_TYPES)
-
-
 def execute_bottom_up(
     expr, scope, aggcontext=None, post_execute_=None, clients=None, **kwargs
 ):
@@ -278,9 +277,7 @@ def execute_bottom_up(
     # figure out what arguments we're able to compute on based on the
     # expressions inputs. things like expressions, None, and scalar types are
     # computable whereas ``list``s are not
-    args = op.inputs
-    is_computable_argument = functools.partial(is_computable_arg, op)
-    computable_args = list(filter(is_computable_argument, args))
+    computable_args = [arg for arg in op.inputs if is_computable_input(arg)]
 
     # recursively compute each node's arguments until we've changed type
     scopes = [
diff --git a/ibis/util.py b/ibis/util.py
--- a/ibis/util.py
+++ b/ibis/util.py
@@ -2,10 +2,12 @@
 
 import collections
 import functools
+import itertools
 import logging
 import operator
 import os
 import types
+from typing import Iterator, Optional, TypeVar
 from uuid import uuid4
 
 import toolz
@@ -234,3 +236,18 @@ def get_logger(name, level=None, format=None, propagate=False):
     )
     logger.addHandler(handler)
     return logger
+
+
+T = TypeVar("T", covariant=True)
+
+
+# taken from the itertools documentation
+def consume(iterator: Iterator[T], n: Optional[int] = None) -> None:
+    """Advance the iterator n-steps ahead. If n is None, consume entirely."""
+    # Use functions that consume iterators at C speed.
+    if n is None:
+        # feed the entire iterator into a zero-length deque
+        collections.deque(iterator, maxlen=0)
+    else:
+        # advance to the empty slice starting at position n
+        next(itertools.islice(iterator, n, n), None)


Test patch:
diff --git a/ibis/pandas/tests/test_core.py b/ibis/pandas/tests/test_core.py
--- a/ibis/pandas/tests/test_core.py
+++ b/ibis/pandas/tests/test_core.py
@@ -1,3 +1,5 @@
+from typing import Any
+
 import pandas as pd
 import pandas.util.testing as tm
 import pytest
@@ -5,8 +7,10 @@
 
 import ibis
 import ibis.common as com
+import ibis.expr.datatypes as dt
 import ibis.expr.operations as ops
 from ibis.pandas.client import PandasClient
+from ibis.pandas.core import is_computable_input
 from ibis.pandas.dispatch import execute_node, post_execute, pre_execute
 
 pytestmark = pytest.mark.pandas
@@ -121,3 +125,47 @@ def tmp_left_join_exe(op, lhs, **kwargs):
     assert result is not None
     assert not result.empty
     assert count[0] == 1
+
+
+def test_is_computable_input():
+    class MyObject:
+        def __init__(self, value: float) -> None:
+            self.value = value
+
+        def __getattr__(self, name: str) -> Any:
+            return getattr(self.value, name)
+
+        def __hash__(self) -> int:
+            return hash((type(self), self.value))
+
+        def __eq__(self, other):
+            return (
+                isinstance(other, type(self))
+                and isinstance(self, type(other))
+                and self.value == other.value
+            )
+
+    @execute_node.register(ops.Add, int, MyObject)
+    def add_int_my_object(op, left, right, **kwargs):
+        return left + right.value
+
+    # This multimethod must be implemented to play nicely with other value
+    # types like columns and literals. In other words, for a custom
+    # non-expression object to play nicely it must somehow map to one of the
+    # types in ibis/expr/datatypes.py
+    @dt.infer.register(MyObject)
+    def infer_my_object(_, **kwargs):
+        return dt.float64
+
+    @is_computable_input.register(MyObject)
+    def is_computable_input_my_object(_):
+        return True
+
+    one = ibis.literal(1)
+    two = MyObject(2.0)
+    assert is_computable_input(two)
+
+    three = one + two
+    four = three + 1
+    result = ibis.pandas.execute(four)
+    assert result == 4.0