relationalai 0.12.6__py3-none-any.whl → 0.12.8__py3-none-any.whl

relationalai/clients/snowflake.py

@@ -851,7 +851,17 @@ Otherwise, remove it from your '{profile}' configuration profile.
                 self.generation
             )
             # If cache is valid (data freshness has not expired), skip polling
-            if not poller.cache.is_valid():
+            if poller.cache.is_valid():
+                cached_sources = len(poller.cache.sources)
+                total_sources = len(sources_list)
+                cached_timestamp = poller.cache._metadata.get("cachedIndices", {}).get(poller.cache.key, {}).get("last_use_index_update_on", "")
+
+                message = f"Using cached data for {cached_sources}/{total_sources} data streams"
+                if cached_timestamp:
+                    print(f"\n{message} (cached at {cached_timestamp})\n")
+                else:
+                    print(f"\n{message}\n")
+            else:
                 return poller.poll()
 
     #--------------------------------------------------
@@ -2842,6 +2852,15 @@ class SnowflakeTable(dsl.Type):
                 else:
                     self(snowflake_id=id).set(**{prop_ident: val})
 
+        # Because we're bypassing a bunch of the normal Type.add machinery here,
+        # we need to manually account for the case where people are using value types.
+        def wrapped(x):
+            if not model._config.get("compiler.use_value_types", False):
+                return x
+            other_id = dsl.create_var()
+            model._action(dsl.build.construct(self._type, [x, other_id]))
+            return other_id
+
         # new UInt128 schema mapping rules
         with model.rule(dynamic=True, globalize=True, source=self._source):
             id = dsl.create_var()
@@ -2851,7 +2870,7 @@ class SnowflakeTable(dsl.Type):
             # for avoiding a non-blocking warning
             edb(dsl.Symbol("METADATA$KEY"), id)
             std.rel.UInt128(id)
-            self.add(id, snowflake_id=id)
+            self.add(wrapped(id), snowflake_id=id)
 
         for prop, prop_type in self._schema["columns"].items():
             _prop = prop
@@ -2873,7 +2892,7 @@ class SnowflakeTable(dsl.Type):
                     model._check_property(_prop._prop)
                     raw_relation = getattr(std.rel, prop_ident)
                     dsl.tag(raw_relation, dsl.Builtins.FunctionAnnotation)
-                    raw_relation.add(id, val)
+                    raw_relation.add(wrapped(id), val)
 
     def namespace(self):
         return f"{self._parent._parent._name}.{self._parent._name}"
@@ -3275,12 +3294,24 @@ class DirectAccessResources(Resources):
             try:
                 response = _send_request()
                 if response.status_code != 200:
+                    # For 404 responses with skip_auto_create=True, return immediately to let caller handle it
+                    # (e.g., get_engine needs to check 404 and return None for auto_create_engine)
+                    # For skip_auto_create=False, continue to auto-creation logic below
+                    if response.status_code == 404 and skip_auto_create:
+                        return response
+
                     try:
                         message = response.json().get("message", "")
                     except requests.exceptions.JSONDecodeError:
-                        raise ResponseStatusException(
-                            f"Failed to parse error response from endpoint {endpoint}.", response
-                        )
+                        # Can't parse JSON response. For skip_auto_create=True (e.g., get_engine),
+                        # this should have been caught by the 404 check above, so this is an error.
+                        # For skip_auto_create=False, we explicitly check status_code below,
+                        # so we don't need to parse the message.
+                        if skip_auto_create:
+                            raise ResponseStatusException(
+                                f"Failed to parse error response from endpoint {endpoint}.", response
+                            )
+                        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
@@ -3473,7 +3504,17 @@ class DirectAccessResources(Resources):
                 generation=self.generation,
             )
             # If cache is valid (data freshness has not expired), skip polling
