relationalai 0.13.5__py3-none-any.whl → 1.0.0a2__py3-none-any.whl

relationalai/semantics/std/re.py

@@ -1,83 +1,101 @@
 from __future__ import annotations
 
-from relationalai.semantics.internal import internal as i
-from relationalai.semantics.metamodel.util import OrderedSet
-from .std import _Integer, _String, _make_expr
-from typing import Literal, Any
-from .. import std
 
+from . import StringValue, IntegerValue, strings, _function_not_implemented
+from ..frontend.base import Library, Expression, Field, Variable
+from ..frontend.core import Number, String, Integer
+from enum import Enum
 
-def escape(regex: _String) -> i.Expression:
-    return _make_expr("escape_regex_metachars", regex, i.String.ref())
+# the front-end library object
+library = Library("re")
 
-class Match(i.Producer):
+#--------------------------------------------------
+# Relationships
+#--------------------------------------------------
 
-    def __init__(self, regex: _String, string: _String, pos: _Integer = 0, _type: Literal["search", "fullmatch", "match"] = "match"):
-        super().__init__(i.find_model([regex, string, pos]))
-        self.regex = regex
-        self.string = string
-        self.pos = pos
+_regex_match_all = library.Relation("regex_match_all", [Field.input("regex", String), Field.input("value", String), Field("pos", Integer), Field("match", String)])
+# _regex_replace = library.Relation("regex_replace", [Field.input("regex", String), Field.input("replacement", String), Field.input("value", String), Field("result", String)])
 
-        if _type == "match":
-            self._expr = _regex_match_all(self.regex, self.string, std.cast_to_int64(self.pos + 1))
-            self._offset, self._full_match = self._expr._arg_ref(2), self._expr._arg_ref(3)
-        elif _type == "search":
-            raise NotImplementedError("`search` is not implemented")
-        elif _type == "fullmatch":
-            _exp = _regex_match_all(self.regex, self.string, std.cast_to_int64(self.pos + 1))
-            self._offset, self._full_match = _exp._arg_ref(2), _exp._arg_ref(3)
-            self._expr = self._full_match == std.strings.substring(self.string, std.cast_to_int64(self.pos), std.strings.len(self.string))
-
-    def group(self, index: _Integer = 0) -> i.Producer:
-        if index == 0:
-            return self._full_match
-        else:
-            return _make_expr("capture_group_by_index", self.regex, self.string, std.cast_to_int64(self.pos + 1), std.cast_to_int64(index), i.String.ref("res"))
+# _regex_escape = library.Relation("regex_escape", [Field.input("regex", String), Field("result", String)])
+# _capture_group_by_index = library.Relation("capture_group_by_index", [Field.input("regex", String), Field.input("string", String), Field.input("pos", Integer), Field.input("index", Integer), Field("result", String)])
+# _capture_group_by_name = library.Relation("capture_group_by_name", [Field.input("regex", String), Field.input("string", String), Field.input("pos", Integer), Field.input("name", String), Field("result", String)])
 
-    def group_by_name(self, name: _String) -> i.Producer:
-        return _make_expr("capture_group_by_name", self.regex, self.string, std.cast_to_int64(self.pos + 1), name, i.String.ref("res"))
 
-    def start(self) -> i.Expression:
-        return self._offset - 1
+#--------------------------------------------------
+# Operations
+#--------------------------------------------------
 
-    def end(self) -> i.Expression:
-        return std.strings.len(self.group(0)) + self.start() - 1
+#  match is ^{REGEX} , search is {REGEX}, and fullmatch is ^{REGEX}$
 
-    def span(self) -> tuple[i.Producer, i.Producer]:
-        return self.start(), self.end()
+def match(regex: StringValue, value: StringValue) -> RegexMatch:
+    """ Check if the regex matches the value from the start. """
+    # REGEXP_LIKE
+    return RegexMatch(regex, value, type=RegexMatchType.MATCH)
+
+def search(regex: StringValue, value: StringValue, pos: IntegerValue = 0) -> RegexMatch:
+    _function_not_implemented("re.search")
+    return RegexMatch(regex, value, pos, type=RegexMatchType.SEARCH)
+
+def fullmatch(regex: StringValue, value: StringValue, pos: IntegerValue = 0) -> RegexMatch:
+    """ Check if the regex matches the entire value starting at the given position. """
+    return RegexMatch(regex, value, pos, type=RegexMatchType.FULLMATCH)
 
