relationalai 1.0.0a2__py3-none-any.whl → 1.0.0a4__py3-none-any.whl

v0/relationalai/clients/resources/snowflake/util.py

@@ -0,0 +1,387 @@
+"""
+Utility functions for Snowflake resources.
+"""
+from __future__ import annotations
+import re
+import decimal
+import base64
+from numbers import Number
+from datetime import datetime, date
+from typing import List, Any, Dict, cast
+
+from .... import dsl
+from ....environments import runtime_env, SnowbookEnvironment
+
+# warehouse-based snowflake notebooks currently don't have hazmat
+crypto_disabled = False
+try:
+    from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+    from cryptography.hazmat.backends import default_backend
+    from cryptography.hazmat.primitives import padding
+except (ModuleNotFoundError, ImportError):
+    crypto_disabled = True
+
+# Constants used by helper functions
+ENGINE_ERRORS = ("engine is suspended", "create/resume", "engine not found", "no engines found", "engine was deleted")
+ENGINE_NOT_READY_MSGS = ("engine is in pending", "engine is provisioning")
+DATABASE_ERRORS = ("database not found",)
+
+# Constants for import/export and transaction processing
+VALID_IMPORT_STATES = ("PENDING", "PROCESSING", "QUARANTINED", "LOADED")
+IMPORT_STREAM_FIELDS = (
+    "ID", "CREATED_AT", "CREATED_BY", "STATUS", "REFERENCE_NAME", "REFERENCE_ALIAS",
+    "FQ_OBJECT_NAME", "RAI_DATABASE", "RAI_RELATION", "DATA_SYNC_STATUS",
+    "PENDING_BATCHES_COUNT", "NEXT_BATCH_STATUS", "NEXT_BATCH_UNLOADED_TIMESTAMP",
+    "NEXT_BATCH_DETAILS", "LAST_BATCH_DETAILS", "LAST_BATCH_UNLOADED_TIMESTAMP", "CDC_STATUS"
+)
+FIELD_MAP = {
+    "database_name": "database",
+    "engine_name": "engine",
+}
+
+
+def process_jinja_template(template: str, indent_spaces: int = 0, **substitutions: Any) -> str:
+    """Process a Jinja-like template.
+
+    Supports:
+    - Variable substitution {{ var }}
+    - Conditional blocks {% if condition %} ... {% endif %}
+    - For loops {% for item in items %} ... {% endfor %}
+    - Comments {# ... #}
+    - Whitespace control with {%- and -%}
+
+    Args:
+        template: The template string
+        indent_spaces: Number of spaces to indent the result
+        **substitutions: Variable substitutions
+    """
+
+    def evaluate_condition(condition: str, context: dict) -> bool:
+        """Safely evaluate a condition string using the context."""
+        # Replace variables with their values
+        for k, v in context.items():
+            if isinstance(v, str):
+                condition = condition.replace(k, f"'{v}'")
+            else:
+                condition = condition.replace(k, str(v))
+        try:
+            return bool(eval(condition, {"__builtins__": {}}, {}))
+        except Exception:
+            return False
+
+    def process_expression(expr: str, context: dict) -> str:
+        """Process a {{ expression }} block."""
+        expr = expr.strip()
+        if expr in context:
+            return str(context[expr])
+        return ""
+
+    def process_block(lines: List[str], context: dict, indent: int = 0) -> List[str]:
+        """Process a block of template lines recursively."""
+        result = []
+        i = 0
+        while i < len(lines):
+            line = lines[i]
+
+            # Handle comments
+            line = re.sub(r'{#.*?#}', '', line)
+
+            # Handle if blocks
+            if_match = re.search(r'{%\s*if\s+(.+?)\s*%}', line)
+            if if_match:
+                condition = if_match.group(1)
+                if_block = []
+                else_block = []
+                i += 1
+                nesting = 1
+                in_else_block = False
+                while i < len(lines) and nesting > 0:
+                    if re.search(r'{%\s*if\s+', lines[i]):
+                        nesting += 1
+                    elif re.search(r'{%\s*endif\s*%}', lines[i]):
+                        nesting -= 1
+                    elif nesting == 1 and re.search(r'{%\s*else\s*%}', lines[i]):
+                        in_else_block = True
+                        i += 1
+                        continue
+
+                    if nesting > 0:
+                        if in_else_block:
+                            else_block.append(lines[i])
+                        else:
+                            if_block.append(lines[i])
+                    i += 1
+                if evaluate_condition(condition, context):
+                    result.extend(process_block(if_block, context, indent))
+                else:
+                    result.extend(process_block(else_block, context, indent))
+                continue
+
+            # Handle for loops
+            for_match = re.search(r'{%\s*for\s+(\w+)\s+in\s+(\w+)\s*%}', line)
+            if for_match:
+                var_name, iterable_name = for_match.groups()
+                for_block = []
+                i += 1
+                nesting = 1
+                while i < len(lines) and nesting > 0:
+                    if re.search(r'{%\s*for\s+', lines[i]):
+                        nesting += 1
+                    elif re.search(r'{%\s*endfor\s*%}', lines[i]):
+                        nesting -= 1
+                    if nesting > 0:
+                        for_block.append(lines[i])
+                    i += 1
+                if iterable_name in context and isinstance(context[iterable_name], (list, tuple)):
+                    for item in context[iterable_name]:
+                        loop_context = dict(context)
+                        loop_context[var_name] = item
+                        result.extend(process_block(for_block, loop_context, indent))
+                continue
+
+            # Handle variable substitution
+            line = re.sub(r'{{\s*(\w+)\s*}}', lambda m: process_expression(m.group(1), context), line)
+
+            # Handle whitespace control
+            line = re.sub(r'{%-', '{%', line)
+            line = re.sub(r'-%}', '%}', line)
+
+            # Add line with proper indentation, preserving blank lines
+            if line.strip():
+                result.append(" " * (indent_spaces + indent) + line)
+            else:
+                result.append("")
+
+            i += 1
+
+        return result
+
+    # Split template into lines and process
+    lines = template.split('\n')
+    processed_lines = process_block(lines, substitutions)
+
+    return '\n'.join(processed_lines)
+
+
+def type_to_sql(type_obj: Any) -> str:
+    if type_obj is str:
+        return "VARCHAR"
+    if type_obj is int:
+        return "NUMBER"
+    if type_obj is Number:
+        return "DECIMAL(38, 15)"
+    if type_obj is float:
+        return "FLOAT"
+    if type_obj is decimal.Decimal:
+        return "DECIMAL(38, 15)"
+    if type_obj is bool:
+        return "BOOLEAN"
+    if type_obj is dict:
+        return "VARIANT"
+    if type_obj is list:
+        return "ARRAY"
+    if type_obj is bytes:
+        return "BINARY"
+    if type_obj is datetime:
+        return "TIMESTAMP"
+    if type_obj is date:
+        return "DATE"
+    if isinstance(type_obj, dsl.Type):
+        return "VARCHAR"
+    raise ValueError(f"Unknown type {type_obj}")
+
+
+def type_to_snowpark(type_obj: Any) -> str:
+    if type_obj is str:
+        return "StringType()"
+    if type_obj is int:
+        return "IntegerType()"
+    if type_obj is float:
+        return "FloatType()"
+    if type_obj is Number:
+        return "DecimalType(38, 15)"
+    if type_obj is decimal.Decimal:
+        return "DecimalType(38, 15)"
+    if type_obj is bool:
+        return "BooleanType()"
+    if type_obj is dict:
+        return "MapType()"
+    if type_obj is list:
+        return "ArrayType()"
+    if type_obj is bytes:
+        return "BinaryType()"
+    if type_obj is datetime:
+        return "TimestampType()"
+    if type_obj is date:
+        return "DateType()"
+    if isinstance(type_obj, dsl.Type):
+        return "StringType()"
+    raise ValueError(f"Unknown type {type_obj}")
+
+
+def sanitize_user_name(user: str) -> str:
+    """Sanitize a user name by extracting the part before '@' and replacing invalid characters."""
+    # Extract the part before the '@'
+    sanitized_user = user.split('@')[0]
+    # Replace any character that is not a letter, number, or underscore with '_'
+    sanitized_user = re.sub(r'[^a-zA-Z0-9_]', '_', sanitized_user)
+    return sanitized_user
+
+
+def is_engine_issue(response_message: str) -> bool:
+    """Check if a response message indicates an engine issue."""
+    return any(kw in response_message.lower() for kw in ENGINE_ERRORS + ENGINE_NOT_READY_MSGS)
+
+
+def is_database_issue(response_message: str) -> bool:
+    """Check if a response message indicates a database issue."""
+    return any(kw in response_message.lower() for kw in DATABASE_ERRORS)
+
+
+def collect_error_messages(e: Exception) -> list[str]:
+    """Collect all error messages from an exception and its chain.
+
+    Extracts messages from:
+    - str(e)
+    - e.message (if present, e.g., RAIException)
+    - e.args (string arguments)
+    - e.__cause__
+    - e.__context__
+    - Nested JavaScript execution errors
+    """
+    messages = [str(e).lower()]
+
+    # Check message attribute (RAIException has this)
+    if hasattr(e, 'message'):
+        msg = getattr(e, 'message', None)
+        if isinstance(msg, str):
+            messages.append(msg.lower())
+
+    # Check args
+    if hasattr(e, 'args') and e.args:
+        for arg in e.args:
+            if isinstance(arg, str):
+                messages.append(arg.lower())
+
+    # Check cause and context
+    if hasattr(e, '__cause__') and e.__cause__:
+        messages.append(str(e.__cause__).lower())
+    if hasattr(e, '__context__') and e.__context__:
+        messages.append(str(e.__context__).lower())
+
+    # Extract nested messages from JavaScript execution errors
+    for msg in messages[:]:  # Copy to avoid modification during iteration
+        if re.search(r"javascript execution error", msg):
+            matches = re.findall(r'"message"\s*:\s*"([^"]+)"', msg, re.IGNORECASE)
+            messages.extend([m.lower() for m in matches])
+
+    return messages
+
+
+#--------------------------------------------------
+# Parameter and Data Transformation Utilities
+#--------------------------------------------------
+
+def normalize_params(params: List[Any] | Any | None) -> List[Any] | None:
+    """Normalize parameters to a list format."""
+    if params is not None and not isinstance(params, list):
+        return cast(List[Any], [params])
+    return params
+
+
+def format_sproc_name(name: str, type_obj: Any) -> str:
+    """Format stored procedure parameter name based on type."""
+    if type_obj is datetime:
+        return f"{name}.astimezone(ZoneInfo('UTC')).isoformat(timespec='milliseconds')"
+    return name
+
+
+def is_azure_url(url: str) -> bool:
+    """Check if a URL is an Azure blob storage URL."""
+    return "blob.core.windows.net" in url
+
+
+def is_container_runtime() -> bool:
+    """Check if running in a container runtime environment."""
+    return isinstance(runtime_env, SnowbookEnvironment) and runtime_env.runner == "container"
+
+
+#--------------------------------------------------
+# Import/Export Utilities
+#--------------------------------------------------
+
+def is_valid_import_state(state: str) -> bool:
+    """Check if an import state is valid."""
+    return state in VALID_IMPORT_STATES
+
+
+def imports_to_dicts(results: List[Any]) -> List[Dict[str, Any]]:
+    """Convert import results to dictionaries with lowercase keys."""
+    parsed_results = [
+        {field.lower(): row[field] for field in IMPORT_STREAM_FIELDS}
+        for row in results
+    ]
+    return parsed_results
+
+
+#--------------------------------------------------
+# Transaction Utilities
+#--------------------------------------------------
+
+def txn_list_to_dicts(transactions: List[Any]) -> List[Dict[str, Any]]:
+    """Convert transaction list to dictionaries with field mapping."""
+    dicts = []
+    for txn in transactions:
+        dict = {}
+        txn_dict = txn.asDict()
+        for key in txn_dict:
+            mapValue = FIELD_MAP.get(key.lower())
+            if mapValue:
+                dict[mapValue] = txn_dict[key]
+            else:
+                dict[key.lower()] = txn_dict[key]
+        dicts.append(dict)
+    return dicts
+
+
+#--------------------------------------------------
+# Encryption Utilities
+#--------------------------------------------------
+
+def decrypt_stream(key: bytes, iv: bytes, src: bytes) -> bytes:
+    """Decrypt the provided stream with PKCS#5 padding handling."""
+    if crypto_disabled:
+        if isinstance(runtime_env, SnowbookEnvironment) and runtime_env.runner == "warehouse":
+            raise Exception("Please open the navigation-bar dropdown labeled *Packages* and select `cryptography` under the *Anaconda Packages* section, and then re-run your query.")
+        else:
+            raise Exception("library `cryptography.hazmat` missing; please install")
+
+    # `type:ignore`s are because of the conditional import, which
+    # we have because warehouse-based snowflake notebooks don't support
+    # the crypto library we're using.
+    cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())  # type: ignore
+    decryptor = cipher.decryptor()
+
+    # Decrypt the data
+    decrypted_padded_data = decryptor.update(src) + decryptor.finalize()
+
+    # Unpad the decrypted data using PKCS#5
+    unpadder = padding.PKCS7(128).unpadder()  # type: ignore # Use 128 directly for AES
+    unpadded_data = unpadder.update(decrypted_padded_data) + unpadder.finalize()
+
+    return unpadded_data
+
+
+def decrypt_artifact(data: bytes, encryption_material: str) -> bytes:
+    """Decrypts the artifact data using provided encryption material."""
+    encryption_material_parts = encryption_material.split("|")
+    assert len(encryption_material_parts) == 3, "Invalid encryption material"
+
+    algorithm, key_base64, iv_base64 = encryption_material_parts
+    assert algorithm == "AES_128_CBC", f"Unsupported encryption algorithm {algorithm}"
+
+    key = base64.standard_b64decode(key_base64)
+    iv = base64.standard_b64decode(iv_base64)
+
+    return decrypt_stream(key, iv, data)
+

