relationalai 0.12.1__py3-none-any.whl → 0.12.3__py3-none-any.whl

relationalai/experimental/solvers.py

@@ -9,11 +9,13 @@ from ..metamodel import Builtins
 from ..tools.cli_controls import Spinner
 from ..tools.constants import DEFAULT_QUERY_TIMEOUT_MINS
 from .. import debugging
+from .. errors import ResponseStatusException
 import uuid
 import 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 ..util.timeout import calc_remaining_timeout_minutes
 
 rel_sv = rel._tagged(Builtins.SingleValued)
@@ -23,7 +25,8 @@ APP_NAME = 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.
-ENGINE_ERRORS = ["worker is suspended", "create/resume", "worker not found", "no workers found", "worker was deleted"]
+WORKER_ERRORS = ["worker is suspended", "create/resume", "worker not found", "no workers found", "worker was deleted"]
+ENGINE_ERRORS = ["engine is suspended", "create/resume", "engine not found", "no engines found", "engine was deleted"]
 ENGINE_NOT_READY_MSGS = ["worker is in pending", "worker is provisioning", "worker is not ready to accept jobs"]
 
 # --------------------------------------------------
@@ -213,13 +216,6 @@ class SolverModel:
         config_file_path = getattr(rai_config, 'file_path', None)
         start_time = time.monotonic()
         remaining_timeout_minutes = query_timeout_mins