-    def _to_keys(self) -> OrderedSet[Any]:
-        return i.find_keys(self._expr)
+def findall(regex: StringValue, value: StringValue) -> tuple[Variable, Variable]:
+    """ Find all non-overlapping matches of the regex in the value. """
+    _function_not_implemented("re.findall")
+    # exp = _regex_match_all(regex, value)
+    # ix, match = exp._arg_ref(2), exp._arg_ref(3)
+    # rank = i.rank(i.asc(ix, match))
+    # return rank, match
+    raise
 
-    def _compile_lookup(self, compiler:i.Compiler, ctx:i.CompilerContext):
-        compiler.lookup(self.regex, ctx)
-        compiler.lookup(self.string, ctx)
-        compiler.lookup(self.pos, ctx)
-        return compiler.lookup(self._expr, ctx)
+def sub(regex: StringValue, repl: StringValue, value: StringValue):
+    """ Replace occurrences of the regex in the value with the replacement string. """
+    _function_not_implemented("re.sub")
+    # return _regex_replace(regex, repl, value)
 
-    def __getattr__(self, name: str) -> Any:
-        return object.__getattribute__(self, name)
+class RegexMatchType(Enum):
+    MATCH = "match"
+    SEARCH = "search"
+    FULLMATCH = "fullmatch"
 
-    def __setattr__(self, name: str, value: Any) -> None:
-        object.__setattr__(self, name, value)
+class RegexMatch(Expression):
 
+    def __init__(self, regex: StringValue, value: StringValue, pos: IntegerValue = 0, type=RegexMatchType.MATCH):
+        if type == RegexMatchType.FULLMATCH:
+            # fullmatch: ^{REGEX}$
+            self.regex = strings.concat(regex, "$")
+        else:
+            self.regex = regex
+        self.value = value
+        # pos is the 0-based index in value where we start matching
+        self.pos = pos
+        # return value is the matched string
+        self.match = String.ref("match")
+        super().__init__(_regex_match_all, (self.regex, value, self.pos, self.match))
+
+    def start(self) -> IntegerValue:
+        return self.pos
 
-def match(regex: _String, string: _String) -> Match:
-    return Match(regex, string)
+    def end(self) -> IntegerValue:
+        return strings.len(self.match) + self.pos - 1
+
+    def span(self) -> tuple[IntegerValue, IntegerValue]:
+        return self.start(), self.end()
 
-def search(regex: _String, string: _String, pos: _Integer = 0) -> Match:
-    return Match(regex, string, pos, _type="search")
+    def group(self, index: IntegerValue = 0) -> Variable:
+        _function_not_implemented("re.RegexMatch.group")
+        raise
 
-def fullmatch(regex: _String, string: _String, pos: _Integer = 0) -> Match:
-    return Match(regex, string, pos, _type="fullmatch")
+    def group_by_name(self, name: StringValue) -> Variable:
+        _function_not_implemented("re.RegexMatch.group_by_name")
+        raise
 
-def findall(regex: _String, string: _String) -> tuple[i.Producer, i.Producer]:
-    exp = _regex_match_all(regex, string)
-    ix, match = exp._arg_ref(2), exp._arg_ref(3)
-    rank = i.rank(i.asc(ix, match))
-    return rank, match
+#--------------------------------------------------
+# Helpers
+#--------------------------------------------------
 
-def _regex_match_all(regex: _String, string: _String, pos: _Integer|None = None) -> i.Expression:
-    if pos is None:
-        pos = i.Int64.ref()
-    return _make_expr("regex_match_all", regex, string, pos, i.String.ref())
+# def _regex_match_all(regex: StringValue, string: StringValue, pos: IntegerValue|None = None) -> i.Expression:
+#     if pos is None:
+#         pos = i.Int64.ref()
+#     return _make_expr("regex_match_all", regex, string, pos, i.StringValue.ref())