v0/relationalai/early_access/dsl/ir/executor.py

@@ -6,10 +6,9 @@ from collections import defaultdict
 from typing import Any, List, Optional
 
 from pandas import DataFrame
-import v0.relationalai as rai
 from v0.relationalai import debugging
 from v0.relationalai.clients import result_helpers
-from v0.relationalai.clients.snowflake import APP_NAME
+from v0.relationalai.clients.resources.snowflake import APP_NAME
 from v0.relationalai.early_access.dsl.ir.compiler import Compiler
 from v0.relationalai.early_access.dsl.ontologies.models import Model
 from v0.relationalai.semantics.metamodel import ir
@@ -37,7 +36,8 @@ class RelExecutor:
         if not self._resources:
             with debugging.span("create_session"):
                 self.dry_run |= bool(self.config.get("compiler.dry_run", False))
-                self._resources = rai.clients.snowflake.Resources(
+                from v0.relationalai.clients.resources.snowflake import Resources
+                self._resources = Resources(
                     dry_run=self.dry_run,
                     config=self.config,
                     generation=Generation.QB,
@@ -257,4 +257,4 @@ class RelExecutor:
             if raw:
                 dataframe, errors = result_helpers.format_results(raw, None, result_cols)
                 self.report_errors(errors)
-            return DataFrame()
+            return DataFrame()

v0/relationalai/early_access/dsl/snow/api.py

@@ -2,13 +2,14 @@ from typing import cast, Optional
 
 import v0.relationalai as rai
 from v0.relationalai import Config
+from v0.relationalai.clients.resources.snowflake import Provider
 from v0.relationalai.early_access.dsl.snow.common import TabularMetadata, ColumnMetadata, SchemaMetadata, \
     ForeignKey, ColumnRef
 
 
 class Executor:
     def __init__(self, config: Optional[Config] = None):
-        self._provider = cast(rai.clients.snowflake.Provider, rai.Provider(config=config))
+        self._provider = cast(Provider, rai.Provider(config=config))
         self._table_meta_cache = {}
         self._schema_fk_cache = {}
 

v0/relationalai/errors.py

@@ -2436,6 +2436,24 @@ class QueryTimeoutExceededException(RAIException):
         Consider increasing the 'query_timeout_mins' parameter in your configuration file{f' (stored in {self.config_file_path})' if self.config_file_path else ''} to allow more time for query execution.
         """)
 
+class GuardRailsException(RAIException):
+    def __init__(self, progress: dict[str, Any]={}):
+        self.name = "Guard Rails Violation"
+        self.message = "Transaction aborted due to guard rails violation."
+        self.progress = progress
+        self.content = self.format_message()
+        super().__init__(self.message, self.name, self.content)
+
+    def format_message(self):
+        messages = [] if self.progress else [self.message]
+        for task in self.progress.get("tasks", {}).values():
+            for warning_type, warning_data in task.get("warnings", {}).items():
+                messages.append(textwrap.dedent(f"""
+                Relation Name: [yellow]{task["task_name"]}[/yellow]
+                Warning: {warning_type}
+                Message: {warning_data["message"]}
+                """))
+        return "\n".join(messages)
 
 #--------------------------------------------------
 # Azure Exceptions

v0/relationalai/experimental/solvers.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 import time
-from typing import Any, List, Optional
+from typing import Any, List, Optional, cast
 from dataclasses import dataclass
 import textwrap
 from .. import dsl, std
@@ -14,14 +14,13 @@ import uuid
 import v0.relationalai
 import json
 from ..clients.util import poll_with_specified_overhead
-from ..clients.snowflake import Resources as SnowflakeResources
-from ..clients.snowflake import DirectAccessClient, DirectAccessResources
+from ..clients.resources.snowflake import Resources as SnowflakeResources, APP_NAME
+from ..clients.resources.snowflake import DirectAccessResources
+from ..clients.direct_access_client import DirectAccessClient
 from ..util.timeout import calc_remaining_timeout_minutes
 
 rel_sv = rel._tagged(Builtins.SingleValued)
 
-APP_NAME = v0.relationalai.clients.snowflake.APP_NAME
-
 ENGINE_TYPE_SOLVER = "SOLVER"
 # TODO (dba) The ERP still uses `worker` instead of `engine`. Change
 # this once we fix this in the ERP.
@@ -741,10 +740,11 @@ class Provider:
     def __init__(self, resources=None):
         if not resources:
             resources = v0.relationalai.Resources()
-        if not isinstance(resources, v0.relationalai.clients.snowflake.Resources):
+        if not isinstance(resources, SnowflakeResources):
             raise Exception("Solvers are only supported on SPCS.")
 
-        self.resources = resources
+        # Type narrowing: resources is confirmed to be SnowflakeResources
+        self.resources: SnowflakeResources = cast(SnowflakeResources, resources)
         self.direct_access_client: Optional[DirectAccessClient] = None
 
         if isinstance(self.resources, DirectAccessResources):

v0/relationalai/semantics/devtools/benchmark_lqp.py

@@ -5,12 +5,11 @@ import argparse
 import os
 import json
 
-from v0.relationalai.clients.snowflake import Resources as snowflake_api
+from v0.relationalai.clients.resources.snowflake import Resources as snowflake_api, APP_NAME
 from v0.relationalai.semantics.lqp.executor import LQPExecutor
 from v0.relationalai.semantics.internal import internal
-from v0.relationalai.clients.use_index_poller import UseIndexPoller as index_poller
+from v0.relationalai.clients.resources.snowflake.use_index_poller import UseIndexPoller as index_poller
 from snowflake.connector.cursor import DictCursor
-from v0.relationalai.clients import snowflake
 
 from enum import Enum
 
@@ -172,7 +171,7 @@ def _exec_snowflake_override(bench_ctx, old_func, marker):
     def new_func(self, code, params, raw=False):
         cur = self._session.connection.cursor(DictCursor)
         try:
-            cur.execute(code.replace(snowflake.APP_NAME, self.get_app_name()), params)
+            cur.execute(code.replace(APP_NAME, self.get_app_name()), params)
             rows = cur.fetchall()
             qid = str(getattr(cur, "sfqid", None))
             assert qid is not None, "Snowflake query ID was not available"
@@ -398,7 +397,7 @@ def get_sf_query_info(bench_ctx):
     return result
 
 def _get_query_info(qids):
-    from v0.relationalai.clients.snowflake import Resources as snowflake_client
+    from v0.relationalai.clients.resources.snowflake import Resources as snowflake_client
     client = snowflake_client()
 
     qids_str = "','".join(qids)

v0/relationalai/semantics/devtools/extract_lqp.py

@@ -7,7 +7,7 @@ import os
 import json
 from contextlib import contextmanager
 
-from v0.relationalai.clients.snowflake import Resources as snowflake_api
+from v0.relationalai.clients.resources.snowflake import Resources as snowflake_api
 from v0.relationalai.semantics.internal import internal
 from typing import Dict, Optional
 

v0/relationalai/semantics/internal/snowflake.py

@@ -61,7 +61,7 @@ def get_session():
             _session = get_active_session()
         except Exception:
             from v0.relationalai import Resources
-            from v0.relationalai.clients.snowflake import Resources as SnowflakeResources
+            from v0.relationalai.clients.resources.snowflake import Resources as SnowflakeResources
             # TODO: we need a better way to handle global config
 
             # using the resource constructor to differentiate between direct access and

v0/relationalai/semantics/lqp/executor.py

@@ -7,7 +7,6 @@ import re
 from pandas import DataFrame
 from typing import Any, Optional, Literal, TYPE_CHECKING
 from snowflake.snowpark import Session
-import v0.relationalai as rai
 
 from v0.relationalai import debugging
 from v0.relationalai.errors import NonDefaultLQPSemanticsVersionWarning
@@ -21,10 +20,10 @@ from lqp import print as lqp_print, ir as lqp_ir
 from lqp.parser import construct_configure
 from v0.relationalai.semantics.lqp.ir import convert_transaction, validate_lqp
 from v0.relationalai.clients.config import Config
-from v0.relationalai.clients.snowflake import APP_NAME
+from v0.relationalai.clients.resources.snowflake import APP_NAME, create_resources_instance
 from v0.relationalai.clients.types import TransactionAsyncResponse
 from v0.relationalai.clients.util import IdentityParser, escape_for_f_string
-from v0.relationalai.tools.constants import USE_DIRECT_ACCESS, QUERY_ATTRIBUTES_HEADER
+from v0.relationalai.tools.constants import QUERY_ATTRIBUTES_HEADER
 from v0.relationalai.tools.query_utils import prepare_metadata_for_headers
 
 if TYPE_CHECKING:
@@ -32,7 +31,9 @@ if TYPE_CHECKING:
 
 # Whenever the logic engine introduces a breaking change in behaviour, we bump this version
 # once the client is ready to handle it.
-DEFAULT_LQP_SEMANTICS_VERSION = "0"
+#
+# [2026-01-09] bumping to 1 to opt-into hard validation errors from the engine
+DEFAULT_LQP_SEMANTICS_VERSION = "1"
 
 class LQPExecutor(e.Executor):
     """Executes LQP using the RAI client."""
@@ -67,17 +68,11 @@ class LQPExecutor(e.Executor):
         if not self._resources:
             with debugging.span("create_session"):
                 self.dry_run |= bool(self.config.get("compiler.dry_run", False))
-                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(
-                    dry_run=self.dry_run,
+                self._resources = create_resources_instance(
                     config=self.config,
-                    generation=rai.Generation.QB,
+                    dry_run=self.dry_run,
                     connection=self.connection,
                     language="lqp",
                 )

v0/relationalai/semantics/lqp/rewrite/extract_keys.py

@@ -118,6 +118,17 @@ class ExtractKeys(Pass):
   the same here).
 """
 class ExtractKeysRewriter(Rewriter):
+    def __init__(self):
+        super().__init__()
+        self.compound_keys: dict[Any, ir.Var] = {}
+
+    def _get_compound_key(self, orig_keys: Iterable[ir.Var]) -> ir.Var:
+        if orig_keys in self.compound_keys:
+            return self.compound_keys[orig_keys]
+        compound_key = f.var("compound_key", types.Hash)
+        self.compound_keys[orig_keys] = compound_key
+        return compound_key
+
     def handle_logical(self, node: ir.Logical, parent: ir.Node, ctx:Optional[Any]=None) -> ir.Logical:
         outputs = [x for x in node.body if isinstance(x, ir.Output) and x.keys]
         # We are not in a logical with an output at this level.
@@ -170,7 +181,7 @@ class ExtractKeysRewriter(Rewriter):
         annos = list(output.annotations)
         annos.append(f.annotation(builtins.output_keys, tuple(output_keys)))
         # Create a compound key that will be used in place of the original keys.
-        compound_key = f.var("compound_key", types.Hash)
+        compound_key = self._get_compound_key(output_keys)
 
         for key_combination in combinations:
             missing_keys = OrderedSet.from_iterable(output_keys)
@@ -192,8 +203,13 @@ class ExtractKeysRewriter(Rewriter):
             # handle the construct node in each clone
             values: list[ir.Value] = [compound_key.type]
             for key in output_keys:
-                assert isinstance(key.type, ir.ScalarType)
-                values.append(ir.Literal(types.String, key.type.name))
+                if isinstance(key.type, ir.UnionType):
+                    # the typer can derive union types when multiple distinct entities flow
+                    # into a relation's field, so use AnyEntity as the type marker
+                    values.append(ir.Literal(types.String, "AnyEntity"))
+                else:
+                    assert isinstance(key.type, ir.ScalarType)
+                    values.append(ir.Literal(types.String, key.type.name))
                 if key in key_combination:
                     values.append(key)
             body.add(ir.Construct(None, tuple(values), compound_key, OrderedSet().frozen()))
@@ -408,6 +424,12 @@ class ExtractKeysRewriter(Rewriter):
                             for arg in args[:-1]:
                                 extended_vars.add(arg)
                             there_is_progress = True
+                    elif isinstance(task, ir.Not):
+                        if isinstance(task.task, ir.Logical):
+                            hoisted = helpers.hoisted_vars(task.task.hoisted)
+                            if var in hoisted:
+                                partitions[var].add(task)
+                                there_is_progress = True
                     else:
                         assert False, f"invalid node kind {type(task)}"