relationalai 0.12.9__py3-none-any.whl → 0.12.10__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_auto_create: bool = False
     ) -> Any:
         # print(f"\n--- sql---\n{code}\n--- end sql---\n")
         if not self._session:
@@ -458,7 +459,6 @@ class Resources(ResourcesBase):
             rai_app = self.config.get("rai_app_name", "")
             current_role = self.config.get("role")
             engine = self.get_default_engine_name()
-            engine_size = self.config.get_default_engine_size()
             assert isinstance(rai_app, str), f"rai_app_name must be a string, not {type(rai_app)}"
             assert isinstance(engine, str), f"engine must be a string, not {type(engine)}"
             print("\n")
@@ -467,15 +467,10 @@ 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 creating the engine if the query is a user transaction. exec_async_v2 will handle that case.
+            if _is_engine_issue(orig_message) and not skip_auto_create:
                 try:
-                    self._poll_use_index(
-                        app_name=self.get_app_name(),
-                        sources=self.sources,
-                        model=self.database,
-                        engine_name=engine,
-                        engine_size=engine_size
-                    )
+                    self.auto_create_engine(engine)
                     return self._exec(code, params, raw=raw, help=help)
                 except EngineNameValidationException as e:
                     raise EngineNameValidationException(engine) from e
@@ -1612,6 +1607,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
         response = self._exec(
             sql_string,
             raw_code,
+            skip_auto_create=True,
         )
         if not response:
             raise Exception("Failed to create transaction")
@@ -1629,6 +1625,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 +1635,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 +1896,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 +1935,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 +1976,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]]:
@@ -3314,19 +3300,10 @@ 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):
-                        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()
-                        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,
-                        )
+                    # Skip auto-retry if skip_auto_create is True to avoid recursion or to let _exec_async_v2 poll the index.
+                    if _is_engine_issue(message) and not skip_auto_create:
+                        engine = payload.get("engine_name", "") if payload else ""
+                        self.auto_create_engine(engine)
                         response = _send_request()
             except requests.exceptions.ConnectionError as e:
                 if "NameResolutionError" in str(e):