relationalai/semantics/std/strings.py

@@ -1,63 +1,118 @@
 from __future__ import annotations
-from typing import Union, Sequence
-
-from relationalai.semantics.internal import internal as b
-from .std import _String, _Integer, _make_expr
-
-# TODO(coey) can we allow multiple args to this, and convert each to string and concatenate? similar to e.g. julia `string("foo", 1, 'b')`. then we don't need a separate concat function (and concat only works on strings, which is inconvenient).
-def string(s: Union[b.Producer, str, float, int]) -> b.Expression:
-    return _make_expr("string", s, b.String.ref("res"))
-
-def concat(s0: _String, s1: _String, *args: _String) -> b.Expression:
-    res = _make_expr("concat", s0, s1, b.String.ref("res0"))
-    for i, s in enumerate(args):
-        res = _make_expr("concat", res, s, b.String.ref(f"res{i+1}"))
+from decimal import Decimal as PyDecimal
+
+from relationalai.semantics import Float
+
+from . import StringValue, IntegerValue,  _function_not_implemented
+from ..frontend.base import Aggregate, Library, Expression, Field, TupleVariable, Variable
+from ..frontend.core import Date, DateTime, Number, String, Integer, Any
+from typing import Sequence
+from .aggregates import string_join
+
+# the front-end library object
+library = Library("strings")
+
+#--------------------------------------------------
+# Relationships
+#--------------------------------------------------
+
+_string = library.Relation("string", [Field.input("s", Any), Field("result", String)], overloads=[
+    [Number, String],
+    [Float, String],
+    [String, String],
+    [DateTime, String],
+    [Date, String],
+])
+_concat = library.Relation("concat", [Field.input("s1", String), Field.input("s2", String), Field("result", String)])
+_contains = library.Relation("contains", [Field.input("s", String), Field.input("substr", String)])
+_ends_with = library.Relation("ends_with", [Field.input("s", String), Field.input("suffix", String)])
+_len = library.Relation("len", [Field.input("s", String), Field("result", Integer)])
+_levenshtein = library.Relation("levenshtein", [Field.input("s1", String), Field.input("s2", String), Field("result", Integer)])
+_like = library.Relation("like", [Field.input("s", String), Field.input("pattern", String)])
+_lower = library.Relation("lower", [Field.input("s", String), Field("result", String)])
+_replace = library.Relation("replace", [Field.input("source", String), Field.input("old", String), Field.input("new", String), Field("result", String)])
+_starts_with = library.Relation("starts_with", [Field.input("s", String), Field.input("prefix", String)])
+_split = library.Relation("split", [Field.input("s", String), Field.input("separator", String), Field("index", Integer), Field("result", String)])
+_strip = library.Relation("strip", [Field.input("s", String), Field("result", String)])
+_substring = library.Relation("substring", [Field.input("s", String), Field.input("start", Integer), Field.input("stop", Integer), Field("result", String)])
+_upper = library.Relation("upper", [Field.input("s", String), Field("result", String)])
+_regex_match = library.Relation("regex_match", [Field.input("regex", String), Field.input("value", String)])
+_join = library.Relation("join", [Field.input("strs", String, is_list=True), Field.input("sep", String), Field("result", String)])
+
+# split_part = f.relation("split_part", [f.input_field("a", types.String), f.input_field("b", types.String), f.field("c", types.Int64), f.field("d", types.String)])
+
+#--------------------------------------------------
+# Operations
+#--------------------------------------------------
+
+def string(s: StringValue|float|PyDecimal) -> Expression:
+    return _string(s)
+
+def concat(s1: StringValue, s2: StringValue, *sn: StringValue) -> Expression:
+    """ Concatenate multiple strings together. """
+    res = _concat(s1, s2, String.ref("res0"))
+    for i, s in enumerate(sn):
+        res = _concat(res, s, String.ref(f"res{i + 1}"))
     return res
 
-def len(s: _String) -> b.Expression:
-    return _make_expr("num_chars", s, b.Int64.ref("res"))
+def join(strs: Sequence[StringValue], separator: str = "") -> Expression:
+    return _join(TupleVariable(strs), separator)
 
