relationalai 0.13.3__tar.gz → 0.13.5__tar.gz

{relationalai-0.13.3 → relationalai-0.13.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: relationalai
-Version: 0.13.3
+Version: 0.13.5
 Summary: RelationalAI Library and CLI
 Author-email: RelationalAI <support@relational.ai>
 License-File: LICENSE
@@ -11,7 +11,7 @@ Requires-Dist: colorama
 Requires-Dist: cryptography
 Requires-Dist: gravis
 Requires-Dist: inquirerpy
-Requires-Dist: lqp==0.2.1
+Requires-Dist: lqp==0.2.3
 Requires-Dist: nicegui==2.16.1
 Requires-Dist: numpy<2
 Requires-Dist: opentelemetry-api

{relationalai-0.13.3 → relationalai-0.13.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = 'relationalai'
-version = '0.13.3'
+version = '0.13.5'
 description = 'RelationalAI Library and CLI'
 readme="docs/pypi/README.md"
 authors = [
@@ -8,7 +8,7 @@ authors = [
 ]
 requires-python = ">= 3.10"
 dependencies = [
-    "lqp==0.2.1",
+    "lqp==0.2.3",
     "snowflake-connector-python[secure-local-storage]",
     "snowflake-snowpark-python>=1.38.0",
     "typing-extensions",

{relationalai-0.13.3 → relationalai-0.13.5}/src/relationalai/clients/exec_txn_poller.py

@@ -5,11 +5,14 @@ from typing import Dict, Optional, TYPE_CHECKING
 
 from relationalai import debugging
 from relationalai.clients.util import poll_with_specified_overhead
+from relationalai.clients.config import Config
 from relationalai.tools.cli_controls import create_progress
 from relationalai.util.format import format_duration
+from relationalai.tools.txn_progress import format_execution_tree
 
 if TYPE_CHECKING:
     from relationalai.clients.resources.snowflake import Resources
+    from relationalai.clients.resources.snowflake.snowflake import TxnStatusResponse
 
 # Polling behavior constants
 POLL_OVERHEAD_RATE = 0.1  # Overhead rate for exponential backoff
@@ -19,6 +22,14 @@ GREEN_COLOR = '\033[92m'
 GRAY_COLOR = '\033[90m'
 ENDC = '\033[0m'
 
+PRINT_TXN_PROGRESS_FLAG = "print_txn_progress"
+PRINT_INTERNAL_TXN_PROGRESS_FLAG = "print_txn_progress_internal"
+
+def should_print_txn_progress(config: Config) -> bool:
+    return bool(config.get(PRINT_TXN_PROGRESS_FLAG, False))
+
+def should_print_internal_txn_progress(config) -> bool:
+    return bool(config.get(PRINT_INTERNAL_TXN_PROGRESS_FLAG, False))
 
 class ExecTxnPoller:
     """
@@ -27,17 +38,19 @@ class ExecTxnPoller:
 
     def __init__(
         self,
-        print_txn_progress: bool,
-        resource: "Resources",
+        config: Config,
+        resource: Optional["Resources"] = None,
         txn_id: Optional[str] = None,
         headers: Optional[Dict] = None,
-        txn_start_time: Optional[float] = None,
+        txn_start_time: Optional[float] = None
     ):
-        self.print_txn_progress = print_txn_progress
+        self.print_txn_progress = should_print_txn_progress(config)
         self.res = resource
         self.txn_id = txn_id
         self.headers = headers or {}
         self.txn_start_time = txn_start_time or time.time()
+        self.print_internal_txn_progress = should_print_internal_txn_progress(config)
+        self.last_status: Optional[TxnStatusResponse] = None
 
     def __enter__(self) -> ExecTxnPoller:
         if not self.print_txn_progress:
@@ -53,17 +66,23 @@ class ExecTxnPoller:
         return self
 
     def __exit__(self, exc_type, exc_value, traceback) -> None:
-        if not self.print_txn_progress or self.txn_id is None:
+        if not self.print_txn_progress:
             return
         # Update to success message with duration
         total_duration = time.time() - self.txn_start_time
         txn_id = self.txn_id
         self.progress.update_main_status(
-            query_complete_message(txn_id, total_duration)
+            query_complete_message(txn_id, total_duration, internal_txn_progress=self._get_internal_progress())
         )
         self.progress.__exit__(exc_type, exc_value, traceback)
         return
 
+    def _get_internal_progress(self) -> Optional[Dict]:
+        """Get internal transaction progress if enabled and available."""
+        if self.print_internal_txn_progress and self.last_status:
+            return self.last_status.progress
+        return None
+
     def poll(self) -> bool:
         """
         Poll for transaction completion with interactive progress display.
@@ -79,44 +98,56 @@ class ExecTxnPoller:
         if self.print_txn_progress:
             # Update the main status to include the new txn_id
             self.progress.update_main_status_fn(
-                lambda: self.description_with_timing(txn_id),
+                lambda: self.description_with_timing(txn_id, self._get_internal_progress()),
             )
 
         # Don't show duration summary - we handle our own completion message
         def check_status() -> bool:
             """Check if transaction is complete."""
-            finished = self.res._check_exec_async_status(txn_id, headers=self.headers)
-            return finished
+            if self.res is None:
+                raise ValueError("Resource must be provided for polling.")
+            self.last_status = self.res._check_exec_async_status(txn_id, headers=self.headers)
+            return self.last_status.finished
 
-        with debugging.span("wait", txn_id=self.txn_id):
+        with debugging.span("wait", txn_id=txn_id):
             poll_with_specified_overhead(check_status, overhead_rate=POLL_OVERHEAD_RATE)
 
-
         return True
 
-    def description_with_timing(self, txn_id: str | None = None) -> str:
+    def description_with_timing(self, txn_id: str | None = None, internal_txn_progress: Dict | None = None) -> str:
         elapsed = time.time() - self.txn_start_time
         if txn_id is None:
             return query_progress_header(elapsed)
         else:
-            return query_progress_message(txn_id, elapsed)
+            return query_progress_message(txn_id, elapsed, internal_txn_progress)
 
 def query_progress_header(duration: float) -> str:
     # Don't print sub-second decimals, because it updates too fast and is distracting.
     duration_str = format_duration(duration, seconds_decimals=False)
     return f"Evaluating Query... {duration_str:>15}\n"
 
-def query_progress_message(id: str, duration: float) -> str:
-    return (
+def query_progress_message(id: str, duration: float, internal_txn_progress: Dict | None = None) -> str:
+    result = (
         query_progress_header(duration) +
         # Print with whitespace to align with the end of the transaction ID
         f"{GRAY_COLOR}ID: {id}{ENDC}"
     )
+    if internal_txn_progress is not None:
+        result += format_execution_tree(internal_txn_progress)
+    return result
 
-def query_complete_message(id: str, duration: float, status_header: bool = False) -> str:
-    return (
+def query_complete_message(id: str | None, duration: float, status_header: bool = False, internal_txn_progress: Dict | None = None) -> str:
+    out = (
         (f"{GREEN_COLOR}✅ " if status_header else "") +
         # Print with whitespace to align with the end of the transaction ID
-        f"Query Complete: {format_duration(duration):>21}\n" +
-        f"{GRAY_COLOR}ID: {id}{ENDC}"
-    )
+        f"Query Complete: {format_duration(duration):>21}"
+        )
+    if id is None:
+        out += ENDC
+    else:
+        out += f"\n{GRAY_COLOR}ID: {id}{ENDC}"
+
+    if internal_txn_progress is not None:
+        out += format_execution_tree(internal_txn_progress)
+
+    return out

{relationalai-0.13.3 → relationalai-0.13.5}/src/relationalai/clients/local.py

@@ -4,6 +4,7 @@ import base64
 import json
 from urllib.parse import quote, urlencode
 import pyarrow as pa
+import time
 import requests
 from email import message_from_bytes, policy
 from email.message import EmailMessage
@@ -18,6 +19,7 @@ from .config import Config
 from .types import TransactionAsyncResponse
 from .util import get_pyrel_version
 from ..errors import ResponseStatusException
+from ..clients.exec_txn_poller import ExecTxnPoller
 from .. import debugging
 
 @dataclass
@@ -112,7 +114,7 @@ class LocalResources(ResourcesBase):
 
     def reset(self):
         raise NotImplementedError("reset not supported in local mode")
-    
+
     #--------------------------------------------------
     # Check direct access is enabled (0 implemented)
     #--------------------------------------------------
@@ -332,7 +334,7 @@ class LocalResources(ResourcesBase):
     #--------------------------------------------------
     # Exec Async
     #--------------------------------------------------
-    
+
     def _parse_multipart_response(self, response: requests.Response) -> Dict[str, Any]:
         response_map = {}
         response_map['results'] = {}
@@ -464,11 +466,17 @@ class LocalResources(ResourcesBase):
             "readonly": readonly,
         }
 
-        parsed_response = self._create_transaction(
-            target_endpoint="create_txn",
-            payload=payload,
-            headers=headers
-        )
+        txn_start_time = time.time()
+        with ExecTxnPoller(
+            self.config,
+            txn_id=None,
+            txn_start_time=txn_start_time
+        ) as _poller:  # unused, except for __enter__ and __exit__ display
+            parsed_response = self._create_transaction(
+                target_endpoint="create_txn",
+                payload=payload,
+                headers=headers
+            )
 
         state = parsed_response["state"]
         if state not in ["COMPLETED", "ABORTED"]:

{relationalai-0.13.3 → relationalai-0.13.5}/src/relationalai/clients/resources/snowflake/__init__.py

@@ -2,7 +2,7 @@
 Snowflake resources module.
 """
 # Import order matters - Resources must be imported first since other classes depend on it
-from .snowflake import Resources, Provider, Graph, SnowflakeClient, APP_NAME, PYREL_ROOT_DB, ExecContext, PrimaryKey, PRINT_TXN_PROGRESS_FLAG
+from .snowflake import Resources, Provider, Graph, SnowflakeClient, APP_NAME, PYREL_ROOT_DB, ExecContext, PrimaryKey
 from .engine_service import EngineType, INTERNAL_ENGINE_SIZES, ENGINE_SIZES_AWS, ENGINE_SIZES_AZURE
 # These imports depend on Resources, so they come after
 from .cli_resources import CLIResources
@@ -14,7 +14,7 @@ __all__ = [
     'Resources', 'DirectAccessResources', 'Provider', 'Graph', 'SnowflakeClient',
     'APP_NAME', 'PYREL_ROOT_DB', 'CLIResources', 'UseIndexResources', 'ExecContext', 'EngineType',
     'INTERNAL_ENGINE_SIZES', 'ENGINE_SIZES_AWS', 'ENGINE_SIZES_AZURE', 'PrimaryKey',
-    'PRINT_TXN_PROGRESS_FLAG', 'create_resources_instance',
+    'create_resources_instance',
 ]
 
 

{relationalai-0.13.3 → relationalai-0.13.5}/src/relationalai/clients/resources/snowflake/direct_access_resources.py

@@ -18,7 +18,7 @@ from snowflake.snowpark import Session
 
 # Import UseIndexResources to enable use_index functionality with direct access
 from .use_index_resources import UseIndexResources
-from .snowflake import TxnCreationResult
+from .snowflake import TxnCreationResult, TxnStatusResponse
 
 # Import helper functions from util
 from .util import is_engine_issue as _is_engine_issue, is_database_issue as _is_database_issue, collect_error_messages
@@ -314,7 +314,7 @@ class DirectAccessResources(UseIndexResources):
 
             return response.json()
 
-    def _check_exec_async_status(self, txn_id: str, headers: Dict[str, str] | None = None) -> bool:
+    def _check_exec_async_status(self, txn_id: str, headers: Dict[str, str] | None = None) -> TxnStatusResponse:
         """Check whether the given transaction has completed."""
 
         with debugging.span("check_status"):
@@ -349,8 +349,12 @@ class DirectAccessResources(UseIndexResources):
             elif reason == TXN_ABORT_REASON_GUARD_RAILS:
                 raise GuardRailsException(response_content.get("progress", {}))
 
-        # @TODO: Find some way to tunnel the ABORT_REASON out. Azure doesn't have this, but it's handy
-        return status == "COMPLETED" or status == "ABORTED"
+        return TxnStatusResponse(
+            txn_id=txn_id,
+            finished=status in ["COMPLETED", "ABORTED"],
+            abort_reason=response_content.get("abort_reason", None),
+            progress=response_content.get("progress", None),
+        )
 
     def _list_exec_async_artifacts(self, txn_id: str, headers: Dict[str, str] | None = None) -> Dict[str, Dict]:
         """Grab the list of artifacts produced in the transaction and the URLs to retrieve their contents."""

{relationalai-0.13.3 → relationalai-0.13.5}/src/relationalai/clients/resources/snowflake/snowflake.py

@@ -15,7 +15,7 @@ import hashlib
 from dataclasses import dataclass
 
 from ....auth.token_handler import TokenHandler
-from relationalai.clients.exec_txn_poller import ExecTxnPoller
+from ....clients.exec_txn_poller import ExecTxnPoller
 import snowflake.snowpark
 
 from ....rel_utils import sanitize_identifier, to_fqn_relation_name
@@ -104,14 +104,16 @@ TERMINAL_TXN_STATES = ["COMPLETED", "ABORTED"]
 TXN_ABORT_REASON_TIMEOUT = "transaction timeout"
 GUARDRAILS_ABORT_REASON = "guard rail violation"
 
-PRINT_TXN_PROGRESS_FLAG = "print_txn_progress"
+ENABLE_GUARD_RAILS_FLAG = "enable_guard_rails"
+
+ENABLE_GUARD_RAILS_HEADER = "X-RAI-Enable-Guard-Rails"
 
 #--------------------------------------------------
 # Helpers
 #--------------------------------------------------
 
-def should_print_txn_progress(config) -> bool:
-    return bool(config.get(PRINT_TXN_PROGRESS_FLAG, False))
+def should_enable_guard_rails(config) -> bool:
+    return bool(config.get(ENABLE_GUARD_RAILS_FLAG, False))
 
 #--------------------------------------------------
 # Resources
@@ -151,6 +153,14 @@ class TxnCreationResult:
     artifact_info: Dict[str, Dict]  # Populated if fast-path (state is COMPLETED/ABORTED)
 
 
+@dataclass
+class TxnStatusResponse:
+    """Transaction progress response for transaction status checks."""
+    txn_id: str
+    finished: bool
+    abort_reason: str | None = None
+    progress: Dict | None = None
+
 class Resources(ResourcesBase):
     def __init__(
         self,
@@ -1403,7 +1413,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
     # Exec Async
     #--------------------------------------------------
 
-    def _check_exec_async_status(self, txn_id: str, headers: Dict | None = None):
+    def _check_exec_async_status(self, txn_id: str, headers: Dict | None = None) -> TxnStatusResponse:
         """Check whether the given transaction has completed."""
         if headers is None:
             headers = {}
@@ -1433,8 +1443,11 @@ Otherwise, remove it from your '{profile}' configuration profile.
             elif response_row.get("ABORT_REASON", "") == GUARDRAILS_ABORT_REASON:
                 raise GuardRailsException()
 
-        # @TODO: Find some way to tunnel the ABORT_REASON out. Azure doesn't have this, but it's handy
-        return status == "COMPLETED" or status == "ABORTED"
+        return TxnStatusResponse(
+            txn_id=txn_id,
+            finished=status in ["COMPLETED", "ABORTED"],
+            abort_reason=response_row.get("ABORT_REASON", None),
+        )
 
 
     def _list_exec_async_artifacts(self, txn_id: str, headers: Dict | None = None) -> Dict[str, Dict]:
@@ -1788,10 +1801,8 @@ Otherwise, remove it from your '{profile}' configuration profile.
 
         with debugging.span("transaction", **query_attrs_dict) as txn_span:
             txn_start_time = time.time()
-            print_txn_progress = should_print_txn_progress(self.config)
-
             with ExecTxnPoller(
-                print_txn_progress=print_txn_progress,
+                config=self.config,
                 resource=self, txn_id=None, headers=request_headers,
                 txn_start_time=txn_start_time
             ) as poller:
@@ -1800,6 +1811,7 @@ Otherwise, remove it from your '{profile}' configuration profile.
                     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 ""
+                    request_headers[ENABLE_GUARD_RAILS_HEADER] = str(should_enable_guard_rails(self.config))
 
                     # Create the transaction
                     result = self._create_v2_txn(

{relationalai-0.13.3 → relationalai-0.13.5}/src/relationalai/clients/resources/snowflake/use_index_poller.py

@@ -66,7 +66,7 @@ ERP_CHECK_FREQUENCY = 15
 
 # Polling behavior constants
 POLL_OVERHEAD_RATE = 0.1  # Overhead rate for exponential backoff
-POLL_MAX_DELAY = 2.5  # Maximum delay between polls in seconds
+POLL_MAX_DELAY = 0  # Maximum delay between polls in seconds
 
 # SQL query template for getting stream column hashes
 # This query calculates a hash of column metadata (name, type, precision, scale, nullable)

{relationalai-0.13.3 → relationalai-0.13.5}/src/relationalai/experimental/solvers.py

@@ -533,6 +533,14 @@ class Solver:
         self.engine_size = engine_size or settings.pop("engine_size", None)
         self.engine_auto_suspend_mins = auto_suspend_mins or settings.pop("auto_suspend_mins", None)
 
+        # Set default CSV store setting if not already configured
+        if "store" not in settings:
+            settings["store"] = {}
+        if "csv" not in settings["store"]:
+            settings["store"]["csv"] = {}
+        if "enabled" not in settings["store"]["csv"]:
+            settings["store"]["csv"]["enabled"] = True
+
         # The settings are used when creating a solver engine, they
         # may configure each individual solver.
         self.engine_settings = settings

{relationalai-0.13.3 → relationalai-0.13.5}/src/relationalai/semantics/internal/snowflake.py

@@ -214,6 +214,7 @@ class Table():
         self._col_names = cols
         self._iceberg_config = config
         self._is_iceberg = config is not None
+        self._skip_cdc = False
         info = self._schemas.get((self._database, self._schema))
         if not info:
             info = self._schemas[(self._database, self._schema)] = SchemaInfo(self._database, self._schema)
@@ -303,7 +304,10 @@ class Table():
 
     def _compile_lookup(self, compiler:b.Compiler, ctx:b.CompilerContext):
         self._lazy_init()
-        Table._used_sources.add(self)
+        if not self._skip_cdc:
+            # Don't do CDC if the underlying data has been loaded
+            # directly via `api.load_data`.
+            Table._used_sources.add(self)
         compiler.lookup(self._rel, ctx)
         return compiler.lookup(b.RelationshipFieldRef(None, self._rel, 0), ctx)
 

{relationalai-0.13.3 → relationalai-0.13.5}/src/relationalai/semantics/lqp/executor.py

@@ -10,8 +10,9 @@ from snowflake.snowpark import Session
 
 from relationalai import debugging
 from relationalai.errors import NonDefaultLQPSemanticsVersionWarning
-from relationalai.semantics.lqp import result_helpers
+from relationalai.semantics.lqp import result_helpers, export_rewriter
 from relationalai.semantics.metamodel import ir, factory as f, executor as e
+from relationalai.semantics.metamodel.visitor import collect_by_type
 
 if TYPE_CHECKING:
     from relationalai.semantics.internal.internal import Model as InternalModel
@@ -368,7 +369,7 @@ class LQPExecutor(e.Executor):
             meta=None,
         )
 
-    def compile_lqp(self, model: ir.Model, task: ir.Task):
+    def compile_lqp(self, model: ir.Model, task: ir.Task, format: Optional[Literal["pandas", "snowpark", "csv"]] = "pandas"):
         configure = self._construct_configure()
         # Merge the epochs into a single transaction. Long term the query bits should all
         # go into a WhatIf action and the intrinsics could be fused with either of them. But
@@ -394,6 +395,21 @@ class LQPExecutor(e.Executor):
             result, final_model = self.compiler.compile_inner(query, options)
             export_info, query_epoch = result
 
+            if format == "csv":
+                # Extract original column names from Output
+                outputs = collect_by_type(ir.Output, task)
+                assert outputs, "No Output found in the task"
+                assert len(outputs) == 1, "Multiple Outputs found in the task"
+                output = outputs[0]
+                original_cols = []
+                for alias, _ in output.aliases:
+                    if not alias:
+                        continue
+                    original_cols.append(alias)
+                # Use rewriter to filter data_columns
+                column_filter = export_rewriter.ExtraColumnsFilter(original_cols)
+                query_epoch = column_filter.filter_epoch(query_epoch)
+
             epochs.append(query_epoch)
             epochs.append(self._compile_undefine_query(query_epoch))
 
@@ -417,8 +433,7 @@ class LQPExecutor(e.Executor):
     ) -> DataFrame:
         self.prepare_data()
         previous_model = self._last_model
-
-        final_model, export_info, txn_proto = self.compile_lqp(model, task)
+        final_model, export_info, txn_proto = self.compile_lqp(model, task, format=format)
 
         if self.dry_run:
             return DataFrame()
@@ -440,11 +455,13 @@ class LQPExecutor(e.Executor):
             nowait_durable=True,
             headers=headers,
         )
-        assert isinstance(raw_results, TransactionAsyncResponse)
+        assert isinstance(raw_results, TransactionAsyncResponse), "Expected TransactionAsyncResponse from LQP execution"
+        assert raw_results.transaction is not None, "Transaction result is missing"
+        txid = raw_results.transaction['id']
 
         try:
-            cols, extra_cols = self._compute_cols(task, final_model)
-            df, errs = result_helpers.format_results(raw_results, cols)
+            cols, extra_cols, key_locs = self._compute_cols(task, final_model)
+            df, errs = result_helpers.format_results(raw_results, cols, key_locs)
             self.report_errors(errs)
 
             # Rename columns if wide outputs is enabled
@@ -453,7 +470,6 @@ class LQPExecutor(e.Executor):
 
             if export_to:
                 assert cols, "No columns found in the output"
-                assert raw_results.transaction, "Invalid transaction result"
                 assert export_info, "Export info should be populated if we are exporting results"
                 result_cols = export_to._col_names
                 if result_cols is not None:
@@ -461,14 +477,18 @@ class LQPExecutor(e.Executor):
                 else:
                     result_cols = [col for col in cols if col not in extra_cols]
                 assert result_cols
-                self._export(raw_results.transaction['id'], export_info, export_to, cols, result_cols, update)
+                self._export(txid, export_info, export_to, cols, result_cols, update)
 
             if format == "csv":
                 if export_info is not None and isinstance(export_info, tuple) and isinstance(export_info[0], str):
-                    return DataFrame([export_info[0]], columns=["path"])
+                    # The full CSV path has two parts. The first part is chosen by the frontend, while
+                    # the second part is chosen by the backend to avoid collisions. We need to ensure
+                    # the second part is synchronized with the future changes in the backend.
+                    full_path = export_info[0] + f"/data_{txid}.gz"
+                    return DataFrame([full_path], columns=["path"])
                 else:
                     raise ValueError("The CSV export was not successful!")
-            
+
             return self._postprocess_df(self.config, df, extra_cols)
 
         except Exception as e:

relationalai-0.13.5/src/relationalai/semantics/lqp/export_rewriter.py

@@ -0,0 +1,40 @@
+#----------------------------------------------------------------------------------------------
+# This is a custom LQP rewriter that filters extra columns from CSV export. It is used in the
+# LQP executor, when the format="csv", to ensure only intended columns are being exported.
+#----------------------------------------------------------------------------------------------
+
+from dataclasses import replace
+from lqp import ir as lqp_ir
+
+class ExtraColumnsFilter:
+
+    def __init__(self, original_cols: list[str]):
+        self.original_cols = set(original_cols)
+
+    def filter_epoch(self, query_epoch: lqp_ir.Epoch) -> lqp_ir.Epoch:
+
+        # Only process epochs with a single read which is dedicated to Export
+        if not (query_epoch.reads and len(query_epoch.reads) == 1):
+            return query_epoch
+
+        old_read = query_epoch.reads[0]
+        if not isinstance(old_read.read_type, lqp_ir.Export):
+            return query_epoch
+
+        config = old_read.read_type.config
+        assert isinstance(config, lqp_ir.ExportCSVConfig) and config.data_columns is not None, \
+            "Expected ExportCSVConfig with data_columns in the read type"
+
+        data_columns = config.data_columns
+
+        # Filter data_columns to only include columns in original_cols
+        new_data_columns = [col for col in data_columns if col.column_name in self.original_cols]
+
+        # Reconstruct the nested structure with filtered data_columns
+        new_config = replace(old_read.read_type.config, data_columns=new_data_columns)
+        new_read_type = replace(old_read.read_type, config=new_config)
+        new_read = replace(old_read, read_type=new_read_type)
+
+        # Return new epoch with updated read
+        remaining_reads = list(query_epoch.reads[1:])
+        return replace(query_epoch, reads=[new_read] + remaining_reads)

{relationalai-0.13.3 → relationalai-0.13.5}/src/relationalai/semantics/lqp/model2lqp.py

@@ -126,34 +126,43 @@ def _translate_to_decls(ctx: TranslationCtx, rule: ir.Logical) -> list[lqp.Decla
 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}")
+        if isinstance(task, ir.Logical):
+            constraint_decls.extend(_translate_to_constraint_decls(ctx, task))
             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)
-
+        else:
+            assert isinstance(task, ir.Require)
+            decl = _translate_to_constraint_decl(ctx, task)
+            if decl is not None:
+                constraint_decls.append(decl)
     return constraint_decls
 
+def _translate_to_constraint_decl(ctx: TranslationCtx, rule: ir.Require) -> Optional[lqp.Declaration]:
+    fd = normalized_fd(rule)
+    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}")
+        return None
+
+    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
+    lqp_id = utils.lqp_hash(fd.canonical_str)
+    lqp_name:lqp.RelationId = lqp.RelationId(id=lqp_id, meta=None)
+
+    return lqp.FunctionalDependency(
+        name=lqp_name,
+        guard=lqp_guard,
+        keys=lqp_keys,
+        values=lqp_values,
+        meta=None
+    )
+
 def _translate_algorithms(ctx: TranslationCtx, task: ir.Logical) -> list[lqp.Declaration]:
     assert is_algorithm_logical(task)
     decls: list[lqp.Declaration] = []