Merge remote-tracking branch 'upstream/main' into lint-fixes

Author: hameerabbasi

src/finch/autoschedule/__init__.py

@@ -18,6 +18,7 @@
 from .optimize import (
     lift_subqueries,
     optimize,
+    propagate_fields,
     propagate_map_queries,
 )
 
@@ -37,6 +38,7 @@
     "Subquery",
     "Table",
     "optimize",
+    "propagate_fields",
     "propagate_map_queries",
     "lift_subqueries",
     "PostOrderDFS",

src/finch/autoschedule/compiler.py

@@ -1,5 +1,5 @@
 from textwrap import dedent
-from typing import Any, TypeVar
+from typing import TypeVar
 
 from ..finch_logic import (
     Alias,
@@ -14,15 +14,17 @@
     Reorder,
     Subquery,
     Table,
+    WithFields,
 )
 from ..symbolic import Term
 
 T = TypeVar("T", bound="LogicNode")
 
 
-def get_or_insert(dictionary: dict[str, LogicNode], key: str, default: Any) -> Any:
-    if key in dictionary:
-        return dictionary[key]
+def get_or_insert(dictionary: dict[str, T], key: str, default: T) -> T:
+    val = dictionary.get(key)
+    if val is not None:
+        return val
     dictionary[key] = default
     return default
 