-            if not poller.cache.is_valid():
+            if poller.cache.is_valid():
+                cached_sources = len(poller.cache.sources)
+                total_sources = len(sources_list)
+                cached_timestamp = poller.cache._metadata.get("cachedIndices", {}).get(poller.cache.key, {}).get("last_use_index_update_on", "")
+
+                message = f"Using cached data for {cached_sources}/{total_sources} data streams"
+                if cached_timestamp:
+                    print(f"\n{message} (cached at {cached_timestamp})\n")
+                else:
+                    print(f"\n{message}\n")
+            else:
                 return poller.poll()
 
     def _check_exec_async_status(self, txn_id: str, headers: Dict[str, str] | None = None) -> bool:

relationalai/clients/use_index_poller.py

@@ -250,7 +250,17 @@ class UseIndexPoller:
             # Cache was used - show how many sources were cached
             total_sources = len(self.cache.sources)
             cached_sources = total_sources - len(self.sources)
-            progress.add_sub_task(f"Using cached data for {cached_sources}/{total_sources} data streams", task_id="cache_usage", category=TASK_CATEGORY_CACHE)
+
+            # Get the timestamp when sources were cached
+            entry = self.cache._metadata.get("cachedIndices", {}).get(self.cache.key, {})
+            cached_timestamp = entry.get("last_use_index_update_on", "")
+
+            message = f"Using cached data for {cached_sources}/{total_sources} data streams"
+            # Format the message with timestamp
+            if cached_timestamp:
+                message += f" (cached at {cached_timestamp})"
+
+            progress.add_sub_task(message, task_id="cache_usage", category=TASK_CATEGORY_CACHE)
             # Complete the subtask immediately since it's just informational
             progress.complete_sub_task("cache_usage")
 

relationalai/early_access/lqp/constructors/__init__.py

@@ -2,12 +2,12 @@ import warnings
 
 from relationalai.semantics.lqp.constructors import (
     mk_abstraction, mk_and, mk_exists, mk_or, mk_pragma, mk_primitive,
-    mk_specialized_value, mk_type, mk_value, mk_var
+    mk_specialized_value, mk_type, mk_value
 )
 
 __all__ = [
     "mk_abstraction", "mk_and", "mk_exists", "mk_or", "mk_pragma", "mk_primitive", "mk_specialized_value", "mk_type",
-    "mk_value", "mk_var"
+    "mk_value"
 ]
 
 warnings.warn(

relationalai/early_access/metamodel/rewrite/__init__.py

@@ -1,7 +1,7 @@
 from relationalai.semantics.metamodel.rewrite import Flatten, \
     DNFUnionSplitter, ExtractNestedLogicals, flatten
 from relationalai.semantics.lqp.rewrite import Splinter,  \
-    ExtractKeys, FDConstraints
+    ExtractKeys, FunctionAnnotations
 
 __all__ = ["Splinter", "Flatten", "DNFUnionSplitter", "ExtractKeys",
-           "ExtractNestedLogicals", "FDConstraints", "flatten"]
+           "ExtractNestedLogicals", "FunctionAnnotations", "flatten"]

relationalai/semantics/internal/internal.py

@@ -2550,8 +2550,6 @@ class Fragment():
         from .snowflake import Table
         assert isinstance(table, Table), "Only Snowflake tables are supported for now"
 
-        result_cols = table._col_names
-
         clone = Fragment(parent=self)
         clone._is_export = True
         qb_model = clone._model or Model("anon")
@@ -2559,8 +2557,7 @@ class Fragment():
         clone._source = runtime_env.get_source_pos()
         with debugging.span("query", dsl=str(clone), **with_source(clone), meta=clone._meta):
             query_task = qb_model._compiler.fragment(clone)
-            qb_model._to_executor().execute(ir_model, query_task, result_cols=result_cols, export_to=table._fqn, update=update, meta=clone._meta)
-
+            qb_model._to_executor().execute(ir_model, query_task, export_to=table, update=update, meta=clone._meta)
 
 #--------------------------------------------------
 # Select / Where

relationalai/semantics/internal/snowflake.py

@@ -12,7 +12,18 @@ from . import internal as b, annotations as anns
 from relationalai import debugging
 from relationalai.errors import UnsupportedColumnTypesWarning
 from snowflake.snowpark.context import get_active_session
