relationalai 1.0.0a3__py3-none-any.whl → 1.0.0a5__py3-none-any.whl

relationalai/shims/executor.py

@@ -3,6 +3,7 @@
 # VERY UNSTABLE API TO EXECUTE FRONTEND PROGRAMS USING V0 EXECUTORS
 #------------------------------------------------------
 from functools import lru_cache
+from pandas import DataFrame
 import json
 from v0.relationalai import debugging
 from v0.relationalai.semantics.lqp.executor import LQPExecutor
@@ -10,6 +11,7 @@ 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.semantics.metamodel import executor as v0executor
 try:
     from v0.relationalai.clients.snowflake import Provider as v0Provider #type: ignore
 except ImportError:
@@ -17,18 +19,17 @@ except ImportError:
 from v0.relationalai.clients.config import Config
 
 # from ..config import Config
-from ..config.shims import DRY_RUN
+from ..config.shims import DRY_RUN, ENFORCE_TYPE_CORRECT
 from ..semantics import Model, Fragment
+from ..semantics.frontend.pprint import format
 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
 
@@ -50,19 +51,21 @@ def with_source(item: mm.Node):
     else:
         return {"file":item.source.file, "line":item.source.line}
 
-def execute(query: Fragment, model: Model|None = None, executor=None, export_to="", update=False):
+def execute(query: Fragment, model: Model|None = None, executor=None, export_to="", update=False, meta=None):
     if not executor:
-        # use_lqp = Config().reasoner.rule.use_lqp
-        use_lqp = bool(get_config().get("reasoner.rule.use_lqp", True))
+        use_lqp = 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)
+    with debugging.span("query", tag=None, export_to=export_to, dsl=format(query), **with_source(mm_query)) as query_span:
+        results = execute_mm(mm_query, mm_model, executor, export_to=export_to, update=update, model=model, meta=meta)
+        query_span["results"] = results
+        return results
 
-def execute_mm(mm_query: mm.Task, mm_model: mm.Model|None = None, executor="lqp", export_to="", update=False, model: Model|None = None):
+def execute_mm(mm_query: mm.Task, mm_model: mm.Model|None = None, executor="lqp", export_to="", update=False, model: Model|None = None, meta=None):
     # perform type inference
-    typer = Typer(enforce=False)
+    typer = Typer(enforce=ENFORCE_TYPE_CORRECT)
     # normalize the metamodel
     normalizer = Normalize()
     # translate the metamodel into a v0 query
@@ -81,22 +84,27 @@ def execute_mm(mm_query: mm.Task, mm_model: mm.Model|None = None, executor="lqp"
             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)
+                with debugging.span("v1") as span:
+                    with debugging.span("Original") as span:
+                        if debugging.DEBUG:
+                            span["metamodel"] = str(mm_model.root)
+                    with debugging.span("Typer") as span:
+                        mm_model = typer.infer_model(mm_model)
+                        if debugging.DEBUG:
+                            span["metamodel"] = str(mm_model.root)
+                            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())}))
+                    with debugging.span("Normalizer") as span:
+                        mm_model = mm_model.mut(root=normalizer.normalize(mm_model.root)) # type: ignore
+                        assert isinstance(mm_model, mm.Model)
+                        if debugging.DEBUG:
+                            span["metamodel"] = str(mm_model.root)
+                    with debugging.span("v0 Translation") as span:
+                        v0_model = translator.translate_model(mm_model)
+                        if debugging.DEBUG:
+                            span["metamodel"] = str(v0_model.root)
 
         #------------------------------------------------------
         # Query processing
@@ -105,23 +113,23 @@ def execute_mm(mm_query: mm.Task, mm_model: mm.Model|None = None, executor="lqp"
         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)