@@ -36,11 +38,14 @@ def get_structure(
         case Alias(name):
             return get_or_insert(aliases, name, Immediate(len(fields) + len(aliases)))
         case Subquery(Alias(name) as lhs, arg):
-            if name in aliases:
-                return aliases[name]
-            return Subquery(
-                get_structure(lhs, fields, aliases), get_structure(arg, fields, aliases)
-            )
+            alias = aliases.get(name)
+            if alias is not None:
+                return alias
+            in_arg = get_structure(arg, fields, aliases)
+            in_lhs = get_structure(lhs, fields, aliases)
+            assert isinstance(in_arg, WithFields)
+            assert isinstance(in_lhs, WithFields)
+            return Subquery(in_lhs, in_arg)
         case Table(tns, idxs):
             assert all(isinstance(idx, Field) for idx in idxs)
             return Table(

src/finch/autoschedule/optimize.py

@@ -1,12 +1,17 @@
+from collections.abc import Iterable
+
 from ..finch_logic import (
     Aggregate,
     Alias,
+    Field,
     LogicNode,
     MapJoin,
     Plan,
     Produces,
     Query,
+    Relabel,
     Subquery,
+    WithFields,
 )
 from ..symbolic import Chain, PostOrderDFS, PostWalk, PreWalk, Rewrite, Term
 from .compiler import LogicCompiler
@@ -89,6 +94,31 @@ def rule_2(ex):
     return Rewrite(PostWalk(rule_2))(root)
 
 
+def _propagate_fields(
+    root: LogicNode, fields: dict[LogicNode, Iterable[Field]]
+) -> LogicNode:
+    match root:
+        case Plan(bodies):
+            return Plan(tuple(_propagate_fields(b, fields) for b in bodies))
+        case Query(lhs, rhs):
+            rhs = _propagate_fields(rhs, fields)
+            assert isinstance(rhs, WithFields)
+            fields[lhs] = rhs.get_fields()
+            return Query(lhs, rhs)
+        case Alias() as a:
+            return Relabel(a, tuple(fields[a]))
+        case node if node.is_expr():
+            return node.make_term(
+                node.head(), *[_propagate_fields(c, fields) for c in node.children()]
+            )
+        case node:
+            return node
+
+
+def propagate_fields(root: LogicNode) -> LogicNode:
+    return _propagate_fields(root, fields={})
+
+
 class DefaultLogicOptimizer:
     def __init__(self, ctx: LogicCompiler):
         self.ctx = ctx

src/finch/finch_logic/__init__.py

@@ -14,6 +14,7 @@
     Reorder,
     Subquery,
     Table,
+    WithFields,
 )
 
 __all__ = [
@@ -34,4 +35,5 @@
     "Reformat",
     "Subquery",
     "Table",
+    "WithFields",
 ]

src/finch/finch_logic/nodes.py

@@ -68,7 +68,14 @@ def make_term(cls, *args: Term) -> Self:
 
 
 @dataclass(eq=True, frozen=True)
-class Immediate(LogicNode):
+class WithFields(LogicNode):
+    @abstractmethod
+    def get_fields(self) -> tuple[Field, ...]:
+        """Get this node's fields."""
+
+
+@dataclass(eq=True, frozen=True)
+class Immediate(WithFields):
     """
     Represents a logical AST expression for the literal value `val`.
 
@@ -97,6 +104,10 @@ def fill_value(self):
     def children(self):
         raise TypeError(f"`{type(self).__name__}` doesn't support `.children()`.")
 
+    def get_fields(self) -> tuple[Field, ...]:
+        """Returns fields of the node."""
+        return ()
+
 
 @dataclass(eq=True, frozen=True)
 class Deferred(LogicNode):
@@ -183,7 +194,7 @@ def children(self):
 
 
 @dataclass(eq=True, frozen=True)
-class Table(LogicNode):
+class Table(WithFields):
     """
     Represents a logical AST expression for a tensor object `tns`, indexed by fields
     `idxs...`. A table is a tensor with named dimensions.
@@ -210,9 +221,17 @@ def children(self):
         """Returns the children of the node."""
         return [self.tns, *self.idxs]
 
+    def get_fields(self) -> tuple[Field, ...]:
+        """Returns fields of the node."""
+        return self.idxs
+
+    @classmethod
+    def make_term(cls, head, tns, *idxs):
+        return head(tns, idxs)
+
 
 @dataclass(eq=True, frozen=True)
-class MapJoin(LogicNode):
+class MapJoin(WithFields):
     """
     Represents a logical AST expression for mapping the function `op` across `args...`.
     Dimensions which are not present are broadcasted. Dimensions which are
@@ -225,7 +244,7 @@ class MapJoin(LogicNode):
     """
 
     op: Immediate
-    args: tuple[LogicNode, ...]
+    args: tuple[WithFields, ...]
 
     @staticmethod
     def is_expr():
@@ -241,8 +260,19 @@ def children(self):
         """Returns the children of the node."""
         return [self.op, *self.args]
 
+    def get_fields(self) -> tuple[Field, ...]:
+        """Returns fields of the node."""
+        # (mtsokol) I'm not sure if this comment still applies - the order is preserved.
+        # TODO: this is wrong here: the overall order should at least be concordant with
+        # the args if the args are concordant
+        fs: list[Field] = []
+        for arg in self.args:
+            fs.extend(arg.get_fields())
+
+        return tuple(fs)
+
     @classmethod
-    def make_term(cls, op: Immediate, *args: LogicNode) -> Self:  # type: ignore[override]
+    def make_term(cls, op: Immediate, *args: WithFields) -> Self:  # type: ignore[override]
         return cls(op, tuple(args))
 
 
@@ -261,7 +291,7 @@ class Aggregate(LogicNode):
 
     op: Immediate
     init: Immediate
-    arg: LogicNode
+    arg: WithFields
     idxs: tuple[Field, ...]
 
     @staticmethod
@@ -278,9 +308,17 @@ def children(self):
         """Returns the children of the node."""
         return [self.op, self.init, self.arg, *self.idxs]
 
+    def get_fields(self) -> tuple[Field, ...]:
+        """Returns fields of the node."""
+        return tuple(field for field in self.arg.get_fields() if field not in self.idxs)
+
+    @classmethod
+    def make_term(cls, head, op, init, arg, *idxs):
+        return head(op, init, arg, idxs)
+
 
 @dataclass(eq=True, frozen=True)
-class Reorder(LogicNode):
+class Reorder(WithFields):
     """
     Represents a logical AST statement that reorders the dimensions of `arg` to be
     `idxs...`. Dimensions known to be length 1 may be dropped. Dimensions that do not
@@ -308,9 +346,17 @@ def children(self):
         """Returns the children of the node."""
         return [self.arg, *self.idxs]
 
+    def get_fields(self) -> tuple[Field, ...]:
+        """Returns fields of the node."""
+        return self.idxs
+
+    @classmethod
+    def make_term(cls, head, arg, *idxs):
+        return head(arg, idxs)
+
 
 @dataclass(eq=True, frozen=True)
-class Relabel(LogicNode):
+class Relabel(WithFields):
     """
     Represents a logical AST statement that relabels the dimensions of `arg` to be
     `idxs...`.
@@ -337,9 +383,13 @@ def children(self):
         """Returns the children of the node."""
         return [self.arg, *self.idxs]
 
+    def get_fields(self) -> tuple[Field, ...]:
+        """Returns fields of the node."""
+        return self.idxs
+
 
 @dataclass(eq=True, frozen=True)
-class Reformat(LogicNode):
+class Reformat(WithFields):
     """
     Represents a logical AST statement that reformats `arg` into the tensor `tns`.
 
@@ -349,7 +399,7 @@ class Reformat(LogicNode):
     """
 
     tns: Immediate
-    arg: LogicNode
+    arg: WithFields
 
     @staticmethod
     def is_expr():
@@ -365,20 +415,24 @@ def children(self):
         """Returns the children of the node."""
         return [self.tns, self.arg]
 
+    def get_fields(self) -> tuple[Field, ...]:
+        """Returns fields of the node."""
+        return self.arg.get_fields()
+
 
 @dataclass(eq=True, frozen=True)
-class Subquery(LogicNode):
+class Subquery(WithFields):
     """
     Represents a logical AST statement that evaluates `rhs`, binding the result to
     `lhs`, and returns `rhs`.
 
     Attributes:
         lhs: The left-hand side of the binding.
-        rhs: The argument to evaluate.
+        arg: The argument to evaluate.
     """
 
     lhs: LogicNode
-    arg: LogicNode
+    arg: WithFields
 
     @staticmethod
     def is_expr():
@@ -394,6 +448,10 @@ def children(self):
         """Returns the children of the node."""
         return [self.lhs, self.arg]
 
+    def get_fields(self) -> tuple[Field, ...]:
+        """Returns fields of the node."""
+        return self.arg.get_fields()
+
 
 @dataclass(eq=True, frozen=True)
 class Query(LogicNode):

src/finch/symbolic/term.py

@@ -40,9 +40,9 @@ def is_expr(self) -> bool:
         `children()` if `True`."""
 
     @abstractmethod
-    def make_term(self, head: Any, *children: Term) -> Self:
+    def make_term(self, *children: Term) -> Self:
         """
-        Construct a new term in the same family of terms with the given head type and
+        Construct a new term in the same family of terms with the given
         children. This function should satisfy
         `x == x.make_term(*x.children())`
         """

tests/test_scheduler.py

@@ -1,13 +1,16 @@
-from finch.autoschedule import lift_subqueries, propagate_map_queries
+from finch.autoschedule import lift_subqueries, propagate_fields, propagate_map_queries
 from finch.finch_logic import (
     Aggregate,
     Alias,
+    Field,
     Immediate,
     MapJoin,
     Plan,
     Produces,
     Query,
+    Relabel,
     Subquery,
+    Table,
 )
 
 
@@ -91,3 +94,40 @@ def test_lift_subqueries():
 
     result = lift_subqueries(plan)
     assert result == expected
+
+
+def test_propagate_fields():
+    plan = Plan(
+        (
+            Query(
+                Alias("A10"),
+                MapJoin(
+                    Immediate("op"),
+                    (
+                        Table(Immediate("tbl1"), (Field("A1"), Field("A2"))),
+                        Table(Immediate("tbl2"), (Field("A2"), Field("A3"))),
+                    ),
+                ),
+            ),
+            Alias("A10"),
+        )
+    )
+
+    expected = Plan(
+        (
+            Query(
+                Alias("A10"),
+                MapJoin(
+                    Immediate("op"),
+                    (
+                        Table(Immediate("tbl1"), (Field("A1"), Field("A2"))),
+                        Table(Immediate("tbl2"), (Field("A2"), Field("A3"))),
+                    ),
+                ),
+            ),
+            Relabel(Alias("A10"), (Field("A1"), Field("A2"), Field("A3"))),
+        )
+    )
+
+    result = propagate_fields(plan)
+    assert result == expected