-def startswith(s0: _String, s1: _String) -> b.Expression:
-    return _make_expr("starts_with", s0, s1)
+def contains(s: StringValue, substr: StringValue) -> Expression:
+    """ Check whether `substr` is contained within `s`. """
+    return _contains(s, substr)
 
-def endswith(s0: _String, s1: _String) -> b.Expression:
-    return _make_expr("ends_with", s0, s1)
+def endswith(s: StringValue, suffix: StringValue) -> Expression:
+    """ Check whether `s` ends with `suffix`. """
+    return _ends_with(s, suffix)
 
-def contains(s0: _String, s1: _String) -> b.Expression:
-    return _make_expr("contains", s0, s1)
+def len(s: StringValue) -> Expression:
+    """ Get the length of the string `s`. """
+    return _len(s)
 
-def substring(s: _String, start: _Integer, stop: _Integer) -> b.Expression:
-    # unlike Python, Rel's range is 1..stop inclusive
-    return _make_expr("substring", s, start+1, stop, b.String.ref("res"))
+def levenshtein(s1: StringValue, s2: StringValue) -> Expression:
+    """ Compute the Levenshtein distance between two strings. """
+    return _levenshtein(s1, s2)
 
-def like(s: _String, pattern: _String) -> b.Expression:
-    return _make_expr("like_match", s, pattern)
+def like(s: StringValue, pattern: StringValue) -> Expression:
+    """ Check whether `s` matches the SQL LIKE pattern `pattern`. """
+    return _like(s, pattern)
 
-def lower(s: _String) -> b.Expression:
-    return _make_expr("lower", s, b.String.ref("res"))
+def lower(s: StringValue) -> Expression:
+    """ Convert the string `s` to lowercase. """
+    return _lower(s)
 
-def upper(s: _String) -> b.Expression:
-    return _make_expr("upper", s, b.String.ref("res"))
+def replace(source: StringValue, old: StringValue, new: StringValue) -> Expression:
+    """ Replace occurrences of `old` with `new` in the string `source`. """
+    return _replace(source, old, new)
 
-def strip(s: _String) -> b.Expression:
-    return _make_expr("strip", s, b.String.ref("res"))
+def startswith(s: StringValue, prefix: StringValue) -> Expression:
+    """ Check whether `s` starts with `prefix`. """
+    return _starts_with(s, prefix)
 
-def levenshtein(s: _String, t: _String) -> b.Expression:
-    return _make_expr("levenshtein", s, t, b.Int64.ref("res"))
+def split(s: StringValue, separator: StringValue) -> tuple[Variable, Variable]:
+    """ Split the string `s` by `separator`, returning variables holding index and result. """
+    idx = Integer.ref("index")
+    res = String.ref("result")
+    exp = _split(s, separator, idx, res)
+    return exp[2], exp[3]
 
-def join(strs: Sequence[_String], separator: _String) -> b.Expression:
-    return _make_expr("join", b.TupleArg(strs), separator, b.String.ref("res"))
+def split_part(s: StringValue, separator: StringValue, index: IntegerValue) -> Expression:
+    """ Get the part of the string `s` at `index` after splitting by `separator`. """
+    return _split(s, separator, index)
 
-def replace(s: _String, old: _String, new: _String) -> b.Expression:
-    return _make_expr("replace", s, old, new, b.String.ref("res"))
+def strip(s: StringValue) -> Expression:
+    """ Strip whitespace from both ends of the string `s`. """
+    return _strip(s)
 
-def split(s: _String, separator: _String) -> tuple[b.Producer, b.Producer]:
-    exp = _make_expr("split", separator, s, b.Int64.ref("idx"), b.String.ref("res"))
-    # indexes are 0-based
-    return exp._arg_ref(2) - 1, exp._arg_ref(3)
+def substring(s: StringValue, start: IntegerValue, stop: IntegerValue) -> Expression:
+    """ Get the substring of `s` from `start` to `stop` (0-based, stop exclusive). """
+    return _substring(s, start, stop)
 