@@ -3340,6 +3317,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
+        )
+
+        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
+                )
+            else:
+                raise ResponseStatusException("Failed to create transaction.", response)
+
+        return response
+
     def _exec_async_v2(
         self,
         database: str,
@@ -3352,6 +3371,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 +3394,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/debugging.py

@@ -26,6 +26,9 @@ find_block_in = find_block_in # re-export
 DEBUG = True
 handled_error = None
 
+# Configurable debug log file location
+DEBUG_LOG_FILE = os.environ.get('RAI_DEBUG_LOG', 'debug.jsonl')
+
 #--------------------------------------------------
 # Log Formatters
 #--------------------------------------------------
@@ -70,7 +73,7 @@ class FlushingFileHandler(logging.FileHandler):
     def emit(self, record):
         if not self._initialized:
             self._initialized = True
-            with open('debug.jsonl', 'w'):
+            with open(DEBUG_LOG_FILE, 'w'):
                 pass
         super().emit(record)
         self.flush()
@@ -78,7 +81,7 @@ class FlushingFileHandler(logging.FileHandler):
 try:
     # keep the old file-based debugger around and working until it's fully replaced.
     if DEBUG:
-        file_handler = FlushingFileHandler('debug.jsonl', mode='a')
+        file_handler = FlushingFileHandler(DEBUG_LOG_FILE, mode='a')
         file_handler.setFormatter(JsonFormatter())
         logger.addHandler(file_handler)
 except Exception:

relationalai/semantics/__init__.py

@@ -8,7 +8,7 @@ __include_in_docs__ = True
 
 from .internal import (
     Model, Concept, Relationship, RelationshipReading, Expression, Fragment, Error, Field,
-    String, Integer, Int64, Int128, Float, Decimal, Bool,
+    AnyEntity, String, Integer, Int64, Int128, Float, Decimal, Bool,
     Date, DateTime,
     RawSource, Hash,
     select, where, require, define, distinct, union, data,
@@ -19,7 +19,7 @@ from .internal import (
 
 __all__ = [
     "Model", "Concept", "Relationship", "RelationshipReading", "Expression", "Fragment", "Error", "Field",
-    "String", "Integer", "Int64", "Int128", "Float", "Decimal", "Bool",
+    "AnyEntity", "String", "Integer", "Int64", "Int128", "Float", "Decimal", "Bool",
     "Date", "DateTime",
     "RawSource", "Hash",
     "select", "where", "require", "define", "distinct", "union", "data",

relationalai/semantics/internal/__init__.py

@@ -4,7 +4,7 @@ API for RelationalAI.
 
 from .internal import (
     Model, Concept, Relationship, RelationshipReading, Expression, Fragment, Error, Field,
-    String, Integer, Int64, Int128, Float, Decimal, Bool,
+    AnyEntity, String, Integer, Int64, Int128, Float, Decimal, Bool,
     Date, DateTime,
     RawSource, Hash,
     select, where, require, define, distinct, union, data,
@@ -15,7 +15,7 @@ from .internal import (
 
 __all__ = [
     "Model", "Concept", "Relationship", "RelationshipReading", "Expression", "Fragment", "Error", "Field",
-    "String", "Integer", "Int64", "Int128", "Float", "Decimal", "Bool",
+    "AnyEntity", "String", "Integer", "Int64", "Int128", "Float", "Decimal", "Bool",
     "Date", "DateTime",
     "RawSource", "Hash",
     "select", "where", "require", "define", "distinct", "union", "data",

relationalai/semantics/internal/internal.py

@@ -621,7 +621,7 @@ class Producer:
 
                 if self._model and self._model._strict:
                     raise AttributeError(f"{self._name} has no relationship `{name}`")
-                if topmost_parent is not concept:
+                if topmost_parent is not concept and topmost_parent not in Concept.builtin_concepts:
                     topmost_parent._relationships[name] = topmost_parent._get_relationship(name)
                     rich.print(f"[red bold][Implicit Subtype Relationship][/red bold] [yellow]{concept}.{name}[/yellow] appended to topmost parent [yellow]{topmost_parent}[/yellow] instead")
 
@@ -1165,7 +1165,10 @@ Primitive = Concept.builtins["Primitive"] = Concept("Primitive")
 Error = Concept.builtins["Error"] = ErrorConcept("Error")
 
 def _register_builtin(name):
-    c = Concept(name, extends=[Primitive])
+    if name == "AnyEntity":
+        c = Concept(name)
+    else:
+        c = Concept(name, extends=[Primitive])
     Concept.builtin_concepts.add(c)
     Concept.builtins[name] = c
 
@@ -1174,6 +1177,7 @@ for builtin in types.builtin_types:
     if isinstance(builtin, ir.ScalarType):
         _register_builtin(builtin.name)
 
+AnyEntity = Concept.builtins["AnyEntity"]
 Float = Concept.builtins["Float"]
 Number = Concept.builtins["Number"]
 Int64 = Concept.builtins["Int64"]
@@ -2896,10 +2900,9 @@ class Compiler():
                 if concept not in self.types:
                     self.to_type(concept)
                     self.to_relation(concept)
-                if concept._extends:
-                    rule = self.concept_inheritance_rule(concept)
-                    if rule:
-                        rules.append(rule)
+                rule = self.concept_inheritance_rule(concept)
+                if rule:
+                    rules.append(rule)
         unresolved = []
         for relationship in model.relationships:
             if relationship not in self.relations:
@@ -3204,8 +3207,11 @@ class Compiler():
         # filter extends to get only non-primitive parents
         parents = []
         for parent in concept._extends:
-            if not parent._is_primitive():
+            if not parent._is_primitive() and parent is not AnyEntity:
                 parents.append(parent)
+        # always extend AnyEntity for non-primitive types that are not built-in
+        if not concept._is_primitive() and concept not in Concept.builtin_concepts:
+            parents.append(AnyEntity)
         # only extends primitive types, no need for inheritance rules
         if not parents:
             return None
@@ -3218,6 +3224,17 @@ class Compiler():
             *[f.derive(self.to_relation(parent), [var]) for parent in parents]
         ])
 
+    def concept_any_entity_rule(self, entities:list[Concept]):
+        """
+        Generate an inheritance rule for all these entities to AnyEntity.
+        """
+        any_entity_relation = self.to_relation(AnyEntity)
+        var = f.var("v", types.Any)
+        return f.logical([
+            f.union([f.lookup(self.to_relation(e), [var]) for e in entities]),
+            f.derive(any_entity_relation, [var])
+        ])
+
     def relation_dict(self, items:dict[Relationship|Concept, Producer], ctx:CompilerContext) -> dict[ir.Relation, list[ir.Var]]:
         return {self.to_relation(k): unwrap_list(self.lookup(v, ctx)) for k, v in items.items()}
 

relationalai/semantics/lqp/README.md

@@ -0,0 +1,34 @@
+# Logic Engine LQP Backend
+
+The logic engine runs the *Logical Query Protocol* (short *LQP*). This module includes a
+compiler from the semantic metamodel to LQP along with an executor.
+
+## Running against a local logic engine
+
+For development and testing, it is possible to run PyRel models against a local logic engine
+server process.
+
+To start your local server, please refer to the [logic engine
+docs](https://github.com/RelationalAI/raicode/tree/master/src/Server#starting-the-server).
+
+With the local server running, add this to your `raiconfig.toml`:
+
+```toml
+[profile.local]
+platform = "local"
+engine = "local"
+host = "localhost"
+port = 8010
+```
+
+Then set `active_profile = "local"` at the top of the file.
+
+**Known limitations:**
+
+Local execution does not support running against Snowflake source tables.
+
+At the moment, locally created databases cannot be cleaned up by the client. Eventually you
+will need to clear your local pager directory.
+
+At the moment, local execution is only supported for fast-path transactions, i.e. those
+which complete in less than 5 seconds. Polling support will be added soon.

relationalai/semantics/lqp/constructors.py

@@ -63,5 +63,6 @@ def mk_attribute(name: str, args: list[lqp.Value]) -> lqp.Attribute:
 def mk_transaction(
     epochs: list[lqp.Epoch],
     configure: lqp.Configure = lqp.construct_configure({}, None),
+    sync = None
 ) -> lqp.Transaction:
-    return lqp.Transaction(epochs=epochs, configure=configure, meta=None)
+    return lqp.Transaction(epochs=epochs, configure=configure, sync=sync, meta=None)

relationalai/semantics/lqp/executor.py

@@ -66,6 +66,8 @@ class LQPExecutor(e.Executor):
                 resource_class = rai.clients.snowflake.Resources
                 if self.config.get("use_direct_access", USE_DIRECT_ACCESS):
                     resource_class = rai.clients.snowflake.DirectAccessResources
+                if self.config.get("platform", "") == "local":
+                    resource_class = rai.clients.local.LocalResources
                 # NOTE: language="lqp" is not strictly required for LQP execution, but it
                 # will significantly improve performance.
                 self._resources = resource_class(
@@ -311,6 +313,12 @@ class LQPExecutor(e.Executor):
             config_dict['ivm.maintenance_level'] = lqp_ir.Value(value=ivm_flag, meta=None)
         return construct_configure(config_dict, None)
 
+    def _should_sync(self, model) :
+        if self._last_model != model:
+            return lqp_ir.Sync(fragments=[], meta=None)
+        else :
+            return None
+
     def _compile_intrinsics(self) -> lqp_ir.Epoch:
         """Construct an epoch that defines a number of built-in definitions used by the
         emitter."""
@@ -334,6 +342,7 @@ class LQPExecutor(e.Executor):
                 meta=None,
             )
 
+    # [RAI-40997] We eagerly undefine query fragments so they are not committed to storage
     def _compile_undefine_query(self, query_epoch: lqp_ir.Epoch) -> lqp_ir.Epoch:
         fragment_ids = []
 
@@ -363,7 +372,9 @@ class LQPExecutor(e.Executor):
         epochs = []
         epochs.append(self._compile_intrinsics())
 
-        if self._last_model != model:
+        sync = self._should_sync(model)
+
+        if sync is not None:
             with debugging.span("compile", metamodel=model) as install_span:
                 install_span["compile_type"] = "model"
                 _, model_epoch = self.compiler.compile(model, {"fragment_id": b"model"})
@@ -383,7 +394,7 @@ class LQPExecutor(e.Executor):
             epochs.append(self._compile_undefine_query(query_epoch))
 
             txn_span["compile_type"] = "query"
-            txn = mk_transaction(epochs=epochs, configure=configure)
+            txn = mk_transaction(epochs=epochs, configure=configure, sync=sync)
             txn_span["lqp"] = lqp_print.to_string(txn, {"print_names": True, "print_debug": False, "print_csv_filename": False})
 
         validate_lqp(txn)

relationalai/semantics/lqp/ir.py

@@ -4,6 +4,7 @@ __all__ = [
     "SourceInfo",
     "LqpNode",
     "Declaration",
+    "FunctionalDependency",
     "Def",
     "Loop",
     "Abstraction",
@@ -45,6 +46,7 @@ __all__ = [
     "Read",
     "Epoch",
     "Transaction",
+    "Sync",
     "DebugInfo",
     "Configure",
     "IVMConfig",
@@ -59,6 +61,7 @@ from lqp.ir import (
     SourceInfo,
     LqpNode,
     Declaration,
+    FunctionalDependency,
     Def,
     Loop,
     Abstraction,
@@ -100,6 +103,7 @@ from lqp.ir import (
     Read,
     Epoch,
     Transaction,
+    Sync,
     DebugInfo,
     Configure,
     IVMConfig,

relationalai/semantics/lqp/model2lqp.py

@@ -11,7 +11,9 @@ from relationalai.semantics.lqp.constructors import (
 )
 from relationalai.semantics.lqp.utils import TranslationCtx, gen_unique_var
 from relationalai.semantics.lqp.validators import assert_valid_input
-
+from relationalai.semantics.lqp.rewrite.functional_dependencies import (
+    normalized_fd, contains_only_declarable_constraints
+)
 from decimal import Decimal as PyDecimal
 from datetime import datetime, date, timezone
 from typing import Tuple, cast, Union, Optional
@@ -102,6 +104,43 @@ def _get_export_reads(export_ids: list[tuple[lqp.RelationId, int, lqp.Type]]) ->
     return (export_filename, col_info, reads)
 
 def _translate_to_decls(ctx: TranslationCtx, rule: ir.Logical) -> list[lqp.Declaration]:
+    if contains_only_declarable_constraints(rule):
+        return _translate_to_constraint_decls(ctx, rule)
+    else:
+        return _translate_to_standard_decl(ctx, rule)
+
+def _translate_to_constraint_decls(ctx: TranslationCtx, rule: ir.Logical) -> list[lqp.Declaration]:
+    constraint_decls: list[lqp.Declaration] = []
+    for task in rule.body:
+        assert isinstance(task, ir.Require)
+        fd = normalized_fd(task)
+        assert fd is not None
+
+        # check for unresolved types
+        if any(types.is_any(var.type) for var in fd.keys + fd.values):
+            warn(f"Ignoring FD with unresolved type: {fd}")
+            continue
+
+        lqp_typed_keys = [_translate_term(ctx, key) for key in fd.keys]
+        lqp_typed_values = [_translate_term(ctx, value) for value in fd.values]
+        lqp_typed_vars:list[Tuple[lqp.Var, lqp.Type]] = lqp_typed_keys + lqp_typed_values # type: ignore
+        lqp_guard_atoms = [_translate_to_atom(ctx, atom) for atom in fd.guard]
+        lqp_guard = mk_abstraction(lqp_typed_vars, mk_and(lqp_guard_atoms))
+        lqp_keys:list[lqp.Var] = [var for (var, _) in lqp_typed_keys] # type: ignore
+        lqp_values:list[lqp.Var] = [var for (var, _) in lqp_typed_values] # type: ignore
+
+        fd_decl = lqp.FunctionalDependency(
+            guard=lqp_guard,
+            keys=lqp_keys,
+            values=lqp_values,
+            meta=None
+        )
+
+        constraint_decls.append(fd_decl)
+
+    return constraint_decls
+
+def _translate_to_standard_decl(ctx: TranslationCtx, rule: ir.Logical) -> list[lqp.Declaration]:
     effects = collect_by_type((ir.Output, ir.Update), rule)
     aggregates = collect_by_type(ir.Aggregate, rule)
     ranks = collect_by_type(ir.Rank, rule)
@@ -452,7 +491,7 @@ def _translate_aggregate(ctx: TranslationCtx, aggr: ir.Aggregate, body: lqp.Form
 
         return mk_exists(result_terms, conjunction)
 
-    # `input_args`` hold the types of the input arguments, but they may have been modified
+    # `input_args` hold the types of the input arguments, but they may have been modified
     # if we're dealing with a count, so we use `abstr_args` to find the type.
     (aggr_arg, aggr_arg_type) = abstr_args[-1]
 

relationalai/semantics/lqp/passes.py

@@ -6,9 +6,11 @@ from relationalai.semantics.metamodel.util import FrozenOrderedSet
 
 from relationalai.semantics.metamodel.rewrite import Flatten
 
-from ..metamodel.rewrite import DischargeConstraints, DNFUnionSplitter, ExtractNestedLogicals, FormatOutputs
-from .rewrite import CDC, ExtractCommon, ExtractKeys, FunctionAnnotations, QuantifyVars, Splinter, SplitMultiCheckRequires
-
+from ..metamodel.rewrite import DNFUnionSplitter, ExtractNestedLogicals, FormatOutputs
+from .rewrite import (
+    AnnotateConstraints, CDC, ExtractCommon, ExtractKeys, FunctionAnnotations, QuantifyVars,
+    Splinter, SplitMultiCheckRequires
+)
 from relationalai.semantics.lqp.utils import output_names
 
 from typing import cast, List, Sequence, Tuple, Union, Optional, Iterable
@@ -20,7 +22,7 @@ def lqp_passes() -> list[Pass]:
     return [
         SplitMultiCheckRequires(),
         FunctionAnnotations(),
-        DischargeConstraints(),
+        AnnotateConstraints(),
         Checker(),
         CDC(), # specialize to physical relations before extracting nested and typing
         ExtractNestedLogicals(), # before InferTypes to avoid extracting casts

relationalai/semantics/lqp/rewrite/__init__.py

@@ -1,3 +1,4 @@
+from .annotate_constraints import AnnotateConstraints
 from .cdc import CDC
 from .extract_common import ExtractCommon
 from .extract_keys import ExtractKeys
@@ -6,6 +7,7 @@ from .quantify_vars import QuantifyVars
 from .splinter import Splinter
 
 __all__ = [
+    "AnnotateConstraints",
     "CDC",
     "ExtractCommon",
     "ExtractKeys",