relationalai 0.12.11__py3-none-any.whl → 0.12.13__py3-none-any.whl

relationalai/clients/snowflake.py

@@ -441,7 +441,8 @@ class Resources(ResourcesBase):
         code: str,
         params: List[Any] | Any | None = None,
         raw: bool = False,
-        help: bool = True
+        help: bool = True,
+        skip_engine_db_error_retry: bool = False
     ) -> Any:
         # print(f"\n--- sql---\n{code}\n--- end sql---\n")
         if not self._session:
@@ -467,7 +468,8 @@ class Resources(ResourcesBase):
             if re.search(f"database '{rai_app}' does not exist or not authorized.".lower(), orig_message):
                 exception = SnowflakeAppMissingException(rai_app, current_role)
                 raise exception from None
-            if _is_engine_issue(orig_message) or _is_database_issue(orig_message):
+            # skip initializing the index if the query is a user transaction. exec_raw/exec_lqp will handle that case with the correct request headers.
+            if (_is_engine_issue(orig_message) or _is_database_issue(orig_message)) and not skip_engine_db_error_retry:
                 try:
                     self._poll_use_index(
                         app_name=self.get_app_name(),
@@ -1609,9 +1611,11 @@ Otherwise, remove it from your '{profile}' configuration profile.
                 sql_string = f"CALL {APP_NAME}.api.exec_async_v2('{database}','{engine}', ?, {inputs}, {readonly}, {nowait_durable}, '{language}', {query_timeout_mins}, {request_headers});"
             else:
                 sql_string = f"CALL {APP_NAME}.api.exec_async_v2('{database}','{engine}', ?, {inputs}, {readonly}, {nowait_durable}, '{language}', {request_headers});"
+        # Don't let exec setup GI on failure, exec_raw and exec_lqp will do that and add the correct headers.
         response = self._exec(
             sql_string,
             raw_code,
+            skip_engine_db_error_retry=True,
         )
         if not response:
             raise Exception("Failed to create transaction")
@@ -1629,6 +1633,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
         bypass_index=False,
         language: str = "rel",
         query_timeout_mins: int | None = None,
+        gi_setup_skipped: bool = False,
     ):
         if inputs is None:
             inputs = {}
@@ -1638,6 +1643,8 @@ Otherwise, remove it from your '{profile}' configuration profile.
         with debugging.span("transaction", **query_attrs_dict) as txn_span:
             with debugging.span("create_v2", **query_attrs_dict) as create_span:
                 request_headers['user-agent'] = get_pyrel_version(self.generation)
+                request_headers['gi_setup_skipped'] = str(gi_setup_skipped)
+                request_headers['pyrel_program_id'] = debugging.get_program_span_id() or ""
                 response = self._exec_rai_app(
                     database=database,
                     engine=engine,
@@ -1897,26 +1904,29 @@ Otherwise, remove it from your '{profile}' configuration profile.
     # Exec
     #--------------------------------------------------
 
-    def exec_lqp(
+    def _exec_with_gi_retry(
         self,
         database: str,
         engine: str | None,
-        raw_code: bytes,
-        readonly=True,
-        *,
-        inputs: Dict | None = None,
-        nowait_durable=False,
-        headers: Dict | None = None,
-        bypass_index=False,
-        query_timeout_mins: int | None = None,
+        raw_code: str,
+        inputs: Dict | None,
+        readonly: bool,
+        nowait_durable: bool,
+        headers: Dict | None,
+        bypass_index: bool,
+        language: str,
+        query_timeout_mins: int | None,
     ):
-        raw_code_b64 = base64.b64encode(raw_code).decode("utf-8")
+        """Execute with graph index retry logic.
 
+        Attempts execution with gi_setup_skipped=True first. If an engine or database
+        issue occurs, polls use_index and retries with gi_setup_skipped=False.
+        """
         try:
             return self._exec_async_v2(
-                database, engine, raw_code_b64, inputs, readonly, nowait_durable,
-                headers=headers, bypass_index=bypass_index, language='lqp',
-                query_timeout_mins=query_timeout_mins,
+                database, engine, raw_code, inputs, readonly, nowait_durable,
+                headers=headers, bypass_index=bypass_index, language=language,
+                query_timeout_mins=query_timeout_mins, gi_setup_skipped=True,
             )
         except Exception as e:
             err_message = str(e).lower()
@@ -1933,13 +1943,32 @@ Otherwise, remove it from your '{profile}' configuration profile.
                 )
 
                 return self._exec_async_v2(
-                    database, engine, raw_code_b64, inputs, readonly, nowait_durable,
-                    headers=headers, bypass_index=bypass_index, language='lqp',
-                    query_timeout_mins=query_timeout_mins,
+                    database, engine, raw_code, inputs, readonly, nowait_durable,
+                    headers=headers, bypass_index=bypass_index, language=language,
+                    query_timeout_mins=query_timeout_mins, gi_setup_skipped=False,
                 )
             else:
                 raise e
 
+    def exec_lqp(
+        self,
+        database: str,
+        engine: str | None,
+        raw_code: bytes,
+        readonly=True,
+        *,
+        inputs: Dict | None = None,
+        nowait_durable=False,
+        headers: Dict | None = None,
+        bypass_index=False,
+        query_timeout_mins: int | None = None,
+    ):
+        raw_code_b64 = base64.b64encode(raw_code).decode("utf-8")
+        return self._exec_with_gi_retry(
+            database, engine, raw_code_b64, inputs, readonly, nowait_durable,
+            headers, bypass_index, 'lqp', query_timeout_mins
+        )
+
 
     def exec_raw(
         self,
@@ -1955,45 +1984,10 @@ Otherwise, remove it from your '{profile}' configuration profile.
         query_timeout_mins: int | None = None,
     ):
         raw_code = raw_code.replace("'", "\\'")
-
-        try:
-            return self._exec_async_v2(
-                database,
-                engine,
-                raw_code,
-                inputs,
-                readonly,
-                nowait_durable,
-                headers=headers,
-                bypass_index=bypass_index,
-                query_timeout_mins=query_timeout_mins,
-            )
-        except Exception as e:
-            err_message = str(e).lower()
-            if _is_engine_issue(err_message) or _is_database_issue(err_message):
-                engine_name = engine or self.get_default_engine_name()
-                engine_size = self.config.get_default_engine_size()
-                self._poll_use_index(
-                    app_name=self.get_app_name(),
-                    sources=self.sources,
-                    model=database,
-                    engine_name=engine_name,
-                    engine_size=engine_size,
-                    headers=headers,
-                )
-                return self._exec_async_v2(
-                    database,
-                    engine,
-                    raw_code,
-                    inputs,
-                    readonly,
-                    nowait_durable,
-                    headers=headers,
-                    bypass_index=bypass_index,
-                    query_timeout_mins=query_timeout_mins,
-                )
-            else:
-                raise e
+        return self._exec_with_gi_retry(
+            database, engine, raw_code, inputs, readonly, nowait_durable,
+            headers, bypass_index, 'rel', query_timeout_mins
+        )
 
 
     def format_results(self, results, task:m.Task|None=None) -> Tuple[DataFrame, List[Any]]:
@@ -3281,6 +3275,7 @@ class DirectAccessResources(Resources):
         path_params: Dict[str, str] | None = None,
         query_params: Dict[str, str] | None = None,
         skip_auto_create: bool = False,
+        skip_engine_db_error_retry: bool = False,
     ) -> requests.Response:
         with debugging.span("direct_access_request"):
             def _send_request():