-def split_part(s: _String, separator: _String, idx: _Integer) -> b.Expression:
-    return _make_expr("split_part", separator, s, idx + 1, b.String.ref("res"))
+def upper(s: StringValue) -> Expression:
+    """ Convert the string `s` to uppercase. """
+    return _upper(s)
 
-def regex_match(s:_String, regex: _String) -> b.Expression:
-    return _make_expr("regex_match", regex, s)
+def regex_match(value: StringValue, regex: StringValue) -> Expression:
+    """ Check if the `value` matches the given `regex`. """
+    return _regex_match(regex, value)

relationalai/shims/executor.py

@@ -0,0 +1,179 @@
+
+#------------------------------------------------------
+# VERY UNSTABLE API TO EXECUTE FRONTEND PROGRAMS USING V0 EXECUTORS
+#------------------------------------------------------
+from functools import lru_cache
+import json
+from v0.relationalai import debugging
+from v0.relationalai.semantics.lqp.executor import LQPExecutor
+from v0.relationalai.semantics.rel.executor import RelExecutor
+from v0.relationalai.semantics.metamodel import ir as v0, factory as v0_factory
+from v0.relationalai.semantics.metamodel.visitor import collect_by_type
+from v0.relationalai.semantics.snowflake import Table as v0Table
+from v0.relationalai.clients.snowflake import Provider as v0Provider
+from v0.relationalai.clients.config import Config
+
+# from ..config import Config
+from ..config.shims import DRY_RUN
+from ..semantics import Model, Fragment
+from ..semantics.metamodel import metamodel as mm
+from ..semantics.metamodel.typer import Typer
+from ..semantics.metamodel.metamodel_analyzer import Normalize
+from .mm2v0 import Translator
+
+DEBUG=False
+PRINT_RESULT=False
+TYPER_DEBUGGER=False
+
+# DEBUG=True
+# PRINT_RESULT=True
+# TYPER_DEBUGGER=True
+
+@lru_cache()
+def get_config():
+    return Config()
+
+def with_source(item: mm.Node):
+    if not hasattr(item, "source"):
+        raise ValueError(f"Item {item} has no source")
+    elif item.source is None:
+        return {}
+    elif debugging.DEBUG:
+        source = item.source.block
+        if source:
+            return { "file": source.file, "line": source.line, "source": source.source }
+        else:
+            return {"file":item.source.file, "line":item.source.line}
+    else:
+        return {"file":item.source.file, "line":item.source.line}
+
+def execute(query: Fragment, model: Model|None = None, executor=None, export_to="", update=False):
+    if not executor:
+        # use_lqp = Config().reasoner.rule.use_lqp
+        use_lqp = bool(get_config().get("reasoner.rule.use_lqp", True))
+        executor = "lqp" if use_lqp else "rel"
+    mm_model = model.to_metamodel() if model else None
+    mm_query = query.to_metamodel()
+    assert isinstance(mm_query, mm.Node)
+    return execute_mm(mm_query, mm_model, executor, export_to=export_to, update=update, model=model)
+
+def execute_mm(mm_query: mm.Task, mm_model: mm.Model|None = None, executor="lqp", export_to="", update=False, model: Model|None = None):
+    # perform type inference
+    typer = Typer(enforce=False)
+    # normalize the metamodel
+    normalizer = Normalize()
+    # translate the metamodel into a v0 query
+    translator = Translator()
+
+    # for typer debugging
+    debugger_msgs = []
+
+    try:
+        #------------------------------------------------------
+        # Model processing
+        #------------------------------------------------------
+        v0_model = None
+        if mm_model:
+            # type inference
+            debugger_msgs.append(json.dumps({ 'id': 'model', 'content': str(mm_model)}))
+            with debugging.span("compile", metamodel=mm_model) as span:
+                span["compile_type"] = "model_v1"
+                mm_model = typer.infer_model(mm_model)
+                span["typed_mm"] = str(mm_model)
+                debugger_msgs.append(json.dumps({ 'id': 'typed_model', 'content': str(mm_model)}))
+                assert(typer.model_net is not None)
+                debugger_msgs.append(json.dumps({ 'id': 'model_net', 'content': str(typer.model_net.to_mermaid())}))
+                # normalization
+                mm_model = mm_model.mut(root=normalizer.normalize(mm_model.root)) # type: ignore
+                assert isinstance(mm_model, mm.Model)
+                if DEBUG:
+                    print("V1 Model:")
+                    print(mm_model)
+                # translation
+                v0_model = translator.translate_model(mm_model)
+                if DEBUG:
+                    print("Translated v0 Model:")
+                    print(v0_model)
+
+        #------------------------------------------------------
+        # Query processing
+        #------------------------------------------------------
+        # type inference
+        debugger_msgs.append(json.dumps({ 'id': 'query', 'content': str(mm_query)}))
+        with debugging.span("compile", metamodel=mm_query) as span:
+            span["compile_type"] = "query_v1"
+            mm_query = typer.infer_query(mm_query)
+            span["typed_mm"] = str(mm_query)
+            debugger_msgs.append(json.dumps({ 'id': 'typed_query', 'content': str(mm_query)}))
+            assert(typer.last_net is not None)
+            debugger_msgs.append(json.dumps({ 'id': 'query_net', 'content': str(typer.last_net.to_mermaid())}))
+            # normalization
+            mm_query = normalizer.normalize(mm_query)  # type: ignore
+            assert isinstance(mm_query, mm.Task)
+            if DEBUG:
+                print("V1 Query:")
+                print(mm_query)
+            # translation
+            v0_query = translator.translate_query(mm_query)
+            if DEBUG:
+                print("Translated v0 Query:")
+                print(v0_query)
+            assert isinstance(v0_query, v0.Task)
+
+        if v0_model is None:
+            # there was no model, so create one from the elements refered to by the query
+            v0_model = v0_factory.model(
+                collect_by_type(v0.Engine, v0_query),
+                collect_by_type(v0.Relation, v0_query),
+                v0_factory._collect_reachable_types(collect_by_type(v0.Type, v0_query)),
+                v0_factory.logical([])
+            )
+    finally:
+        if TYPER_DEBUGGER:
+            with open("typer_debug.jsonl", "w") as f:
+                for msg in debugger_msgs:
+                    f.write(msg)
+                    f.write('\n')
+
+    if DRY_RUN or get_config().get("compiler.dry_run", False):
+        results = []
+    else:
+        # create snowflake tables for all the tables that have been used
+        ts = [v0Table(t.name) for t in translator.used_tables if not t.uri.startswith("dataframe://")]
+        for t in ts:
+            t._lazy_init()
+            v0Table._used_sources.add(t)
+
+        export_table = None
+        if export_to:
+            export_table = v0Table(export_to)
+
+        # get an executor and execute
+        executor = _get_executor(executor, model.name if model else "")
+        with debugging.span("query", tag=None, export_to=export_to, dsl="", **with_source(mm_query)) as query_span:
+            if isinstance(executor, (LQPExecutor, RelExecutor)):
+                results = executor.execute(v0_model, v0_query, export_to=export_table, update=update)
+            else:
+                results = executor.execute(v0_model, v0_query)
+            query_span["results"] = results
+        if DEBUG or PRINT_RESULT:
+            print(results)
+    return results
+
+@lru_cache()
+def _get_executor(name: str, database: str = "ttb_test"):
+    if name == "duckdb":
+        from v0.relationalai.semantics.sql.executor.duck_db import DuckDBExecutor
+        return DuckDBExecutor()
+    elif name == "lqp":
+        return LQPExecutor(database)
+    elif name == "rel":
+        return RelExecutor(database)
+    elif name == "snowflake":
+        from v0.relationalai.semantics.sql.executor.snowflake import SnowflakeExecutor
+        return SnowflakeExecutor(database="TEST_DB", schema="PUBLIC")
+    else:
+        raise ValueError(f"Unknown executor: {name}")
+
+def get_provider():
+    return v0Provider()