+            with debugging.span("v1") as span:
+                with debugging.span("Original") as span:
+                    span["metamodel"] = str(mm_query)
+                with debugging.span("Typer") as span:
+                    mm_query = typer.infer_query(mm_query)
+                    span["metamodel"] = str(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())}))
+                with debugging.span("Normalizer") as span:
+                    mm_query = normalizer.normalize(mm_query)  # type: ignore
+                    assert isinstance(mm_query, mm.Task)
+                with debugging.span("v0 Translation") as span:
+                    v0_query = translator.translate_query(mm_query)
+                    span["metamodel"] = str(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
@@ -139,13 +147,14 @@ def execute_mm(mm_query: mm.Task, mm_model: mm.Model|None = None, executor="lqp"
                     f.write('\n')
 
     if DRY_RUN or get_config().get("compiler.dry_run", False):
-        results = []
+        results = DataFrame()
     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)
+            if not any(v0t._fqn == t._fqn for v0t in v0Table._used_sources):
+                t._lazy_init()
+                v0Table._used_sources.add(t)
 
         export_table = None
         if export_to:
@@ -153,18 +162,17 @@ def execute_mm(mm_query: mm.Task, mm_model: mm.Model|None = None, executor="lqp"
 
         # 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:
+        if isinstance(executor, (LQPExecutor, RelExecutor)):
+            results = executor.execute(v0_model, v0_query, export_to=export_table, update=update, meta=meta)
+        else:
+            results = executor.execute(v0_model, v0_query)
+        if PRINT_RESULT:
             print(results)
     return results
 
 @lru_cache()
 def _get_executor(name: str, database: str = "ttb_test"):
+    v0executor.SUPPRESS_TYPE_ERRORS = True
     if name == "duckdb":
         from v0.relationalai.semantics.sql.executor.duck_db import DuckDBExecutor
         return DuckDBExecutor()

relationalai/shims/fixtures.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+import os
+
+import v0.relationalai as rai
+from v0.relationalai.clients import config as cfg
+
+def create_engine(engine_name: str, size: str,  use_direct_access=False):
+    print('create_engine: about to call make_config')
+    config = make_config(engine_name, use_direct_access=use_direct_access)
+
+    provider = rai.Resources(config=config)
+    print(f"Creating engine {engine_name}")
+    provider.create_engine(name=engine_name, type="LOGIC", size=size)
+    print(f"Engine {engine_name} created")
+
+def delete_engine(engine_name: str, use_direct_access=False):
+    print(f"Deleting engine {engine_name}")
+    config = make_config(engine_name, use_direct_access=use_direct_access)
+    provider = rai.Resources(config=config)
+    provider.delete_engine(engine_name, "LOGIC")
+    print(f"Engine {engine_name} deleted")
+
+def make_config(engine_name: str | None = None, fetch_profile: bool = True, use_direct_access = False, show_full_traces = True, show_debug_logs = True, reuse_model = False) -> cfg.Config:
+    # First try to load from raiconfig.toml
+    try:
+        config = cfg.Config()
+        if config.file_path is not None:
+            # Set test defaults
+            config.set("reuse_model", reuse_model)
+            rai.Resources(config=config)
+            return config
+    except Exception as e:
+        print(f"Could not load config from file: {e}, trying environment variables")
+    # If that fails, construct from environment variables
+    cloud_provider = os.getenv("RAI_CLOUD_PROVIDER")
+
+    if cloud_provider is None:
+        cloud_provider = os.getenv("RAI_CLOUD_PROVIDER")
+
+    print('cloud provider:', cloud_provider)
+
+    if cloud_provider is None:
+        raise ValueError("RAI_CLOUD_PROVIDER must be set")
+    elif cloud_provider == "snowflake":
+        sf_username = os.getenv("SF_TEST_ACCOUNT_USERNAME")
+        sf_password = os.getenv("SF_TEST_ACCOUNT_PASSWORD")
+        sf_account = os.getenv("SF_TEST_ACCOUNT_NAME")
+        sf_role = os.getenv("SF_TEST_ROLE_NAME", "RAI_USER")
+        sf_warehouse = os.getenv("SF_TEST_WAREHOUSE_NAME")
+        sf_app_name = os.getenv("SF_TEST_APP_NAME")
+        if sf_username is None or sf_password is None:
+            raise ValueError(
+                "SF_TEST_ACCOUNT_USERNAME, SF_TEST_ACCOUNT_PASSWORD, SF_TEST_ACCOUNT_NAME must be set if RAI_CLOUD_PROVIDER is set to 'snowflake'"
+            )
+
+        current_config = {
+            "platform": "snowflake",
+            "user": sf_username,
+            "password": sf_password,
+            "account": sf_account,
+            "role": sf_role,
+            "warehouse": sf_warehouse,
+            "rai_app_name": sf_app_name,
+            "use_direct_access": use_direct_access,
+            "show_full_traces": show_full_traces,
+            "show_debug_logs": show_debug_logs,
+            "reuse_model": reuse_model,
+        }
+        if engine_name:
+            current_config["engine"] = engine_name
+
+        # For direct access, we use key-pair as the default for testing and need to configure additional parameters.
+        if use_direct_access:
+            authenticator=os.getenv("AUTHENTICATOR","")
+            private_key_file=os.getenv("PRIVATE_KEY_FILE","")
+            private_key_file_pwd=os.getenv("SF_TEST_ACCOUNT_KEY_PASSPHRASE","")
+
+            current_config["authenticator"] = authenticator
+            current_config["private_key_file"] = private_key_file
+            current_config["private_key_file_pwd"]=private_key_file_pwd
+
+        return cfg.Config(current_config, fetch=fetch_profile)
+
+    else:
+        raise ValueError(f"Unsupported cloud provider: {cloud_provider}")

relationalai/shims/helpers.py

@@ -1,10 +1,8 @@
 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
 
@@ -105,6 +103,32 @@ class VarFinder(ContainerWalker):
     def var(self, node:mm.Var):
         self.vars.add(node)
 
+class NodeFinder(ContainerWalker):
+
+    def is_descendant(self, node:mm.Node, ancestor:mm.Task) -> bool:
+        """ Ture if node is a descendant of the ancestor task, i.e. one of its children,
+        or a child of one of its children, etc.
+        """
+        self.node = node
+        self.found = False
+        try:
+            self(ancestor)
+            return self.found
+        finally:
+            self.node = None
+            self.found = False
+
+    def enter_container(self, container, children):
+        if self.found:
+            return NO_WALK
+        return None
+
+    def _visit_node(self, node: mm.Node):
+        if node == self.node:
+            self.found = True
+            return node
+        return super()._visit_node(node)
+
 #------------------------------------------------------
 # Helper functions shared across shim modules
 #------------------------------------------------------

relationalai/shims/hoister.py

@@ -1,15 +1,13 @@
 from __future__ import annotations
 
 from collections import defaultdict
-import itertools
 from dataclasses import dataclass
-from typing import  Iterable, Sequence as seq
+from typing import  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
-from .helpers import is_output_update, is_main_output, ContainerWalker
+from .helpers import is_output_update, is_main_output, ContainerWalker, NodeFinder
 
 
 @dataclass(frozen=True)
@@ -60,6 +58,8 @@ class Scope:
 class Hoister(ContainerWalker):
     def __init__(self):
         super().__init__()
+        # helper for finding children nodes
+        self.finder = NodeFinder()
         # the result of the whole analysis
         self.hoists: dict[mm.Container, OrderedSet[mm.Var]] = defaultdict(OrderedSet[mm.Var])
         self.container_vars: dict[mm.Container, set[mm.Var]] = defaultdict(set)
@@ -130,6 +130,15 @@ class Hoister(ContainerWalker):
                     self.hoists[child] = parent_hoists
                     self.allowed_vars[child] = allowed
             self.allowed_vars[container] = allowed
+            # this might be a match/union that just acts as a filter for some sibling logical,
+            # so we need to create an artificial demand for the allowed vars in the parent container
+            # to hoist them up if there's no other output providing them.
+            # ex: define(foo(Person)).where(union(Person.age > 10, Person.age < 5))
+            # parent = self.scope.parent.get(container, None)
+            # if parent:
+            #     for allowed_var in allowed:
+            #         self._register_use(allowed_var, input=True, container=parent)
+            #         self.scope.var_refs[allowed_var]
 
         if container.scope:
             # when leaving the scope, compute hoisted vars
@@ -141,6 +150,13 @@ class Hoister(ContainerWalker):
                 for ref in var_refs:
                     if ref.is_input() and not any(self.scope.is_ancestor(other.node, ref.node) for other in var_refs if other != ref and other.is_output()):
                         requires_hoist.add(ref)
+                    # if this is a require, the ref is in the check or error, and there's another
+                    # ref in the domain, we need to hoist in the domain to make it available to
+                    # check/error
+                    elif isinstance(container, mm.Require):
+                        if self.finder.is_descendant(ref.node, container.check) or ( container.error and self.finder.is_descendant(ref.node, container.error)):
+                            if any(self.finder.is_descendant(other.node, container.domain) for other in var_refs if other != ref):
+                                requires_hoist.add(ref)
 
                 # for references that require a hoist, find some other output and hoist
                 # that up until the least common ancestor of the two references
@@ -150,7 +166,9 @@ class Hoister(ContainerWalker):
                             # find least common ancestor of the two references
                             lca = self.scope.least_common_ancestor(ref.node, other.node)
                             assert(lca is not None)
-                            if var not in self.container_vars[lca]:
+                            # we should never hoist vars where the common ancestor is a match/union because that
+                            # means we're hopping variables across branches
+                            if var not in self.container_vars[lca] and not isinstance(lca, (mm.Match, mm.Union)):
                                 self._hoist_until(var, other.node, lca)
 
         # pop container from stack
@@ -188,9 +206,9 @@ class Hoister(ContainerWalker):
             self.scope.var_refs[var].add(VarRef(parent, input))
             self.scope.available_vars[parent].add(var)
 
-    def lookup(self, l: mm.Lookup):
-        for arg, field in zip(l.args, l.relation.fields):
-            if not field.input or l.relation == b.core["="]:
+    def lookup(self, lookup_node: mm.Lookup):
+        for arg, field in zip(lookup_node.args, lookup_node.relation.fields):
+            if not field.input or lookup_node.relation == b.core["="]:
                 self._register_use(arg, input=False)
             else:
                 self._register_use(arg, input=True)
@@ -209,7 +227,8 @@ class Hoister(ContainerWalker):
                     self._register_use(arg, input=True, container=main_output_container)
 
     def aggregate(self, a: mm.Aggregate):
-        self._register_use(a.projection, input=True)
+        for p in a.projection:
+            self._register_use(p, input=True)
         for g in a.group:
             self._register_use(g, input=True)
         for arg, field in zip(a.args, a.aggregation.fields):