-        # 1. Materialize the model and store it.
-        # TODO(coey) Currently we must run a dummy query to install the pyrel rules in a separate txn
-        # to the solve_output updates. Ideally pyrel would offer an option to flush the rules separately.
-        self.graph.exec_raw("", query_timeout_mins=remaining_timeout_minutes)
-        remaining_timeout_minutes = calc_remaining_timeout_minutes(
-            start_time, query_timeout_mins, config_file_path=config_file_path,
-        )
         response = self.graph.exec_raw(
             textwrap.dedent(f"""
             @inline
@@ -265,18 +261,7 @@ class SolverModel:
         remaining_timeout_minutes = calc_remaining_timeout_minutes(
             start_time, query_timeout_mins, config_file_path=config_file_path
         )
-        try:
-            job_id = solver._exec_job(payload, log_to_console=log_to_console, query_timeout_mins=remaining_timeout_minutes)
-        except Exception as e:
-            err_message = str(e).lower()
-            if any(kw in err_message.lower() for kw in ENGINE_ERRORS + ENGINE_NOT_READY_MSGS):
-                solver._auto_create_solver_async()
-                remaining_timeout_minutes = calc_remaining_timeout_minutes(
-                    start_time, query_timeout_mins, config_file_path=config_file_path
-                )
-                job_id = solver._exec_job(payload, log_to_console=log_to_console, query_timeout_mins=remaining_timeout_minutes)
-            else:
-                raise e
+        job_id = solver._exec_job(payload, log_to_console=log_to_console, query_timeout_mins=remaining_timeout_minutes)
 
         # 3. Extract result.
         remaining_timeout_minutes = calc_remaining_timeout_minutes(
@@ -553,7 +538,11 @@ class Solver:
         # may configure each individual solver.
         self.engine_settings = settings
 
-        return self._auto_create_solver_async()
+        # Optimistically set the engine object to a `READY` engine to
+        # avoid checking the engine status on each execution.
+        self.engine:Optional[dict[str,Any]] = {"name": engine_name, "state": "READY"}
+
+        return None
 
     # --------------------------------------------------
     # Helper
@@ -572,6 +561,7 @@ class Solver:
             assert len(engines) == 1 or len(engines) == 0
             if len(engines) != 0:
                 engine = engines[0]
+
             if engine:
                 # TODO (dba) Logic engines support altering the
                 # auto_suspend_mins setting. Currently, we don't have
@@ -653,31 +643,32 @@ class Solver:
 
             self.engine = engine
 
-    def _exec_job_async(self, payload, query_timeout_mins: Optional[int]=None):
-        payload_json = json.dumps(payload)
-        engine_name = self.engine["name"]
-        if query_timeout_mins is None and (timeout_value := self.rai_config.get("query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS)) is not None:
-            query_timeout_mins = int(timeout_value)
-        if query_timeout_mins is not None:
-            sql_string = textwrap.dedent(f"""
-            CALL {APP_NAME}.experimental.exec_job_async('{ENGINE_TYPE_SOLVER}', '{engine_name}', '{payload_json}', null, {query_timeout_mins})
-            """)
-        else:
-            sql_string = textwrap.dedent(f"""
-            CALL {APP_NAME}.experimental.exec_job_async('{ENGINE_TYPE_SOLVER}', '{engine_name}', '{payload_json}')
-            """)
-        res = self.provider.resources._exec(sql_string)
-        return res[0]["ID"]
-
     def _exec_job(self, payload, log_to_console=True, query_timeout_mins: Optional[int]=None):
-        # Make sure the engine is ready.
-        if self.engine["state"] != "READY":
-            poll_with_specified_overhead(lambda: self._is_solver_ready(), 0.1)
+        if self.engine is None:
+            raise Exception("Engine not initialized.")
 
         with debugging.span("job") as job_span:
-            job_id = self._exec_job_async(payload, query_timeout_mins=query_timeout_mins)
+            # Retry logic. If creating a job fails with an engine
+            # related error we will create/resume/... the engine and
+            # retry.
+            try:
+                job_id = self.provider.create_job_async(self.engine["name"], payload, query_timeout_mins=query_timeout_mins)
+            except Exception as e:
+                err_message = str(e).lower()
+                if isinstance(e, ResponseStatusException):
+                    err_message = e.response.json().get("message", "")
+                if any(kw in err_message.lower() for kw in ENGINE_ERRORS + WORKER_ERRORS + ENGINE_NOT_READY_MSGS):
+                    self._auto_create_solver_async()
+                    # Wait until the engine is ready.
+                    poll_with_specified_overhead(lambda: self._is_solver_ready(), 0.1)
+                    job_id = self.provider.create_job_async(self.engine["name"], payload, query_timeout_mins=query_timeout_mins)
+                else:
+                    raise e
+
             job_span["job_id"] = job_id
             debugging.event("job_created", job_span, job_id=job_id, engine_name=self.engine["name"], job_type=ENGINE_TYPE_SOLVER)
+            if not isinstance(job_id, str):
+                job_id = ""
             polling_state = PollingState(job_id, "", False, log_to_console)
 
             try:
@@ -693,7 +684,14 @@ class Solver:
             return job_id
 
     def _is_solver_ready(self):
+        if self.engine is None:
+            raise Exception("Engine not initialized.")
+
         result = self.provider.get_solver(self.engine["name"])
+
+        if result is None:
+            raise Exception("No engine available.")
+            
         self.engine = result
         state = result["state"]
         if state != "READY" and state != "PENDING":
@@ -711,20 +709,11 @@ class Solver:
 
         return status == "COMPLETED" or status == "FAILED" or status == "CANCELED"
 
-    def _get_job_events(self, job_id: str, continuation_token: str = ""):
-        results = self.provider.resources._exec(
-            f"SELECT {APP_NAME}.experimental.get_job_events('{ENGINE_TYPE_SOLVER}', '{job_id}', '{continuation_token}');"
-        )
-        if not results:
-            return {"events": [], "continuation_token": None}
-        row = results[0][0]
-        return json.loads(row)
-
     def _print_solver_logs(self, state: PollingState):
         if state.is_done:
             return
 
-        resp = self._get_job_events(state.job_id, state.continuation_token)
+        resp = self.provider.get_job_events(state.job_id, state.continuation_token)
 
         # Print solver logs to stdout.
         for event in resp["events"]:
@@ -754,7 +743,12 @@ class Provider:
             resources = relationalai.Resources()
         if not isinstance(resources, relationalai.clients.snowflake.Resources):
             raise Exception("Solvers are only supported on SPCS.")
+
         self.resources = resources
+        self.direct_access_client: Optional[DirectAccessClient] = None
+
+        if isinstance(self.resources, DirectAccessResources):
+            self.direct_access_client = self.resources.direct_access_client
 
     def create_solver(
         self,
@@ -770,75 +764,285 @@ class Provider:
         engine_config: dict[str, Any] = {"settings": settings}
         if auto_suspend_mins is not None:
             engine_config["auto_suspend_mins"] = auto_suspend_mins
-        self.resources._exec(
-            f"CALL {APP_NAME}.experimental.create_engine('{ENGINE_TYPE_SOLVER}', '{name}', '{size}', {engine_config});"
+        self.resources._exec_sql(
+            f"CALL {APP_NAME}.experimental.create_engine('{ENGINE_TYPE_SOLVER}', '{name}', '{size}', {engine_config});", None
         )
 
     def create_solver_async(
         self,
         name: str,
         size: str | None = None,
-        settings: dict = {},
+        settings: dict | None = None,
         auto_suspend_mins: int | None = None,
     ):
         if size is None:
             size = "HIGHMEM_X64_S"
-        if settings is None:
-            settings = ""
-        engine_config: dict[str, Any] = {"settings": settings}
-        if auto_suspend_mins is not None:
-            engine_config["auto_suspend_mins"] = auto_suspend_mins
-        self.resources._exec(
-            f"CALL {APP_NAME}.experimental.create_engine_async('{ENGINE_TYPE_SOLVER}', '{name}', '{size}', {engine_config});"
-        )
+
+        if self.direct_access_client is not None:
+            payload:dict[str, Any] = {
+                "name": name,
+                "settings": settings,
+               }
+            if auto_suspend_mins is not None:
+                payload["auto_suspend_mins"] = auto_suspend_mins
+            if size is not None:
+                payload["size"] = size
+            response = self.direct_access_client.request(
+                "create_engine",
+                payload=payload,
+                path_params={"engine_type": "solver"},
+               )
+            if response.status_code != 200:
+                raise ResponseStatusException(
+                    f"Failed to create engine {name} with size {size}.", response
+                )
+        else:
+            engine_config: dict[str, Any] = {}
+            if settings is not None:
+                engine_config["settings"] = settings
+            if auto_suspend_mins is not None:
+                engine_config["auto_suspend_mins"] = auto_suspend_mins
+            self.resources._exec_sql(
+                f"CALL {APP_NAME}.experimental.create_engine_async('{ENGINE_TYPE_SOLVER}', '{name}', '{size}', {engine_config});",
+                None
+               )
 
     def delete_solver(self, name: str):
-        self.resources._exec(
-            f"CALL {APP_NAME}.experimental.delete_engine('{ENGINE_TYPE_SOLVER}', '{name}');"
-        )
+        if self.direct_access_client is not None:
+            response = self.direct_access_client.request(
+                "delete_engine", path_params = {"engine_type": ENGINE_TYPE_SOLVER, "engine_name": name}
+            )
+            if response.status_code != 200:
+                raise ResponseStatusException("Failed to delete engine.", response)
+            return None
+        else:
+            self.resources._exec_sql(
+                f"CALL {APP_NAME}.experimental.delete_engine('{ENGINE_TYPE_SOLVER}', '{name}');",
+                None
+            )
 
     def resume_solver_async(self, name: str):
-        self.resources._exec(
-            f"CALL {APP_NAME}.experimental.resume_engine_async('{ENGINE_TYPE_SOLVER}', '{name}');"
-        )
+        if self.direct_access_client is not None:
+            response = self.direct_access_client.request(
+                "resume_engine", path_params = {"engine_type": ENGINE_TYPE_SOLVER, "engine_name": name}
+            )
+            if response.status_code != 200:
+                raise ResponseStatusException("Failed to resume engine.", response)
+            return None
+        else:
+            self.resources._exec_sql(
+                f"CALL {APP_NAME}.experimental.resume_engine_async('{ENGINE_TYPE_SOLVER}', '{name}');",
+                None
+            )
+            return None
 
     def get_solver(self, name: str):
-        results = self.resources._exec(
-            f"CALL {APP_NAME}.experimental.get_engine('{ENGINE_TYPE_SOLVER}', '{name}');"
-        )
-        return solver_list_to_dicts(results)[0]
+        if self.direct_access_client is not None:
+            response = self.direct_access_client.request(
+                "get_engine", path_params = {"engine_type": ENGINE_TYPE_SOLVER, "engine_name": name}
+            )
+            if response.status_code != 200:
+                raise ResponseStatusException("Failed to get engine.", response)
+            solver = response.json()
+            if not solver :
+                return None
+            solver_state = {
+                "name": solver["name"],
+                "id": solver["id"],
+                "size": solver["size"],
+                "state": solver["status"], # callers are expecting 'state'
+                "created_by": solver["created_by"],
+                "created_on": solver["created_on"],
+                "updated_on": solver["updated_on"],
+                "version": solver["version"],
+                "auto_suspend": solver["auto_suspend_mins"],
+                "suspends_at": solver["suspends_at"],
+                "solvers": []
+                if solver["settings"] == ""
+                else [
+                        k
+                        for (k,v) in json.loads(solver["settings"]).items()
+                        if isinstance(v, dict) and v.get("enabled", False)
+                ],
+                }
+            return solver_state
+        else:
+            results = self.resources._exec_sql(
+                f"CALL {APP_NAME}.experimental.get_engine('{ENGINE_TYPE_SOLVER}', '{name}');",
+                None
+               )
+            return solver_list_to_dicts(results)[0]
 
     def list_solvers(self, state: str | None = None):
-        where_clause = f"WHERE TYPE='{ENGINE_TYPE_SOLVER}'"
-        where_clause = (
-            f"{where_clause} AND STATUS = '{state.upper()}'" if state else where_clause
-        )
-        statement = f"SELECT NAME,ID,SIZE,STATUS,CREATED_BY,CREATED_ON,UPDATED_ON,AUTO_SUSPEND_MINS,SETTINGS FROM {APP_NAME}.experimental.engines {where_clause};"
-        results = self.resources._exec(statement)
-        return solver_list_to_dicts(results)
+        if self.direct_access_client is not None:
+            response = self.direct_access_client.request(
+                "list_engines"
+            )
+            if response.status_code != 200:
+                raise ResponseStatusException("Failed to list engines.", response)
+            response_content = response.json()
+            if not response_content:
+                return []
+            engines = [
+                {
+                    "name": engine["name"],
+                    "id": engine["id"],
+                    "size": engine["size"],
+                    "state": engine["status"], # callers are expecting 'state'
+                    "created_by": engine["created_by"],
+                    "created_on": engine["created_on"],
+                    "updated_on": engine["updated_on"],
+                    "auto_suspend_mins": engine["auto_suspend_mins"],
+                    "solvers": []
+                    if engine["settings"] == ""
+                    else [
+                        k
+                        for (k, v) in json.loads(engine["settings"]).items()
+                        if isinstance(v, dict) and v.get("enabled", False)
+                           ],
+                }
+                for engine in response_content.get("engines", [])
+                if (state is None or engine.get("status") == state) and (engine.get("type") == ENGINE_TYPE_SOLVER)
+               ]
+            return sorted(engines, key=lambda x: x["name"])
+        else:
+            where_clause = f"WHERE TYPE='{ENGINE_TYPE_SOLVER}'"
+            where_clause = (
+                f"{where_clause} AND STATUS = '{state.upper()}'" if state else where_clause
+               )
+            statement = f"SELECT NAME,ID,SIZE,STATUS,CREATED_BY,CREATED_ON,UPDATED_ON,AUTO_SUSPEND_MINS,SETTINGS FROM {APP_NAME}.experimental.engines {where_clause};"
+            results = self.resources._exec_sql(statement, None)
+            return solver_list_to_dicts(results)
 
     # --------------------------------------------------
     # Job API
     # --------------------------------------------------
 
+    def create_job_async(self, engine_name, payload, query_timeout_mins: Optional[int]=None):
+        payload_json = json.dumps(payload)
+
+        if query_timeout_mins is None and (timeout_value := self.resources.config.get("query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS)) is not None:
+            query_timeout_mins = int(timeout_value)
+            
+        if self.direct_access_client is not None:
+            job = {
+                "job_type":ENGINE_TYPE_SOLVER,
+                "worker_name": engine_name,
+                "timeout_mins": query_timeout_mins,
+                "payload": payload_json,
+            }
+            response = self.direct_access_client.request(
+                "create_job",
+                payload=job,
+            )
+            if response.status_code != 200:
+                raise ResponseStatusException("Failed to create job.", response)
+            response_content = response.json()
+            return response_content["id"]
+        else:
+            if query_timeout_mins is not None:
+                sql_string = textwrap.dedent(f"""
+                CALL {APP_NAME}.experimental.exec_job_async('{ENGINE_TYPE_SOLVER}', '{engine_name}', '{payload_json}', null, {query_timeout_mins})
+                """)
+            else:
+                sql_string = textwrap.dedent(f"""
+                CALL {APP_NAME}.experimental.exec_job_async('{ENGINE_TYPE_SOLVER}', '{engine_name}', '{payload_json}')
+                """)
+            res = self.resources._exec_sql(sql_string, None)
+            return res[0]["ID"]
+
     def list_jobs(self, state=None, limit=None):
-        state_clause = f"AND STATE = '{state.upper()}'" if state else ""
-        limit_clause = f"LIMIT {limit}" if limit else ""
-        results = self.resources._exec(
-            f"SELECT ID,STATE,CREATED_BY,CREATED_ON,FINISHED_AT,DURATION,PAYLOAD,ENGINE_NAME FROM {APP_NAME}.experimental.jobs where type='{ENGINE_TYPE_SOLVER}' {state_clause} ORDER BY created_on DESC {limit_clause};"
-        )
-        return job_list_to_dicts(results)
+        if self.direct_access_client is not None:
+            response = self.direct_access_client.request(
+                "list_jobs"
+            )
+            if response.status_code != 200:
+                raise ResponseStatusException("Failed to list jobs.", response)
+            response_content = response.json()
+            if not response_content:
+                return []
+            jobs = [
+                {
+                    "id": job["id"],
+                    "state": job["state"],
+                    "created_by": job["created_by"],
+                    "created_on": job["created_on"],
+                    "finished_at": job.get("finished_at", None),
+                    "duration": job["duration"] if "duration" in job else 0,
+                    "solver": json.loads(job["payload"]).get("solver", ""),
+                    "engine": job.get("engine_name", job["worker_name"]),
+                }
+                for job in response_content.get("jobs", [])
+                if state is None or job.get("state") == state
+               ]
+            return sorted(jobs, key=lambda x: x["created_on"], reverse=True)
+        else:
+            state_clause = f"AND STATE = '{state.upper()}'" if state else ""
+            limit_clause = f"LIMIT {limit}" if limit else ""
+            results = self.resources._exec_sql(
+                f"SELECT ID,STATE,CREATED_BY,CREATED_ON,FINISHED_AT,DURATION,PAYLOAD,ENGINE_NAME FROM {APP_NAME}.experimental.jobs where type='{ENGINE_TYPE_SOLVER}' {state_clause} ORDER BY created_on DESC {limit_clause};",
+                None
+               )
+            return job_list_to_dicts(results)
 
     def get_job(self, id: str):
-        results = self.resources._exec(
-            f"CALL {APP_NAME}.experimental.get_job('{ENGINE_TYPE_SOLVER}', '{id}');"
-        )
-        return job_list_to_dicts(results)[0]
+        if self.direct_access_client is not None:
+            response = self.direct_access_client.request(
+                "get_job", path_params = {"job_type": ENGINE_TYPE_SOLVER, "job_id": id}
+            )
+            if response.status_code != 200:
+                raise ResponseStatusException("Failed to get job.", response)
+            response_content = response.json()
+            return response_content["job"]
+        else:
+            results = self.resources._exec_sql(
+                f"CALL {APP_NAME}.experimental.get_job('{ENGINE_TYPE_SOLVER}', '{id}');",
+                None
+            )
+            return job_list_to_dicts(results)[0]
+
+    def get_job_events(self, job_id: str, continuation_token: str = ""):
+        if self.direct_access_client is not None:
+            response = self.direct_access_client.request(
+                "get_job_events",
+                path_params = {"job_type": ENGINE_TYPE_SOLVER, "job_id": job_id, "stream_name": "progress"},
+                query_params={"continuation_token": continuation_token},
+            )
+            if response.status_code != 200:
+                raise ResponseStatusException("Failed to get job events.", response)
+            response_content = response.json()
+            if not response_content:
+                return {
+                    "events": [],
+                    "continuation_token": None
+                }
+            return response_content
+        else:
+            results = self.resources._exec_sql(
+                f"SELECT {APP_NAME}.experimental.get_job_events('{ENGINE_TYPE_SOLVER}', '{job_id}', '{continuation_token}');",
+                None
+               )
+            if not results:
+                return {"events": [], "continuation_token": None}
+            row = results[0][0]
+            if not isinstance(row, str):
+                row = ""
+            return json.loads(row)
 
     def cancel_job(self, id: str):
-        self.resources._exec(
-            f"CALL {APP_NAME}.experimental.cancel_job('{ENGINE_TYPE_SOLVER}', '{id}');"
-        )
+        if self.direct_access_client is not None:
+            response = self.direct_access_client.request(
+                "cancel_job", path_params = {"job_type": ENGINE_TYPE_SOLVER, "job_id": id}
+            )
+            if response.status_code != 200:
+                raise ResponseStatusException("Failed to cancel job.", response)
+            return None
+        else:
+            self.resources._exec_sql(
+                f"CALL {APP_NAME}.experimental.cancel_job('{ENGINE_TYPE_SOLVER}', '{id}');",
+                None
+            )
+            return None
 
 
 def solver_list_to_dicts(results):
@@ -865,7 +1069,6 @@ def solver_list_to_dicts(results):
         for row in results
     ]
 
-
 def job_list_to_dicts(results):
     if not results:
         return []

relationalai/semantics/internal/internal.py

@@ -346,7 +346,7 @@ def find_select_keys(item: Any, keys:OrderedSet[Key]|None = None, enable_primiti
 
     if isinstance(item, (list, tuple)):
         for it in item:
-            find_select_keys(it, keys)
+            find_select_keys(it, keys, enable_primitive_key=enable_primitive_key)
 
     elif isinstance(item, (Relationship, RelationshipReading)) and item._parent:
         find_select_keys(item._parent, keys)
@@ -390,7 +390,7 @@ def find_select_keys(item: Any, keys:OrderedSet[Key]|None = None, enable_primiti
         find_select_keys(item._arg, keys)
 
     elif isinstance(item, Alias):
-        find_select_keys(item._thing, keys)
+        find_select_keys(item._thing, keys, enable_primitive_key=enable_primitive_key)
 
     elif isinstance(item, Aggregate):
         keys.update( Key(k, True) for k in item._group )
@@ -2418,21 +2418,21 @@ class Fragment():
 
     def meta(self, **kwargs: Any) -> Fragment:
         """Add metadata to the query.
-        
+
         Metadata can be used for debugging and observability purposes.
-        
+
         Args:
             **kwargs: Metadata key-value pairs
-        
+
         Returns:
             Fragment: Returns self for method chaining
-            
+
         Example:
             select(Person.name).meta(workload_name="test", priority=1, enabled=True)
         """
         if not kwargs:
             raise ValueError("meta() requires at least one argument")
-            
+
         self._meta.update(kwargs)
         return self
 
@@ -2560,7 +2560,7 @@ class Fragment():
         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)
-    
+
 
 #--------------------------------------------------
 # Select / Where

relationalai/semantics/lqp/executor.py

@@ -60,7 +60,7 @@ class LQPExecutor(e.Executor):
                 if not self.dry_run:
                     self.engine = self._resources.get_default_engine_name()
                     if not self.keep_model:
-                        atexit.register(self._resources.delete_graph, self.database, True)
+                        atexit.register(self._resources.delete_graph, self.database, True, "lqp")
         return self._resources
 
     # Checks the graph index and updates it if necessary
@@ -88,7 +88,15 @@ class LQPExecutor(e.Executor):
         assert self.engine is not None
 
         with debugging.span("poll_use_index", sources=sources, model=model, engine=engine_name):
-            resources.poll_use_index(app_name, sources, model, self.engine, engine_size, program_span_id)
+            resources.poll_use_index(
+                app_name=app_name,
+                sources=sources,
+                model=model,
+                engine_name=self.engine,
+                engine_size=engine_size,
+                language="lqp",
+                program_span_id=program_span_id,
+            )
 
     def report_errors(self, problems: list[dict[str, Any]], abort_on_error=True):
         from relationalai import errors
@@ -292,6 +300,27 @@ class LQPExecutor(e.Executor):
                 meta=None,
             )
 
+    def _compile_undefine_query(self, query_epoch: lqp_ir.Epoch) -> lqp_ir.Epoch:
+        fragment_ids = []
+
+        for write in query_epoch.writes:
+            if isinstance(write.write_type, lqp_ir.Define):
+                fragment_ids.append(write.write_type.fragment.id)
+
+        # Construct new Epoch with Undefine operations for all collected fragment IDs
+        undefine_writes = [
+            lqp_ir.Write(
+                write_type=lqp_ir.Undefine(fragment_id=frag_id, meta=None),
+                meta=None
+            )
+            for frag_id in fragment_ids
+        ]
+
+        return lqp_ir.Epoch(
+            writes=undefine_writes,
+            meta=None,
+        )
+
     def compile_lqp(self, model: ir.Model, task: ir.Task):
         configure = self._construct_configure()
 
@@ -326,7 +355,11 @@ class LQPExecutor(e.Executor):
         if model_txn is not None:
             epochs.append(model_txn.epochs[0])
 
-        epochs.append(query_txn.epochs[0])
+        query_txn_epoch = query_txn.epochs[0]
+
+        epochs.append(query_txn_epoch)
+
+        epochs.append(self._compile_undefine_query(query_txn_epoch))
 
         txn = lqp_ir.Transaction(epochs=epochs, configure=configure, meta=None)
 

relationalai/semantics/lqp/model2lqp.py

@@ -580,8 +580,9 @@ def get_relation_id(ctx: TranslationCtx, relation: ir.Relation, projection: list
     if relation.id in ctx.def_names.id_to_name:
         unique_name = ctx.def_names.id_to_name[relation.id]
     else:
-        prefix = helpers.relation_name_prefix(relation)
-        unique_name = ctx.def_names.get_name_by_id(relation.id, prefix + relation.name)
+        name = helpers.relation_name_prefix(relation) + relation.name
+        name = helpers.sanitize(name)
+        unique_name = ctx.def_names.get_name_by_id(relation.id, name)
 
     return utils.gen_rel_id(ctx, unique_name, types)
 

relationalai/semantics/lqp/rewrite/extract_common.py

@@ -265,10 +265,12 @@ class ExtractCommon(Pass):
         for child in common_body:
             body_output_vars.update(ctx.info.task_outputs(child))
 
-        # Compute the union of input vars across all composites, intersected with output
+        # Compute the union of input vars across all non-extracted tasks (basically
+        # composites and binders left behind), intersected with output
         # vars of the common body
         exposed_vars = OrderedSet.from_iterable(ctx.info.task_inputs(sample)) & body_output_vars
-        for composite in composites:
+        non_extracted_tasks = (binders - common_body) | composites
+        for composite in non_extracted_tasks:
             if composite is sample:
                 continue
             # compute common input vars