relationalai/shims/helpers.py

@@ -0,0 +1,126 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import  Iterable, Sequence as seq
+
+from relationalai.semantics.metamodel import metamodel as mm
+from relationalai.semantics.metamodel.builtins import builtins as b
+from relationalai.semantics.metamodel.rewriter import NO_WALK, Walker
+from relationalai.util.structures import OrderedSet
+
+
+class ContainerWalker(Walker):
+    """ Walker that redirects all container nodes to a common handler. """
+
+    def enter_container(self, container: mm.Container, children: seq[mm.Node]) -> object |  None:
+        return None
+
+    def exit_container(self, container: mm.Container, children: seq[mm.Node]) -> object |  None:
+        return container
+
+    #------------------------------------------------------
+    # Redirect containers to common handlers
+    #------------------------------------------------------
+
+    def enter_logical(self, logical: mm.Logical):
+        return self.enter_container(logical, logical.body)
+
+    def logical(self, logical: mm.Logical):
+        return self.exit_container(logical, logical.body)
+
+    def enter_sequence(self, sequence: mm.Sequence):
+        return self.enter_container(sequence, sequence.tasks)
+
+    def sequence(self, sequence: mm.Sequence):
+        return self.exit_container(sequence, sequence.tasks)
+
+    def enter_union(self, union: mm.Union):
+        return self.enter_container(union, union.tasks)
+
+    def union(self, union: mm.Union):
+        return self.exit_container(union, union.tasks)
+
+    def enter_match(self, match: mm.Match):
+        return self.enter_container(match, match.tasks)
+
+    def match(self, match: mm.Match):
+        return self.exit_container(match, match.tasks)
+
+    def enter_until(self, until: mm.Until):
+        return self.enter_container(until, [until.check, until.body])
+
+    def until(self, until: mm.Until):
+        self.exit_container(until, [until.check, until.body])
+
+    def enter_wait(self, wait: mm.Wait):
+        return self.enter_container(wait, [wait.check])
+
+    def wait(self, wait: mm.Wait):
+        return self.exit_container(wait, [wait.check])
+
+    def enter_loop(self, loop: mm.Loop):
+        return self.enter_container(loop, [loop.body])
+
+    def loop(self, loop: mm.Loop):
+        return self.exit_container(loop, [loop.body])
+
+    def enter_require(self, require: mm.Require):
+        return self.enter_container(require, [require.domain, require.check])
+
+    def require(self, require: mm.Require):
+        return self.exit_container(require, [require.domain, require.check])
+
+    def enter_not(self, not_: mm.Not):
+        return self.enter_container(not_, [not_.task])
+
+    def not_(self, not_: mm.Not):
+        return self.exit_container(not_, [not_.task])
+
+    def enter_exists(self, exists: mm.Exists):
+        return self.enter_container(exists, [exists.task])
+
+    def exists(self, exists: mm.Exists):
+        return self.exit_container(exists, [exists.task])
+
+
+class VarFinder(ContainerWalker):
+    """ Find all variables used in a set of nodes. """
+
+    def enter_container(self, container, children):
+         if self.shallow:
+            return NO_WALK
+         return None
+
+    def find_vars(self, node:mm.Node|Iterable[mm.Node], positive_only=False, shallow=False) -> OrderedSet[mm.Var]:
+        self.vars:OrderedSet[mm.Var] = OrderedSet()
+        self.positive_only = positive_only
+        self.shallow = shallow
+        self(node)
+        return self.vars
+
+    def enter_not(self, not_:mm.Not):
+        if self.positive_only:
+            return NO_WALK
+
+    def var(self, node:mm.Var):
+        self.vars.add(node)
+
+#------------------------------------------------------
+# Helper functions shared across shim modules
+#------------------------------------------------------
+
+def is_output_update(u: mm.Task) -> bool:
+    return (
+        isinstance(u, mm.Update) and
+        len(u.args) > 0 and
+        isinstance(u.args[0], mm.Var) and
+        isinstance(u.args[0].type, mm.Table) and
+        (u.relation is u.args[0].type or u.relation in u.args[0].type.columns)
+        # In v0 all outputs, including exports, end up as v0.Output so we actually
+        # don't want to limit this to just arrow results
+        # u.args[0].type.uri.startswith("dataframe://")
+    )
+
+def is_main_output(u: mm.Task) -> bool:
+    """ This task is an update to dataframe, and represents the main output (it has only keys, no values)."""
+    return isinstance(u, mm.Update) and is_output_update(u) and isinstance(u.relation, mm.Table)