@@ -3314,8 +3309,8 @@ class DirectAccessResources(Resources):
                         message = ""  # Not used when we check status_code directly
 
                     # fix engine on engine error and retry
-                    # Skip auto-retry if skip_auto_create is True to avoid recursion
-                    if (_is_engine_issue(message) and not skip_auto_create) or _is_database_issue(message):
+                    # Skip setting up GI if skip_auto_create is True to avoid recursion or skip_engine_db_error_retry is true to let _exec_async_v2 perform the retry with the correct headers.
+                    if ((_is_engine_issue(message) and not skip_auto_create) or _is_database_issue(message)) and not skip_engine_db_error_retry:
                         engine_name = payload.get("caller_engine_name", "") if payload else ""
                         engine_name = engine_name or self.get_default_engine_name()
                         engine_size = self.config.get_default_engine_size()
@@ -3340,6 +3335,48 @@ class DirectAccessResources(Resources):
                 raise e
             return response
 
+    def _txn_request_with_gi_retry(
+        self,
+        payload: Dict,
+        headers: Dict[str, str],
+        query_params: Dict,
+        engine: Union[str, None],
+    ):
+        """Make request with graph index retry logic.
+
+        Attempts request with gi_setup_skipped=True first. If an engine or database
+        issue occurs, polls use_index and retries with gi_setup_skipped=False.
+        """
+        response = self.request(
+            "create_txn", payload=payload, headers=headers, query_params=query_params, skip_auto_create=True, skip_engine_db_error_retry=True
+        )
+
+        if response.status_code != 200:
+            try:
+                message = response.json().get("message", "")
+            except requests.exceptions.JSONDecodeError:
+                message = ""
+
+            if _is_engine_issue(message) or _is_database_issue(message):
+                engine_name = engine or self.get_default_engine_name()
+                engine_size = self.config.get_default_engine_size()
+                self._poll_use_index(
+                    app_name=self.get_app_name(),
+                    sources=self.sources,
+                    model=self.database,
+                    engine_name=engine_name,
+                    engine_size=engine_size,
+                    headers=headers,
+                )
+                headers['gi_setup_skipped'] = 'False'
+                response = self.request(
+                    "create_txn", payload=payload, headers=headers, query_params=query_params, skip_auto_create=True, skip_engine_db_error_retry=True
+                )
+            else:
+                raise ResponseStatusException("Failed to create transaction.", response)
+
+        return response
+
     def _exec_async_v2(
         self,
         database: str,
@@ -3352,6 +3389,7 @@ class DirectAccessResources(Resources):
         bypass_index=False,
         language: str = "rel",
         query_timeout_mins: int | None = None,
+        gi_setup_skipped: bool = False,
     ):
 
         with debugging.span("transaction") as txn_span:
@@ -3374,12 +3412,15 @@ class DirectAccessResources(Resources):
                     payload["timeout_mins"] = query_timeout_mins
                 query_params={"use_graph_index": str(use_graph_index and not bypass_index)}
 
-                response = self.request(
-                    "create_txn", payload=payload, headers=headers, query_params=query_params,
-                )
+                # Add gi_setup_skipped to headers
+                if headers is None:
+                    headers = {}
+                headers["gi_setup_skipped"] = str(gi_setup_skipped)
+                headers['pyrel_program_id'] = debugging.get_program_span_id() or ""
 
-                if response.status_code != 200:
-                    raise ResponseStatusException("Failed to create transaction.", response)
+                response = self._txn_request_with_gi_retry(
+                    payload, headers, query_params, engine
+                )
 
                 artifact_info = {}
                 response_content = response.json()

relationalai/semantics/lqp/passes.py

@@ -29,8 +29,8 @@ def lqp_passes() -> list[Pass]:
         InferTypes(),
         DNFUnionSplitter(),
         ExtractKeys(),
-        ExtractCommon(),
         FormatOutputs(),
+        ExtractCommon(), # Extracts tasks that will become common after Flatten into their own definition
         Flatten(),
         Splinter(), # Splits multi-headed rules into multiple rules
         QuantifyVars(), # Adds missing existentials

relationalai/semantics/lqp/rewrite/annotate_constraints.py

