Implementation of cansplitpush, repeat_operator, replace_and_remove_nodes and find_lowest_nodes

src/finchlite/algebra/__init__.py

@@ -15,12 +15,14 @@
 )
 from .operator import (
     InitWrite,
+    cansplitpush,
     conjugate,
     first_arg,
     identity,
     overwrite,
     promote_max,
     promote_min,
+    repeat_operator,
 )
 from .tensor import (
     Tensor,
@@ -37,6 +39,7 @@
     "Tensor",
     "TensorFType",
     "TensorPlaceholder",
+    "cansplitpush",
     "conjugate",
     "element_type",
     "fill_value",
@@ -56,6 +59,7 @@
     "promote_type",
     "query_property",
     "register_property",
+    "repeat_operator",
     "return_type",
     "shape_type",
 ]

src/finchlite/algebra/algebra.py

@@ -434,6 +434,9 @@ def is_distributive(op, other_op):
     (np.logical_or, lambda op, other_op: other_op == np.logical_and),
     (operator.pow, lambda op, other_op: False),
     (operator.truediv, lambda op, other_op: False),
+    (operator.add, lambda op, other_op: False),
+    (max, lambda op, other_op: False),
+    (min, lambda op, other_op: False),
 ]:
     register_property(fn, "__call__", "is_distributive", func)
 

src/finchlite/algebra/operator.py

@@ -1,4 +1,8 @@
+import math
+import operator
+
 from . import algebra
+from .algebra import is_associative, is_commutative, is_idempotent
 
 
 def and_test(a, b):
@@ -147,3 +151,73 @@ def identity(x):
     "return_type",
     lambda op, x: x,
 )
+
+
+def repeat_operator(x):
+    """
+    If there exists an operator g such that
+    f(x, x, ..., x)  (n times)  is equal to g(x, n),
+    then return g.
+    """
+    if not callable(x):
+        raise TypeError("Can't check repeat operator of non-callable objects!")
+
+    if is_idempotent(x):
+        return None
+
+    if x is operator.add:
+        return operator.mul
+
+    if x is operator.mul:
+        return math.exp
+
+    return None
+
+
+for fn in [
+    operator.and_,
+    operator.or_,
+    min,
+    max,
+]:
+    algebra.register_property(
+        fn,
+        "__call__",
+        "repeat_operator",
+        lambda op: None,
+    )
+
+algebra.register_property(
+    operator.add,
+    "__call__",
+    "repeat_operator",
+    lambda op: operator.mul,
+)
+
+algebra.register_property(
+    operator.mul,
+    "__call__",
+    "repeat_operator",
+    lambda op: math.exp,
+)
+
+
+def cansplitpush(x, y):
+    """
+    Return True if a reduction with operator `x` can be 'split-pushed' through
+    a pointwise operator `y`.
+
+    We allow split-push when:
+      - x has a known repeat operator (repeat_operator(x) is not None),
+      - x and y are the same operator,
+      - and x is both commutative and associative.
+    """
+    if not callable(x) or not callable(y):
+        raise TypeError("Can't check splitpush of non-callable operators!")
+
+    return (
+        repeat_operator(x) is not None
+        and x == y
+        and is_commutative(x)
+        and is_associative(x)
+    )

src/finchlite/galley/LogicalOptimizer/__init__.py

@@ -1,13 +1,17 @@
 from .annotated_query import (
     AnnotatedQuery,
+    find_lowest_roots,
     get_idx_connected_components,
     get_reducible_idxs,
+    replace_and_remove_nodes,
 )
 from .logic_to_stats import insert_statistics
 
 __all__ = [
     "AnnotatedQuery",
+    "find_lowest_roots",
     "get_idx_connected_components",
     "get_reducible_idxs",
     "insert_statistics",
+    "replace_and_remove_nodes",
 ]