+from typing import ClassVar, Optional
 
+#--------------------------------------------------
+# Iceberg Configuration
+#--------------------------------------------------
+@dataclass
+class IcebergConfig:
+    """Configuration for exporting to Iceberg tables."""
+    external_volume: str | None = None
+    default: ClassVar[Optional["IcebergConfig"]]
+
+IcebergConfig.default = IcebergConfig()
 #--------------------------------------------------
 # Helpers
 #--------------------------------------------------
@@ -191,7 +202,7 @@ class Table():
     _schemas:dict[tuple[str, str], SchemaInfo] = {}
     _used_sources:OrderedSet[Table] = ordered_set()
 
-    def __init__(self, fqn:str, cols:list[str]|None=None, schema:dict[str, str|b.Concept]|None=None) -> None:
+    def __init__(self, fqn:str, cols:list[str]|None=None, schema:dict[str, str|b.Concept]|None=None, config: IcebergConfig|None=None) -> None:
         self._fqn = fqn
         parser = IdentityParser(fqn, require_all_parts=True)
         self._database, self._schema, self._table, self._fqn = parser.to_list()
@@ -201,6 +212,8 @@ class Table():
         self._ref = self._concept.ref("row_id")
         self._cols = {}
         self._col_names = cols
+        self._iceberg_config = config
+        self._is_iceberg = config is not None
         info = self._schemas.get((self._database, self._schema))
         if not info:
             info = self._schemas[(self._database, self._schema)] = SchemaInfo(self._database, self._schema)

relationalai/semantics/lqp/constructors.py

@@ -1,6 +1,5 @@
 from typing import Tuple
 from relationalai.semantics.lqp import ir as lqp
-from relationalai.semantics.metamodel.ir import sanitize
 
 def mk_and(args: list[lqp.Formula]) -> lqp.Formula:
     # Flatten nested conjunctions
@@ -49,10 +48,6 @@ def mk_specialized_value(value) -> lqp.SpecializedValue:
 def mk_value(value) -> lqp.Value:
     return lqp.Value(value=value, meta=None)
 
-def mk_var(name: str) -> lqp.Var:
-    _name = '_' if name == '_' else sanitize(name)
-    return lqp.Var(name=_name, meta=None)
-
 def mk_type(typename: lqp.TypeName, parameters: list[lqp.Value]=[]) -> lqp.Type:
     return lqp.Type(type_name=typename, parameters=parameters, meta=None)
 

relationalai/semantics/lqp/executor.py

@@ -4,7 +4,7 @@ import atexit
 import re
 
 from pandas import DataFrame
-from typing import Any, Optional, Literal
+from typing import Any, Optional, Literal, TYPE_CHECKING
 from snowflake.snowpark import Session
 import relationalai as rai
 
@@ -20,10 +20,14 @@ from relationalai.semantics.lqp.ir import convert_transaction, validate_lqp
 from relationalai.clients.config import Config
 from relationalai.clients.snowflake import APP_NAME
 from relationalai.clients.types import TransactionAsyncResponse
-from relationalai.clients.util import IdentityParser
+from relationalai.clients.util import IdentityParser, escape_for_f_string
 from relationalai.tools.constants import USE_DIRECT_ACCESS, QUERY_ATTRIBUTES_HEADER
 from relationalai.tools.query_utils import prepare_metadata_for_headers
 
+if TYPE_CHECKING:
+    from relationalai.semantics.snowflake import Table
+
+
 class LQPExecutor(e.Executor):
     """Executes LQP using the RAI client."""
 
@@ -172,12 +176,12 @@ class LQPExecutor(e.Executor):
         elif len(all_errors) > 1:
             raise errors.RAIExceptionSet(all_errors)
 
-    def _export(self, txn_id: str, export_info: tuple, dest_fqn: str, actual_cols: list[str], declared_cols: list[str], update:bool):
+    def _export(self, txn_id: str, export_info: tuple, dest: Table, actual_cols: list[str], declared_cols: list[str], update: bool):
         # At this point of the export, we assume that a CSV file has already been written
         # to the Snowflake Native App stage area. Thus, the purpose of this method is to
         # copy the data from the CSV file to the destination table.
         _exec = self.resources._exec