@@ -10,7 +10,7 @@ from relationalai.semantics.lqp.rewrite.functional_dependencies import (
     is_valid_unique_constraint, normalized_fd
 )
 
-
+_DISABLE_CONSTRAINT_DECLARATIONS = True
 
 class AnnotateConstraints(Pass):
     """
@@ -36,6 +36,8 @@ class AnnotateConstraintsRewriter(DischargeConstraintsVisitor):
     """
 
     def _should_be_declarable_constraint(self, node: Require) -> bool:
+        if _DISABLE_CONSTRAINT_DECLARATIONS:
+            return False
         if not is_valid_unique_constraint(node):
             return False
         # Currently, we only declare non-structural functional dependencies.

relationalai/semantics/lqp/rewrite/extract_common.py

@@ -3,11 +3,13 @@ from __future__ import annotations
 from dataclasses import dataclass
 from typing import Optional
 from relationalai.semantics.metamodel import ir, factory as f, helpers, visitor
-from relationalai.semantics.metamodel.compiler import Pass, group_tasks
+from relationalai.semantics.metamodel.compiler import Pass
 from relationalai.semantics.metamodel.util import OrderedSet, ordered_set
 from relationalai.semantics.metamodel import dependency
 from relationalai.semantics.metamodel import builtins
 