-        dest_database, dest_schema, dest_table, _ = IdentityParser(dest_fqn, require_all_parts=True).to_list()
+        dest_database, dest_schema, dest_table, _ = IdentityParser(dest._fqn, require_all_parts=True).to_list()
         filename = export_info[0]
         result_table_name = filename + "_table"
 
@@ -203,8 +207,28 @@ class LQPExecutor(e.Executor):
                 # destination table. This step also cleans up the result table.
                 out_sample = _exec(f"select * from {APP_NAME}.results.{result_table_name} limit 1;")
                 names = self._build_projection(declared_cols, actual_cols, column_fields, out_sample)
+                dest_fqn = dest._fqn
                 try:
                     if not update:
+                        createTableLogic = f"""
+                                        CREATE TABLE {dest_fqn} AS
+                                        SELECT {names}
+                                        FROM {APP_NAME}.results.{result_table_name};
+                        """
+                        if dest._is_iceberg:
+                            assert dest._iceberg_config is not None
+                            external_volume_clause = ""
+                            if dest._iceberg_config.external_volume:
+                                external_volume_clause = f"EXTERNAL_VOLUME = '{dest._iceberg_config.external_volume}'"
+                            createTableLogic = f"""
+                                        CREATE ICEBERG TABLE {dest_fqn}
+                                        CATALOG = "SNOWFLAKE"
+                                        {external_volume_clause}
+                                        AS
+                                        SELECT {names}
+                                        FROM {APP_NAME}.results.{result_table_name};
+                            """
+
                         _exec(f"""
                             BEGIN
                                 -- Check if table exists
@@ -227,9 +251,7 @@ class LQPExecutor(e.Executor):
                                 ELSE
                                     -- Create table based on the SELECT
                                     EXECUTE IMMEDIATE '
-                                        CREATE TABLE {dest_fqn} AS
-                                        SELECT {names}
-                                        FROM {APP_NAME}.results.{result_table_name};
+                                        {escape_for_f_string(createTableLogic)}
                                     ';
                                 END IF;
                             END;
@@ -376,7 +398,7 @@ class LQPExecutor(e.Executor):
         return final_model, export_info, txn_proto
 
     # TODO (azreika): This should probably be split up into exporting and other processing. There are quite a lot of arguments here...
-    def _process_results(self, task: ir.Task, final_model: ir.Model, raw_results: TransactionAsyncResponse, result_cols: Optional[list[str]], export_info: Optional[tuple], export_to: Optional[str], update: bool) -> DataFrame:
+    def _process_results(self, task: ir.Task, final_model: ir.Model, raw_results: TransactionAsyncResponse, export_info: Optional[tuple], export_to: Optional[Table], update: bool) -> DataFrame:
         cols, extra_cols = self._compute_cols(task, final_model)
 
         df, errs = result_helpers.format_results(raw_results, cols)
@@ -391,6 +413,8 @@ class LQPExecutor(e.Executor):
             assert cols, "No columns found in the output"
             assert isinstance(raw_results, TransactionAsyncResponse) and raw_results.transaction, "Invalid transaction result"
 
+            result_cols = export_to._col_names
+
             if result_cols is not None:
                 assert all(col in result_cols or col in extra_cols for col in cols)
             else:
@@ -403,7 +427,7 @@ class LQPExecutor(e.Executor):
         return self._postprocess_df(self.config, df, extra_cols)
 
     def execute(self, model: ir.Model, task: ir.Task, format: Literal["pandas", "snowpark"] = "pandas",
-                result_cols: Optional[list[str]] = None, export_to: Optional[str] = None,
+                export_to: Optional[Table] = None,
                 update: bool = False, meta: dict[str, Any] | None = None) -> DataFrame:
         self.prepare_data()
         previous_model = self._last_model
@@ -433,7 +457,7 @@ class LQPExecutor(e.Executor):
         assert isinstance(raw_results, TransactionAsyncResponse)
 
         try:
-            return self._process_results(task, final_model, raw_results, result_cols, export_info, export_to, update)
+            return self._process_results(task, final_model, raw_results, export_info, export_to, update)
         except Exception as e:
             # If processing the results failed, revert to the previous model.
             self._last_model = previous_model

relationalai/semantics/lqp/intrinsics.py

@@ -1,13 +1,13 @@
 from datetime import datetime, timezone
 
 from relationalai.semantics.lqp import ir as lqp
-from relationalai.semantics.lqp.constructors import mk_abstraction, mk_value, mk_var, mk_type, mk_primitive
+from relationalai.semantics.lqp.constructors import mk_abstraction, mk_value, mk_type, mk_primitive
 from relationalai.semantics.lqp.utils import lqp_hash
 
 def mk_intrinsic_datetime_now() -> lqp.Def:
     """Constructs a definition of the current datetime."""
     id = lqp_hash("__pyrel_lqp_intrinsic_datetime_now")
-    out = mk_var("out")
+    out = lqp.Var(name="out", meta=None)
     out_type = mk_type(lqp.TypeName.DATETIME)
     now = mk_value(lqp.DateTimeValue(value=datetime.now(timezone.utc), meta=None))
     datetime_now = mk_abstraction(

relationalai/semantics/lqp/model2lqp.py

@@ -7,7 +7,7 @@ from relationalai.semantics.lqp.pragmas import pragma_to_lqp_name
 from relationalai.semantics.lqp.types import meta_type_to_lqp
 from relationalai.semantics.lqp.constructors import (
     mk_abstraction, mk_and, mk_exists, mk_or, mk_pragma, mk_primitive,
-    mk_specialized_value, mk_type, mk_value, mk_var, mk_attribute
+    mk_specialized_value, mk_type, mk_value, mk_attribute
 )
 from relationalai.semantics.lqp.utils import TranslationCtx, gen_unique_var
 from relationalai.semantics.lqp.validators import assert_valid_input
@@ -253,7 +253,7 @@ def _translate_rank(ctx: TranslationCtx, rank: ir.Rank, body: lqp.Formula) -> lq
     # to convert it to Int128.
     result_var, _ = _translate_term(ctx, rank.result)
     # The primitive will return an Int64 result, so we need a var to hold the intermediary.
-    result_64_var = gen_unique_var(ctx, "rank_result_64")
+    result_64_var = gen_unique_var(ctx, "v_rank")
     result_64_type = mk_type(lqp.TypeName.INT)
 
     cast = lqp.Cast(input=result_64_var, result=result_var, meta=None)
@@ -306,7 +306,7 @@ def _translate_descending_rank(ctx: TranslationCtx, limit: int, result: lqp.Var,
     aggr_abstr_args = new_abstr_args + [(count_var, count_type)]
     count_aggr = lqp.Reduce(
         op=lqp_operator(
-            ctx.var_names,
+            ctx,
             "count",
             "count",
             mk_type(lqp.TypeName.INT)
@@ -431,7 +431,7 @@ def _translate_aggregate(ctx: TranslationCtx, aggr: ir.Aggregate, body: lqp.Form
         (sum_var, sum_type) = abstr_args[-2]
 
         result = lqp.Reduce(
-            op=lqp_avg_op(ctx.var_names, aggr.aggregation.name, sum_var.name, sum_type),
+            op=lqp_avg_op(ctx, aggr.aggregation.name, sum_var.name, sum_type),
             body=mk_abstraction(abstr_args, body),
             terms=[sum_result, count_result],
             meta=None,
@@ -464,7 +464,7 @@ def _translate_aggregate(ctx: TranslationCtx, aggr: ir.Aggregate, body: lqp.Form
 
     # Group-bys do not need to be handled at all, since they are introduced outside already
     reduce = lqp.Reduce(
-        op=lqp_operator(ctx.var_names, aggr.aggregation.name, aggr_arg.name, aggr_arg_type),
+        op=lqp_operator(ctx, aggr.aggregation.name, aggr_arg.name, aggr_arg_type),
         body=mk_abstraction(abstr_args, body),
         terms=output_vars,
         meta=None
@@ -523,9 +523,8 @@ def _translate_term(ctx: TranslationCtx, term: ir.Value) -> Tuple[lqp.Term, lqp.
         # TODO: ScalarType is not like other terms, should be handled separately.
         return to_lqp_value(term.name, types.String), meta_type_to_lqp(types.String)
     elif isinstance(term, ir.Var):
-        name = ctx.var_names.get_name_by_id(term.id, term.name)
         t = meta_type_to_lqp(term.type)
-        return mk_var(name), t
+        return _translate_var(ctx, term), t
     else:
         assert isinstance(term, ir.Literal), f"Cannot translate value {term!r} of type {type(term)} to LQP Term; neither Var nor Literal."
         v = to_lqp_value(term.value, term.type)
@@ -801,3 +800,7 @@ def _translate_join(ctx: TranslationCtx, task: ir.Lookup) -> lqp.Formula:
     output_term = _translate_term(ctx, target)[0]
 
     return lqp.Reduce(meta=None, op=op, body=body, terms=[output_term])
+
+def _translate_var(ctx: TranslationCtx, term: ir.Var) -> lqp.Var:
+    name = ctx.var_names.get_name_by_id(term.id, term.name)
+    return lqp.Var(name=name, meta=None)

relationalai/semantics/lqp/passes.py

@@ -6,17 +6,19 @@ from relationalai.semantics.metamodel.util import FrozenOrderedSet
 
 from relationalai.semantics.metamodel.rewrite import Flatten
 
-from ..metamodel.rewrite import DischargeConstraints, DNFUnionSplitter, ExtractNestedLogicals
-from .rewrite import CDC, ExtractCommon, ExtractKeys, FDConstraints, QuantifyVars, Splinter
+from ..metamodel.rewrite import DischargeConstraints, DNFUnionSplitter, ExtractNestedLogicals, FormatOutputs
+from .rewrite import CDC, ExtractCommon, ExtractKeys, FunctionAnnotations, QuantifyVars, Splinter
 
 from relationalai.semantics.lqp.utils import output_names
 
 from typing import cast, List, Sequence, Tuple, Union, Optional, Iterable
 from collections import defaultdict
+import pandas as pd
+import hashlib
 
 def lqp_passes() -> list[Pass]:
     return [
-        FDConstraints(),
+        FunctionAnnotations(),
         DischargeConstraints(),
         Checker(),
         CDC(), # specialize to physical relations before extracting nested and typing
@@ -25,6 +27,7 @@ def lqp_passes() -> list[Pass]:
         DNFUnionSplitter(),
         ExtractKeys(),
         ExtractCommon(),
+        FormatOutputs(),
         Flatten(),
         Splinter(), # Splits multi-headed rules into multiple rules
         QuantifyVars(), # Adds missing existentials
@@ -337,7 +340,7 @@ class UnifyDefinitions(Pass):
         )
 
 # Creates intermediary relations for all Data nodes and replaces said Data nodes
-# with a Lookup into these created relations.
+# with a Lookup into these created relations. Reuse duplicate created relations.
 class EliminateData(Pass):
     def rewrite(self, model: ir.Model, options:dict={}) -> ir.Model:
         r = self.DataRewriter()
@@ -350,17 +353,25 @@ class EliminateData(Pass):
         # Counter for naming new relations.
         # It must be that new_count == len new_updates == len new_relations.
         new_count: int
+        # Cache for Data nodes to avoid creating duplicate intermediary relations
+        data_cache: dict[str, ir.Relation]
 
         def __init__(self):
             self.new_relations = []
             self.new_updates = []
             self.new_count = 0
+            self.data_cache = {}
             super().__init__()
 
-        # Create a new intermediary relation representing the Data (and pop it in
-        # new_updates/new_relations) and replace this Data with a Lookup of said
-        # intermediary.
-        def handle_data(self, node: ir.Data, parent: ir.Node) -> ir.Lookup:
+        # Create a cache key for a Data node based on its structure and content
+        def _data_cache_key(self, node: ir.Data) -> str:
+            values = pd.util.hash_pandas_object(node.data).values
+            return hashlib.sha256(bytes(values)).hexdigest()
+
+        def _intermediary_relation(self, node: ir.Data) -> ir.Relation:
+            cache_key = self._data_cache_key(node)
+            if cache_key in self.data_cache:
+                return self.data_cache[cache_key]
             self.new_count += 1
             intermediary_name = f"formerly_Data_{self.new_count}"
 
@@ -379,7 +390,6 @@ class EliminateData(Pass):
                                 f.lookup(rel_builtins.eq, [f.literal(val), var])
                                 for (val, var) in zip(row, node.vars)
                             ],
-                            hoisted = node.vars,
                         )
                         for row in node
                     ],
@@ -390,6 +400,16 @@ class EliminateData(Pass):
             ])
             self.new_updates.append(intermediary_update)
 
+            # Cache the result for reuse
+            self.data_cache[cache_key] = intermediary_relation
+
+            return intermediary_relation
+
+        # Create a new intermediary relation representing the Data (and pop it in
+        # new_updates/new_relations) and replace this Data with a Lookup of said
+        # intermediary.
+        def handle_data(self, node: ir.Data, parent: ir.Node) -> ir.Lookup:
+            intermediary_relation = self._intermediary_relation(node)
             replacement_lookup = f.lookup(intermediary_relation, node.vars)
 
             return replacement_lookup

relationalai/semantics/lqp/primitives.py

@@ -1,8 +1,8 @@
 from relationalai.semantics.metamodel.types import digits_to_bits
 from relationalai.semantics.lqp import ir as lqp
 from relationalai.semantics.lqp.types import is_numeric
-from relationalai.semantics.lqp.utils import UniqueNames, lqp_hash
-from relationalai.semantics.lqp.constructors import mk_primitive, mk_specialized_value, mk_type, mk_value, mk_var
+from relationalai.semantics.lqp.utils import TranslationCtx, gen_unique_var, lqp_hash
+from relationalai.semantics.lqp.constructors import mk_primitive, mk_specialized_value, mk_type, mk_value
 
 rel_to_lqp = {
     "=": "rel_primitive_eq",
@@ -205,15 +205,15 @@ def is_monotype(name: str) -> bool:
 
 # We take the name and type of the variable that we're summing over, so that we can generate
 # recognizable names for the variables in the reduce operation and preserve the type.
-def lqp_avg_op(names: UniqueNames, op_name: str, sum_name: str, sum_type: lqp.Type) -> lqp.Abstraction:
+def lqp_avg_op(ctx: TranslationCtx, op_name: str, sum_name: str, sum_type: lqp.Type) -> lqp.Abstraction:
     count_type = mk_type(lqp.TypeName.INT)
     vars = [
-        (mk_var(names.get_name(sum_name)), sum_type),
-        (mk_var(names.get_name("counter")), count_type),
-        (mk_var(names.get_name(sum_name)), sum_type),
-        (mk_var(names.get_name("one")), count_type),
-        (mk_var(names.get_name("sum")), sum_type),
-        (mk_var(names.get_name("count")), count_type),
+        (gen_unique_var(ctx, sum_name), sum_type),
+        (gen_unique_var(ctx, "counter"), count_type),
+        (gen_unique_var(ctx, sum_name), sum_type),
+        (gen_unique_var(ctx, "one"), count_type),
+        (gen_unique_var(ctx, "sum"), sum_type),
+        (gen_unique_var(ctx, "count"), count_type),
     ]
 
     x1 = vars[0][0]
@@ -233,10 +233,10 @@ def lqp_avg_op(names: UniqueNames, op_name: str, sum_name: str, sum_type: lqp.Ty
     return lqp.Abstraction(vars=vars, value=body, meta=None)
 
 # Default handler for aggregation operations in LQP.
-def lqp_agg_op(names: UniqueNames, op_name: str, aggr_arg_name: str, aggr_arg_type: lqp.Type) -> lqp.Abstraction:
-    x = mk_var(names.get_name(f"x_{aggr_arg_name}"))
-    y = mk_var(names.get_name(f"y_{aggr_arg_name}"))
-    z = mk_var(names.get_name(f"z_{aggr_arg_name}"))
+def lqp_agg_op(ctx: TranslationCtx, op_name: str, aggr_arg_name: str, aggr_arg_type: lqp.Type) -> lqp.Abstraction:
+    x = gen_unique_var(ctx, f"x_{aggr_arg_name}")
+    y = gen_unique_var(ctx, f"y_{aggr_arg_name}")
+    z = gen_unique_var(ctx, f"z_{aggr_arg_name}")
     ts = [(x, aggr_arg_type), (y, aggr_arg_type), (z, aggr_arg_type)]
 
     name = agg_to_lqp.get(op_name, op_name)
@@ -244,9 +244,9 @@ def lqp_agg_op(names: UniqueNames, op_name: str, aggr_arg_name: str, aggr_arg_ty
 
     return lqp.Abstraction(vars=ts, value=body, meta=None)
 
-def lqp_operator(names: UniqueNames, op_name: str, aggr_arg_name: str, aggr_arg_type: lqp.Type) -> lqp.Abstraction:
+def lqp_operator(ctx: TranslationCtx, op_name: str, aggr_arg_name: str, aggr_arg_type: lqp.Type) -> lqp.Abstraction:
     # TODO: Can we just pass through unknown operations?
     if op_name not in agg_to_lqp:
         raise NotImplementedError(f"Unsupported aggregation: {op_name}")
 
-    return lqp_agg_op(names, op_name, aggr_arg_name, aggr_arg_type)
+    return lqp_agg_op(ctx, op_name, aggr_arg_name, aggr_arg_type)

relationalai/semantics/lqp/rewrite/__init__.py

@@ -1,7 +1,7 @@
 from .cdc import CDC
 from .extract_common import ExtractCommon
 from .extract_keys import ExtractKeys
-from .fd_constraints import FDConstraints
+from .function_annotations import FunctionAnnotations
 from .quantify_vars import QuantifyVars
 from .splinter import Splinter
 
@@ -9,7 +9,7 @@ __all__ = [
     "CDC",
     "ExtractCommon",
     "ExtractKeys",
-    "FDConstraints",
+    "FunctionAnnotations",
     "QuantifyVars",
     "Splinter",
 ]

relationalai/semantics/lqp/rewrite/{fd_constraints.py → function_annotations.py}

@@ -5,7 +5,7 @@ from relationalai.semantics.metamodel import ir, compiler as c, visitor as v, bu
 from relationalai.semantics.metamodel.util import OrderedSet, ordered_set
 
 
-class FDConstraints(c.Pass):
+class FunctionAnnotations(c.Pass):
     """
     Pass marks all appropriate relations with `function` annotation.
     Criteria:
@@ -17,7 +17,7 @@ class FDConstraints(c.Pass):
         collect_fd = CollectFunctionalRelationsVisitor()
         new_model = collect_fd.walk(model)
         # mark relations collected by previous visitor with `@function` annotation
-        return FDConstraintsVisitor(collect_fd.functional_relations).walk(new_model)
+        return FunctionalAnnotationsVisitor(collect_fd.functional_relations).walk(new_model)
 
 
 @dataclass
@@ -57,9 +57,9 @@ class CollectFunctionalRelationsVisitor(v.Rewriter):
 
 
 @dataclass
-class FDConstraintsVisitor(v.Rewriter):
+class FunctionalAnnotationsVisitor(v.Rewriter):
     """
-    This visitor marks functional_relations with `functional` annotation.
+    This visitor marks functional_relations with `function` annotation.
     """
 
     def __init__(self, functional_relations: OrderedSet):