+from functools import reduce
+
 class ExtractCommon(Pass):
     """
     Pass to analyze Logical bodies and extract lookups in their own Logical if it makes
@@ -71,7 +73,7 @@ class ExtractCommon(Pass):
         )
 
     #--------------------------------------------------
-    # IR handlers
+    # Extra classes
     #--------------------------------------------------
 
     class Context():
@@ -80,8 +82,26 @@ class ExtractCommon(Pass):
             self.info = dependency.analyze(model.root)
             self.options = options
 
+    @dataclass
+    class ExtractionPlan():
+        # tasks to extract to the body of the common logical
+        common_body: OrderedSet[ir.Task]
+        # tasks to remain in the original body
+        remaining_body: OrderedSet[ir.Task]
+        # variables to be exposed by the common logical
+        exposed_vars: OrderedSet[ir.Var]
+        # map from nested composite to the tasks in the common body that still need to be
+        # included in its body, because it contains variables not exposed by the common logical
+        local_dependencies: dict[ir.Task, OrderedSet[ir.Task]]
+        # a reference to the common connection created for this plan, if any
+        common_reference: Optional[ir.Lookup] = None
+
+    #--------------------------------------------------
+    # IR handlers
+    #--------------------------------------------------
+
     def handle(self, task: ir.Task, ctx: Context):
-        # currently we only extract if it's a sequence of Logicals, but we could in the
+        # Currently we only extract if it's a sequence of Logicals, but we could in the
         # future support other intermediate nodes
         if isinstance(task, ir.Logical):
             return self.handle_logical(task, ctx)
@@ -89,66 +109,52 @@ class ExtractCommon(Pass):
             return task
 
     def handle_logical(self, task: ir.Logical, ctx: Context):
-        # Process the original body to find binders and extractables. The Flatten pass later
-        # will extract out both composites and effects, so we group them together here
-        groups = group_tasks(task.body, {
-            "binders": helpers.BINDERS,
-            "composites_and_effects": helpers.COMPOSITES + helpers.EFFECTS,
-        })
-        binders = groups["binders"]
-        composites_and_effects = groups["composites_and_effects"]
-
-        # the new body of the rewritten task
+        # Process the original body to find candidates for extraction. Each task is in one
+        # of three categories:
+        # - Binders: tasks that bind variables. These are candidates for extracting into
+        #   the common body.
+        # - Flattenables: tasks that will later be extracted by the Flatten pass
+        # - Other: tasks that are neither binders nor flattenables; these will remain
+        #   in the body as-is.
+
+        binders = ordered_set()
+        flattenables = ordered_set()
+        other = ordered_set()
+
+        for child in task.body:
+            if _is_binder(child):
+                binders.add(child)
+            elif _is_flattenable(ctx, child):
+                flattenables.add(child)
+            else:
+                other.add(child)
+
+        # The new body of the rewritten task
         body:OrderedSet[ir.Task] = ordered_set()
 
-        # quick check to see if it's worth doing more analysis; we only want to extract
-        # common binders if there are multiple, and there are also multiple composites
+        # Quick check to see if it's worth doing more analysis; we only want to extract
+        # common binders if there are multiple, and there are also multiple flattenables
         # that will be extracted by the flatten pass later (so that they can share the
         # extracted logic).
-        plan = None
-        if len(binders) > 1 and composites_and_effects:
-            extractables = self._get_extractables(ctx, composites_and_effects)
-            # only makes sense to extract common if at least one nested composite will be
-            # extracted during Flatten
-            if extractables:
-                # make a plan to extract common tasks from the logical
-                plan = self._create_extraction_plan(binders, composites_and_effects, extractables, ctx)
-                if plan and len(composites_and_effects) > 1:
-                    # plan is worthwhile and there are multiple composites, extract the common body and add the connection to the body
-                    exposed_vars = plan.exposed_vars.get_list()
-                    plan.common_reference = f.lookup(helpers.extract(task, plan.common_body, exposed_vars, ctx.rewrite_ctx, "common"), exposed_vars)
-                    # if we are not distributing the reference, add to the main body
-                    if not plan.distribute_common_reference:
-                        body.add(plan.common_reference)
-
-        # if we have a plan and will distribute the common reference, keep track of
-        # variables still needed by the remaining tasks, as they need to be hoisted by
-        # the remaining composites that get the common reference
-        remaining_vars = None
-        if plan and plan.distribute_common_reference:
-            # add variables hoisted by this logical that are in the exposed vars, to
-            # make sure they are hoisted all the way through
-            remaining_vars = OrderedSet.from_iterable(helpers.hoisted_vars(task.hoisted)) & plan.exposed_vars
-            for child in task.body:
-                if child in groups["other"] or child not in plan.remaining_body or child in composites_and_effects:
-                    continue
-                remaining_vars.update(ctx.info.task_inputs(child))
-                remaining_vars.update(ctx.info.task_outputs(child))
-            remaining_vars = remaining_vars & plan.exposed_vars
-
-        # if the plan was not used in one of the cases above, ignore it completely, we
-        # are not extracting common nor distributing it around
-        if plan and not plan.distribute_common_reference and not len(composites_and_effects) > 1:
-            plan = None
+        plan: Optional[ExtractCommon.ExtractionPlan] = None
+        if len(binders) > 1 and len(flattenables) > 1:
+            plan = _create_extraction_plan(ctx, binders, flattenables, other)
+            if plan:
+                # plan is worthwhile, extract the common body and add the connection to the body
+                exposed_vars = plan.exposed_vars.get_list()
+                plan.common_reference = f.lookup(helpers.extract(task, plan.common_body, exposed_vars, ctx.rewrite_ctx, "common"), exposed_vars)
+
+                # Add plan common reference to the body.
+                body.add(plan.common_reference)
 
         # recursively handle children
         for child in task.body:
             # skip children that were extracted
-            if plan and child not in groups["other"] and child not in plan.remaining_body and child not in composites_and_effects:
+            if plan and child not in other and child not in plan.remaining_body and child not in flattenables:
                 continue
 
             # no plan or child is not a composite, so just add the handled to the body
-            if not plan or child not in composites_and_effects:
+            if not plan or child not in flattenables:
                 body.add(self.handle(child, ctx))
                 continue
 
@@ -156,7 +162,7 @@ class ExtractCommon(Pass):
             replacement = self.handle(child, ctx)
 
             # this child needs either extra local dependencies or the common reference
-            if child in plan.local_dependencies or plan.distribute_common_reference:
+            if child in plan.local_dependencies:
                 # the new body will have maybe the common reference and the local deps
                 replacement_body = ordered_set()
 
@@ -172,19 +178,6 @@ class ExtractCommon(Pass):
                         dep_outputs.update(ctx.info.task_outputs(d))
                     hoisted.update(dep_outputs & ctx.info.task_inputs(replacement))
 
-                if plan.distribute_common_reference:
-                    if len(composites_and_effects) == 1:
-                        # if there's a single composite, just insert the whole common body into it
-                        replacement_body.update(plan.common_body)
-                    else:
-                        # otherwise insert a clone of the reference on the extracted rule
-                        assert(plan.common_reference)
-                        replacement_body.add(plan.common_reference.clone())
-                    # add remaining vars to hoisted, making sure there's no duplicates (due to VarOrDefault)
-                    hoisted_vars = helpers.hoisted_vars(hoisted)
-                    if remaining_vars:
-                        hoisted = OrderedSet.from_iterable(filter(lambda v: v not in hoisted_vars, remaining_vars)) | hoisted
-
                 if child in plan.local_dependencies:
                     for local_dep in plan.local_dependencies[child]:
                         replacement_body.add(local_dep.clone())
@@ -208,177 +201,138 @@ class ExtractCommon(Pass):
 
         return ir.Logical(task.engine, task.hoisted, tuple(body))
 
-    @dataclass
-    class ExtractionPlan():
-        # tasks to extract to the body of the common logical
-        common_body: OrderedSet[ir.Task]
-        # tasks to remain in the original body
-        remaining_body: OrderedSet[ir.Task]
-        # variables to be exposed by the common logical
-        exposed_vars: OrderedSet[ir.Var]
-        # map from nested composite to the tasks in the common body that still need to be
-        # included in its body, because it contains variables not exposed by the common logical
-        local_dependencies: dict[ir.Task, OrderedSet[ir.Task]]
-        # whether the common reference should be distributed to composites
-        distribute_common_reference: bool
-        # a reference to the common connection created for this plan, if any
-        common_reference: Optional[ir.Lookup] = None
 
+#--------------------------------------------------
+# Utilities
+#--------------------------------------------------
 
-    def _create_extraction_plan(self, binders: OrderedSet[ir.Task], composites: OrderedSet[ir.Task], extractables: list[ir.Task], ctx: Context):
-        """
-        Compute a plan to extract tasks in this frame that are common dependencies
-        across these composite tasks.
-        """
-        # If there are any pragma lookups, then don't extract anything. Pragma lookups are
-        # designed to control execution order, and extracting them may affect their
-        # semantics.
-        for b in binders:
-            if isinstance(b, ir.Lookup) and builtins.is_pragma(b.relation):
-                return None
-
-        # Compute intersection of task dependencies
-        sample = composites.some()
-        deps = ctx.info.task_dependencies(sample)
-        if deps is None:
+def _create_extraction_plan(ctx: ExtractCommon.Context, binders: OrderedSet[ir.Task], flattenables: OrderedSet[ir.Task], others: OrderedSet[ir.Task]) -> Optional[ExtractCommon.ExtractionPlan]:
+    """
+    Compute a plan to extract tasks in this frame that are common dependencies
+    across these composite tasks.
+    """
+    # If there are any pragma lookups, then don't extract anything. Pragma lookups are
+    # designed to control execution order, and extracting them may affect their
+    # semantics.
+    for b in binders:
+        if isinstance(b, ir.Lookup) and builtins.is_pragma(b.relation):
             return None
-        # only get sibling dependencies
-        common_body = binders & deps
-
-        # For other composites, remove their sibling dependencies so that we end up with
-        # the intersection of dependencies
-        for composite in composites:
-            if composite is sample:
-                continue
 
-            # compute sibling dependencies
-            deps = ctx.info.task_dependencies(composite)
+    # Compute intersection of task dependencies
+    all_deps = [ctx.info.task_dependencies(f) for f in flattenables]
+    deps = reduce(lambda a, b: a & b, all_deps)
+    common_body = binders & deps
+
+    # We don't need to extract anything if there's only zero or one common tasks
+    if len(common_body) < 2:
+        return None
+
+    # Keep track of remaining tasks that are not extracted in the common body
+    remaining = ordered_set()
+
+    # Compute the vars that should be output from the common body. These are the union of
+    # all input vars across all non-extracted tasks, intersected with output vars of
+    # the common body.
+
+    # First, compute the output vars of the common body
+    common_body_output_vars = OrderedSet()
+    for child in common_body:
+        common_body_output_vars.update(ctx.info.task_outputs(child))
+
+    # Next, compute the union of the input vars of all non-extracted tasks
+    non_extracted_tasks = (binders - common_body) | flattenables | others
+    all_exposed_vars: list[OrderedSet[ir.Var]] = []
+    for t in non_extracted_tasks:
+        input_vars = ctx.info.task_inputs(t)
+        all_exposed_vars.append(input_vars if input_vars else OrderedSet())
+
+    exposed_vars = reduce(lambda a, b: a | b, all_exposed_vars) & common_body_output_vars
+
+    # If there are no vars in common, then it's not worth extracting
+    if not exposed_vars:
+        return None
+
+    # Make sure that all local dependencies of the common body are included in the common
+    # body. This is important for the safety of this rewrite.
+    for task in common_body:
+        local_deps = ctx.info.local_dependencies(task)
+        if local_deps:
+            common_body.update(local_deps & binders)
+
+    # check which of the original binders remain, and make sure their dependencies also stay
+    for binder in binders:
+        if binder not in common_body:
+            remaining.add(binder)
+            deps = _compute_local_dependencies(ctx, binders, binder, exposed_vars)
             if deps:
-                for task in common_body:
-                    if task not in deps:
-                        common_body.remove(task)
-
-        # Compute union of input vars
-        # Start with the output vars of the common body. We only want to expose vars that
-        # are output from the common body
-        body_output_vars = OrderedSet()
-        for child in common_body:
-            body_output_vars.update(ctx.info.task_outputs(child))
-
-        # Compute the union of input vars across all non-extracted tasks (basically
-        # composites and binders left behind), intersected with output
-        # vars of the common body
-        exposed_vars = OrderedSet.from_iterable(ctx.info.task_inputs(sample)) & body_output_vars
-        non_extracted_tasks = (binders - common_body) | composites
-        for composite in non_extracted_tasks:
-            if composite is sample:
-                continue
-            # compute common input vars
-            t_inputs = OrderedSet.from_iterable(ctx.info.task_inputs(composite))
-            exposed_vars.update(t_inputs & body_output_vars)
+                remaining.update(deps)
 
-        # no vars in common, not worth to extract
-        if not exposed_vars:
-            return None
+    # for each composite, check if there are additional tasks needed, because the task
+    # depends on it but it is not exposed by the vars
+    local_dependencies: dict[ir.Task, OrderedSet[ir.Task]] = dict()
+    for flattenable in flattenables:
+        local = _compute_local_dependencies(ctx, binders, flattenable, exposed_vars)
+        if local:
+            local_dependencies[flattenable] = local
 
-        for task in common_body:
-            local_deps = ctx.info.local_dependencies(task)
-            if local_deps:
-                common_body.update(local_deps & binders)
+    return ExtractCommon.ExtractionPlan(common_body, remaining, exposed_vars, local_dependencies)
 
-        # not useful to extract common tasks if there's a single one
-        if len(common_body) < 2:
-            return None
-
-        # check if some variable used in the common body is needed by some binder that is
-        # not going to be extracted. In that case, we need to expose this variable from the
-        # common body
-        common_vars = ordered_set()
-        for task in common_body:
-            common_vars.update(ctx.info.task_outputs(task))
-        common_vars = common_vars - exposed_vars
-        for v in common_vars:
-            for binder in binders:
-                if binder not in common_body and ctx.info.task_inputs(binder) and v in ctx.info.task_inputs(binder):
-                    exposed_vars.add(v)
-                    break
-
-        # check which of the original binders remain, and make sure their dependencies also stay
-        remaining = ordered_set()
-        for binder in binders:
-            if binder not in common_body:
-                remaining.add(binder)
-                deps = self._compute_local_dependencies(ctx, binders, binder, exposed_vars)
-                if deps:
-                    remaining.update(deps)
-
-        # for each composite, check if there are additional tasks needed, because the task
-        # depends on it but it is not exposed by the vars
-        local_dependencies: dict[ir.Task, OrderedSet[ir.Task]] = dict()
-        for composite in composites:
-            local = self._compute_local_dependencies(ctx, binders, composite, exposed_vars)
-            if local:
-                local_dependencies[composite] = local
-
-        # distribute the common reference only if all of the composites are extractable and there's nothing else remaining
-        distribute_common_reference = len(extractables) == len(composites) and not remaining
-
-        return ExtractCommon.ExtractionPlan(common_body, remaining, exposed_vars, local_dependencies, distribute_common_reference)
-
-    def _compute_local_dependencies(self, ctx: Context, binders: OrderedSet[ir.Task], composite: ir.Task, exposed_vars: OrderedSet[ir.Var]):
-        """
-        The tasks in common_body will be extracted into a logical that will expose the exposed_vars.
-        Compute which additional dependencies are needed specifically for this composite, because
-        it depends on some tasks that are extracted to common_body but not exposed by exposed_vars.
-        """
-
-        # working list of vars we still need to fulfill
-        inputs = ctx.info.task_inputs(composite)
-        if not inputs:
-            return None
-
-        # vars exposed by exposed vars + tasks added to the local body
-        vars_exposed = OrderedSet.from_iterable(exposed_vars)
-        vars_needed = (inputs - vars_exposed)
-        if not vars_needed:
-            return None
+def _compute_local_dependencies(ctx: ExtractCommon.Context, binders: OrderedSet[ir.Task], composite: ir.Task, exposed_vars: OrderedSet[ir.Var]):
+    """
+    The tasks in common_body will be extracted into a logical that will expose the exposed_vars.
+    Compute which additional dependencies are needed specifically for this composite, because
+    it depends on some tasks that are extracted to common_body but not exposed by exposed_vars.
+    """
 
-        # this is a greedy algorithm that uses the first task in the common body that provides
-        # a variable needed; it may result in sub-optimal extraction, but should be correct
-        local_body = ordered_set()
-        while(vars_needed):
-            v = vars_needed.pop()
-            for x in binders:
-                if x not in local_body:
-                    # an x that is not yet in local_body can fulfill v
-                    x_outputs = ctx.info.task_outputs(x)
-                    if x_outputs and v in x_outputs:
-                        # add it to local_body and add its outputs to vars exposed
-                        local_body.add(x)
-                        vars_exposed.add(x_outputs)
-                        # but add its inputs the vars now needed
-                        inputs = ctx.info.task_inputs(x)
-                        if inputs:
-                            vars_needed.update(inputs - vars_exposed)
-        return local_body
-
-    def _get_extractables(self, ctx: Context, composites: OrderedSet[ir.Task]):
-        """
-        Extractables are tasks that will eventually be extracted by the Flatten pass later.
-        Given a set of tasks, return the extractable ones.
-        """
-        def _extractable(t: ir.Task):
-            # With GNF outputs (i.e., wide_outputs = False), the output tasks will be
-            # extracted into separate top-level single-column outputs.
-            if isinstance(t, ir.Output) and not ctx.options.get("wide_outputs", False):
-                return True
-
-            extractable_types = (ir.Update, ir.Aggregate, ir.Match, ir.Union, ir.Rank)
-            return isinstance(t, ir.Logical) and len(visitor.collect_by_type(extractable_types, t)) > 0
-
-        extractables = []
-        for t in composites:
-            if _extractable(t):
-                extractables.append(t)
-
-        return extractables
+    # working list of vars we still need to fulfill
+    inputs = ctx.info.task_inputs(composite)
+    if not inputs:
+        return None
+
+    # vars exposed by exposed vars + tasks added to the local body
+    vars_exposed = OrderedSet.from_iterable(exposed_vars)
+    vars_needed = (inputs - vars_exposed)
+    if not vars_needed:
+        return None
+
+    # this is a greedy algorithm that uses the first task in the common body that provides
+    # a variable needed; it may result in sub-optimal extraction, but should be correct
+    local_body = ordered_set()
+    while(vars_needed):
+        v = vars_needed.pop()
+        for x in binders:
+            if x not in local_body:
+                # an x that is not yet in local_body can fulfill v
+                x_outputs = ctx.info.task_outputs(x)
+                if x_outputs and v in x_outputs:
+                    # add it to local_body and add its outputs to vars exposed
+                    local_body.add(x)
+                    vars_exposed.add(x_outputs)
+                    # but add its inputs the vars now needed
+                    inputs = ctx.info.task_inputs(x)
+                    if inputs:
+                        vars_needed.update(inputs - vars_exposed)
+    return local_body
+
+def _is_binder(task: ir.Task):
+    # If the task itself is a binder
+    if any(isinstance(task, binder) for binder in (ir.Lookup, ir.Construct, ir.Exists, ir.Data, ir.Not)):
+        return True
+
+    # If the task is a Logical containing only binders
+    if isinstance(task, ir.Logical) and all(_is_binder(c) for c in task.body):
+        return True
+
+    # If the task is a Union containing only binders
+    if isinstance(task, ir.Union) and all(_is_binder(c) for c in task.tasks):
+        return True
+
+    return False
+
+def _is_flattenable(ctx: ExtractCommon.Context, task: ir.Task):
+    # Each output will be flattened into its own top-level def
+    if isinstance(task, ir.Output):
+        return True
+
+    extractable_types = (ir.Update, ir.Aggregate, ir.Match, ir.Rank)
+    return isinstance(task, ir.Logical) and len(visitor.collect_by_type(extractable_types, task)) > 0

relationalai/semantics/metamodel/rewrite/dnf_union_splitter.py

@@ -140,13 +140,13 @@ class DNFExtractor(Visitor):
                         for new_task in replacement_tasks:
                             # copy to mutate
                             new_body = list(body)
-                            new_body.append(new_task)
+                            new_body.append(new_task.clone())
                             new_replacement_bodies.append(new_body)
                     replacement_bodies = new_replacement_bodies
 
                 else:
                     for new_body in replacement_bodies:
-                        new_body.append(task)
+                        new_body.append(task.clone())
 
             replacement_tasks: list[ir.Task] = []
             for body in replacement_bodies:

relationalai/semantics/metamodel/rewrite/flatten.py

@@ -211,6 +211,17 @@ class Flatten(Pass):
                 if child in composites:
                     all_composites_removed = False
 
+        # Filter out empty logicals from the body
+        new_body:OrderedSet[ir.Task] = ordered_set()
+        for b in body:
+            if isinstance(b, ir.Logical):
+                if not b.body:
+                    # empty logical, skip
+                    continue
+            new_body.add(b)
+
+        body = new_body
+
         # all children were extracted or all composites were removed without any effects
         # left and no outputs (so no way for outer dependencies), drop this logical
         if not body or (all_composites_removed and not any([isinstance(t, helpers.EFFECTS) for t in body]) and not ctx.info.task_outputs(task)):

relationalai/semantics/rel/compiler.py

@@ -33,8 +33,8 @@ class Compiler(c.Compiler):
             InferTypes(),
             DNFUnionSplitter(),
             ExtractKeys(),
-            ExtractCommon(),
             FormatOutputs(),
+            ExtractCommon(),
             Flatten(),
             Splinter(),
             QuantifyVars(),

{relationalai-0.12.11.dist-info → relationalai-0.12.13.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: relationalai
-Version: 0.12.11
+Version: 0.12.13
 Summary: RelationalAI Library and CLI
 Author-email: RelationalAI <support@relational.ai>
 License-File: LICENSE

{relationalai-0.12.11.dist-info → relationalai-0.12.13.dist-info}/RECORD

@@ -29,7 +29,7 @@ relationalai/clients/hash_util.py,sha256=pZVR1FX3q4G_19p_r6wpIR2tIM8_WUlfAR7AVZJ
 relationalai/clients/local.py,sha256=uX1Or2WO0juDuqa6TCCvm3G2ieP6p0PtYp_jfrCMlVc,23577
 relationalai/clients/profile_polling.py,sha256=pUH7WKH4nYDD0SlQtg3wsWdj0K7qt6nZqUw8jTthCBs,2565
 relationalai/clients/result_helpers.py,sha256=wDSD02Ngx6W-YQqBIGKnpXD4Ju3pA1e9Nz6ORRI6SRI,17808
-relationalai/clients/snowflake.py,sha256=QPJrRolINDkjZb8xyuvaLJIWOEYXdeALLSNqIAAXUEY,166006
+relationalai/clients/snowflake.py,sha256=pB_4t70vTmfGPUJnJpbrrsNH5pojPJyqodsMtUL6b9c,168261
 relationalai/clients/types.py,sha256=eNo6akcMTbnBFbBbHd5IgVeY-zuAgtXlOs8Bo1SWmVU,2890
 relationalai/clients/use_index_poller.py,sha256=rrkg35xiHqY0-2dZlPkgixEGENrIrl7bf_2TboX_qew,46794
 relationalai/clients/util.py,sha256=NJC8fnrWHR01NydwESPSetIHRWf7jQJURYpaWJjmDyE,12311
@@ -314,7 +314,7 @@ relationalai/semantics/lqp/executor.py,sha256=GuPiSJpLaMYqjZOlVEFLOUilyuLcfkQgZ6
 relationalai/semantics/lqp/intrinsics.py,sha256=oKPIcW8PYgU-yPTO21iSF00RBsFKPFFP5MICe6izjKk,871
 relationalai/semantics/lqp/ir.py,sha256=6W9mUH0W7u5eIfF1S3o33uSOfQuM3UcqEkxrxpr1X_8,1867
 relationalai/semantics/lqp/model2lqp.py,sha256=PTe2PtHioWAclu6tdkCu5iGaHJxPIe7zIG1Ew72qxDo,36957
-relationalai/semantics/lqp/passes.py,sha256=i5LrS92Wn-Aa6Hase4UXlogUEGL_tUYHNAYDr16CVw0,28540
+relationalai/semantics/lqp/passes.py,sha256=bpRCdssqcSlVw77tam6EIjI0BypDoFKGWJix8dTywhU,28621
 relationalai/semantics/lqp/pragmas.py,sha256=FzzldrJEAZ1AIcEw6D-FfaVg3CoahRYgPCFo7xHfg1g,375
 relationalai/semantics/lqp/primitives.py,sha256=9Hjow-Yp06jt0xatuUrH1dw0ErnzknIr9K0TB_AwdjU,11029
 relationalai/semantics/lqp/result_helpers.py,sha256=oYpLoTBnzsiyOVIWA2rLMHlgs7P7BoEkqthQ2aMosnk,10123
@@ -322,9 +322,9 @@ relationalai/semantics/lqp/types.py,sha256=3TZ61ybwNV8lDyUMujZIWNFz3Fgn4uifsJb8E
 relationalai/semantics/lqp/utils.py,sha256=iOoS-f8kyFjrgAnpK4cWDvAA-WmPgDRggSKUXm_JdTc,6317
 relationalai/semantics/lqp/validators.py,sha256=FlKMKclHj0L71QUtl0aqKknqksSWM-di4N9bjGDJvnY,1561
 relationalai/semantics/lqp/rewrite/__init__.py,sha256=V9ERED9qdh4VvY9Ud_M8Zn8lhVANdOGIgW03l55sGj0,492
-relationalai/semantics/lqp/rewrite/annotate_constraints.py,sha256=W0_pQYWuq7ocEWNp5rWnq7YjRuJwsEHan1gIrAiRCXU,2201
+relationalai/semantics/lqp/rewrite/annotate_constraints.py,sha256=b_Ly4_80dQpRzWbeLC72JVfxzhwOPBpiCdEqtBiEiwM,2310
 relationalai/semantics/lqp/rewrite/cdc.py,sha256=I6DeMOZScx-3UAVoSCMn9cuOgLzwdvJVKNwsgFa6R_k,10390
-relationalai/semantics/lqp/rewrite/extract_common.py,sha256=sbihURqk4wtc1ekDWXWltq9LrO42XTLfOHl5D6nT5vw,18371
+relationalai/semantics/lqp/rewrite/extract_common.py,sha256=ZRvmeYHN8JEkU-j3fRx1e0_JK-46n6NqhxtwZe6L10c,14690
 relationalai/semantics/lqp/rewrite/extract_keys.py,sha256=iSbwGQG9p8j-erknEwl2pZkunJEpXTlnL9ohr4KVS8M,19317
 relationalai/semantics/lqp/rewrite/function_annotations.py,sha256=9ZzLASvXh_OgQ04eup0AyoMIh2HxWHkoRETLm1-XtWs,4660
 relationalai/semantics/lqp/rewrite/functional_dependencies.py,sha256=4oQcVQtAGDqY850B1bNszigQopf6y9Y_CaUyWx42PtM,12718
@@ -344,9 +344,9 @@ relationalai/semantics/metamodel/util.py,sha256=cmSmeww34JVMqcFudwVAY820IPM2ETSE
 relationalai/semantics/metamodel/visitor.py,sha256=DFY0DACLhxlZ0e4p0vWqbK6ZJr_GWEvH66CU_HVuoTk,35527
 relationalai/semantics/metamodel/rewrite/__init__.py,sha256=9ONWFSdMPHkWpObDMSljt8DywhpFf4Ehsq1aT3fTPt8,344
 relationalai/semantics/metamodel/rewrite/discharge_constraints.py,sha256=0v613BqCLlo4sgWuZjcLSxxakp3d34mYWbG4ldhzGno,1949
-relationalai/semantics/metamodel/rewrite/dnf_union_splitter.py,sha256=ZXX190gCKXhdB-Iyi4MGowc4FS9P0PIJTtTT0LrTr6A,7970
+relationalai/semantics/metamodel/rewrite/dnf_union_splitter.py,sha256=piV8FEr4GHKSDcCcsu_TLHJHrsX7blTq4TUvGpGrN8Q,7986
 relationalai/semantics/metamodel/rewrite/extract_nested_logicals.py,sha256=vQ0-7t_GORskB1ZG50KuzM4phm6YNPvehfFn3v_LbgI,3354
-relationalai/semantics/metamodel/rewrite/flatten.py,sha256=SWvDA3NJqU0ZEJ3kzaXaJ6ZC9xeHuZ8Ju9S1jxU1Nyc,22189
+relationalai/semantics/metamodel/rewrite/flatten.py,sha256=CMCFrMCPIQJs9Ln8TJBJYbZaw6EpIcyWh6N6HaZlLQA,22513
 relationalai/semantics/metamodel/rewrite/format_outputs.py,sha256=n0IxC3RL3UMly6MWsq342EGfL2yGj3vOgVG_wg7kt-o,6225
 relationalai/semantics/metamodel/typer/__init__.py,sha256=E3ydmhWRdm-cAqWsNR24_Qd3NcwiHx8ElO2tzNysAXc,143
 relationalai/semantics/metamodel/typer/checker.py,sha256=frY0gilDO6skbDiYFiIpDUOWyt9s9jAJsRBs848DcG0,19184
@@ -364,7 +364,7 @@ relationalai/semantics/reasoners/optimization/solvers_dev.py,sha256=lbw3c8Z6PlHR
 relationalai/semantics/reasoners/optimization/solvers_pb.py,sha256=ESwraHU9c4NCEVRZ16tnBZsUCmJg7lUhy-v0-GGq0qo,48000
 relationalai/semantics/rel/__init__.py,sha256=pMlVTC_TbQ45mP1LpzwFBBgPxpKc0H3uJDvvDXEWzvs,55
 relationalai/semantics/rel/builtins.py,sha256=kQToiELc4NnvCmXyFtu9CsGZNdTQtSzTB-nuyIfQcsM,1562
-relationalai/semantics/rel/compiler.py,sha256=-gyu_kL1p4Z-SXI2fwLXSbGfUy-ejQcKVrc66rVvOyg,43044
+relationalai/semantics/rel/compiler.py,sha256=pFkEbuPKVd8AI4tiklcv06LbNnK8KfoV4FwmY9Lrhqo,43044
 relationalai/semantics/rel/executor.py,sha256=v-yHl9R8AV0AA2xnm5YZDzue83pr8j2Q97Ky1MKkU70,17309
 relationalai/semantics/rel/rel.py,sha256=9I_V6dQ83QRaLzq04Tt-KjBWhmNxNO3tFzeornBK4zc,15738
 relationalai/semantics/rel/rel_utils.py,sha256=EH-NBROA4vIJXajLKniapt4Dxt7cXSqY4NEjD-wD8Mc,9566
@@ -442,8 +442,8 @@ frontend/debugger/dist/index.html,sha256=0wIQ1Pm7BclVV1wna6Mj8OmgU73B9rSEGPVX-Wo
 frontend/debugger/dist/assets/favicon-Dy0ZgA6N.png,sha256=tPXOEhOrM4tJyZVJQVBO_yFgNAlgooY38ZsjyrFstgg,620
 frontend/debugger/dist/assets/index-Cssla-O7.js,sha256=MxgIGfdKQyBWgufck1xYggQNhW5nj6BPjCF6Wleo-f0,298886
 frontend/debugger/dist/assets/index-DlHsYx1V.css,sha256=21pZtAjKCcHLFjbjfBQTF6y7QmOic-4FYaKNmwdNZVE,60141
-relationalai-0.12.11.dist-info/METADATA,sha256=Gaq2yLT7R1yQ5rf6a7IwR1fWRaFaZ-sMLV5cMdQmeWg,2563
-relationalai-0.12.11.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-relationalai-0.12.11.dist-info/entry_points.txt,sha256=fo_oLFJih3PUgYuHXsk7RnCjBm9cqRNR--ab6DgI6-0,88
-relationalai-0.12.11.dist-info/licenses/LICENSE,sha256=pPyTVXFYhirkEW9VsnHIgUjT0Vg8_xsE6olrF5SIgpc,11343
-relationalai-0.12.11.dist-info/RECORD,,
+relationalai-0.12.13.dist-info/METADATA,sha256=qiiOC2n_SXtlfEHafFp8hVqY1M2VSzGOl5fqOPt1_ZQ,2563
+relationalai-0.12.13.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+relationalai-0.12.13.dist-info/entry_points.txt,sha256=fo_oLFJih3PUgYuHXsk7RnCjBm9cqRNR--ab6DgI6-0,88
+relationalai-0.12.13.dist-info/licenses/LICENSE,sha256=pPyTVXFYhirkEW9VsnHIgUjT0Vg8_xsE6olrF5SIgpc,11343
+relationalai-0.12.13.dist-info/RECORD,,