relationalai 0.13.0__py3-none-any.whl → 0.13.0.dev0__py3-none-any.whl

relationalai/semantics/rel/executor.py

@@ -1,362 +0,0 @@
-from __future__ import annotations
-import atexit
-from collections import defaultdict
-import json
-import re
-import textwrap
-import uuid
-
-from pandas import DataFrame
-from typing import Any, Optional, Literal, TYPE_CHECKING
-from snowflake.snowpark import Session
-
-from relationalai import debugging
-from relationalai.clients import result_helpers
-from relationalai.clients.util import IdentityParser, escape_for_f_string
-from relationalai.clients.resources.snowflake import APP_NAME, create_resources_instance
-from relationalai.semantics.metamodel import ir, executor as e, factory as f
-from relationalai.semantics.rel import Compiler
-from relationalai.clients.config import Config
-from relationalai.tools.constants import Generation, QUERY_ATTRIBUTES_HEADER
-from relationalai.tools.query_utils import prepare_metadata_for_headers
-
-if TYPE_CHECKING:
-    from relationalai.semantics.snowflake import Table
-
-class RelExecutor(e.Executor):
-    """Executes Rel code using the RAI client."""
-
-    def __init__(
-        self,
-        database: str,
-        dry_run: bool = False,
-        keep_model: bool = True,
-        wide_outputs: bool = False,
-        connection: Session | None = None,
-        config: Config | None = None,
-    ) -> None:
-        super().__init__()
-        self.database = database
-        self.dry_run = dry_run
-        self.keep_model = keep_model
-        self.wide_outputs = wide_outputs
-        self.compiler = Compiler()
-        self.connection = connection
-        self.config = config or Config()
-        self._resources = None
-        self._last_model = None
-        self._last_sources_version = (-1, None)
-
-    @property
-    def resources(self):
-        if not self._resources:
-            with debugging.span("create_session"):
-                self.dry_run |= bool(self.config.get("compiler.dry_run", False))
-                # NOTE: language="rel" is required for Rel execution. It is the default, but
-                # we set it explicitly here to be sure.
-                self._resources = create_resources_instance(
-                    config=self.config,
-                    dry_run=self.dry_run,
-                    connection=self.connection,
-                    language="rel",
-                )
-                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, "rel")
-        return self._resources
-
-    def check_graph_index(self, headers: dict[str, Any] | None = None):
-        # Has to happen first, so self.dry_run is populated.
-        resources = self.resources
-
-        if self.dry_run:
-            return
-
-        from relationalai.semantics.snowflake import Table
-        table_sources = Table._used_sources
-        if not table_sources.has_changed(self._last_sources_version):
-            return
-
-        model = self.database
-        app_name = resources.get_app_name()
-        engine_name = self.engine
-        engine_size = self.resources.config.get_default_engine_size()
-
-        program_span_id = debugging.get_program_span_id()
-        sources = [t._fqn for t in Table._used_sources]
-        self._last_sources_version = Table._used_sources.version()
-
-        assert self.engine is not None
-
-        with debugging.span("poll_use_index", sources=sources, model=model, engine=engine_name):
-            resources.maybe_poll_use_index(
-                app_name=app_name,
-                sources=sources,
-                model=model,
-                engine_name=self.engine,
-                engine_size=engine_size,
-                program_span_id=program_span_id,
-                headers=headers,
-            )
-
-    def report_errors(self, problems: list[dict[str, Any]], abort_on_error=True):
-        from relationalai import errors
-        all_errors = []
-        undefineds = []
-        pyrel_errors = defaultdict(list)
-        pyrel_warnings = defaultdict(list)
-
-        for problem in problems:
-            message = problem.get("message", "")
-            report = problem.get("report", "")
-            # TODO: we need to build source maps
-            # path = problem.get("path", "")
-            # source_task = self._install_batch.line_to_task(path, problem["start_line"]) or task
-            # source = debugging.get_source(source_task) or debugging.SourceInfo()
-            source = debugging.SourceInfo()
-            severity = problem.get("severity", "warning")
-            code = problem.get("code")
-
-            if severity in ["error", "exception"]:
-                if code == "UNDEFINED_IDENTIFIER":
-                    match = re.search(r'`(.+?)` is undefined', message)
-                    if match:
-                        undefineds.append((match.group(1), source))
-                    else:
-                        all_errors.append(errors.RelQueryError(problem, source))
-                elif "overflowed" in report:
-                    all_errors.append(errors.NumericOverflow(problem, source))
-                elif code == "PYREL_ERROR":
-                    pyrel_errors[problem["props"]["pyrel_id"]].append(problem)
-                elif abort_on_error:
-                    all_errors.append(errors.RelQueryError(problem, source))
-            else:
-                if code == "ARITY_MISMATCH":
-                    errors.ArityMismatch(problem, source)
-                elif code == "IC_VIOLATION":
-                    all_errors.append(errors.IntegrityConstraintViolation(problem, source))
-                elif code == "PYREL_ERROR":
-                    pyrel_warnings[problem["props"]["pyrel_id"]].append(problem)
-                else:
-                    errors.RelQueryWarning(problem, source)
-
-        if abort_on_error and len(undefineds):
-            all_errors.append(errors.UninitializedPropertyException(undefineds))
-
-        if abort_on_error:
-            for pyrel_id, pyrel_problems in pyrel_errors.items():
-                all_errors.append(errors.ModelError(pyrel_problems))
-
-        for pyrel_id, pyrel_problems in pyrel_warnings.items():
-            errors.ModelWarning(pyrel_problems)
-
-
-        if len(all_errors) == 1:
-            raise all_errors[0]
-        elif len(all_errors) > 1:
-            raise errors.RAIExceptionSet(all_errors)
-
-    def _export(self, raw_code: str, dest: Table, actual_cols: list[str], declared_cols: list[str], update:bool, headers: dict[str, Any] | None = None):
-        # _export is Snowflake-specific and requires Snowflake Resources
-        # It calls Snowflake stored procedures (APP_NAME.api.exec_into_table, etc.)
-        # LocalResources doesn't support this functionality
-        from relationalai.clients.local import LocalResources
-        if isinstance(self.resources, LocalResources):
-            raise NotImplementedError("Export functionality is not supported in local mode. Use Snowflake Resources instead.")
-        
-        _exec = self.resources._exec
-        output_table = "out" + str(uuid.uuid4()).replace("-", "_")
-        txn_id = None
-        artifacts = None
-        dest_database, dest_schema, dest_table, _ = IdentityParser(dest._fqn, require_all_parts=True).to_list()
-        dest_fqn = dest._fqn
-        assert self.resources._session  # All Snowflake Resources have _session
-        with debugging.span("transaction"):
-            try:
-                with debugging.span("exec_format") as span:
-                    res = _exec(f"call {APP_NAME}.api.exec_into_table(?, ?, ?, ?, ?, NULL, ?, {headers}, ?, ?);", [self.database, self.engine, raw_code, output_table, False, True, False, None])
-                    txn_id = json.loads(res[0]["EXEC_INTO_TABLE"])["rai_transaction_id"]
-                    span["txn_id"] = txn_id
-
-                with debugging.span("write_table", txn_id=txn_id):
-                    out_sample = _exec(f"select * from {APP_NAME}.results.{output_table} limit 1;")
-                    if out_sample:
-                        keys = set([k.lower() for k in out_sample[0].as_dict().keys()])
-                        fields = []
-                        ix = 0
-                        for name in declared_cols:
-                            if name in actual_cols:
-                                field = f"col{ix:03} as \"{name}\"" if f"col{ix:03}" in keys else f"NULL as {name}"
-                                ix += 1
-                            else:
-                                field = f"NULL as \"{name}\""
-                            fields.append(field)
-                        names = ", ".join(fields)
-                        if not update:
-
-                            createTableLogic = f"""
-                                        CREATE TABLE {dest_fqn} AS
-                                        SELECT {names}
-                                        FROM {APP_NAME}.results.{output_table};
-                            """
-                            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.{output_table};
-                                """
-
-                            _exec(f"""
-                                BEGIN
-                                    -- Check if table exists
-                                    IF (EXISTS (
-                                        SELECT 1
-                                        FROM {dest_database}.INFORMATION_SCHEMA.TABLES
-                                        WHERE table_schema = '{dest_schema}'
-                                        AND table_name = '{dest_table}'
-                                    )) THEN
-                                        -- Insert into existing table
-                                        EXECUTE IMMEDIATE '
-                                            BEGIN
-                                                TRUNCATE TABLE {dest_fqn};
-                                                INSERT INTO {dest_fqn}
-                                                SELECT {names}
-                                                FROM {APP_NAME}.results.{output_table};
-                                            END;
-                                        ';
-                                    ELSE
-                                        -- Create table based on the SELECT
-                                        EXECUTE IMMEDIATE '
-                                            {escape_for_f_string(createTableLogic)}
-                                        ';
-                                    END IF;
-                                    CALL {APP_NAME}.api.drop_result_table('{output_table}');
-                                END;
-                            """)
-                        else:
-                            _exec(f"""
-                                begin
-                                    INSERT INTO {dest_fqn} SELECT {names} FROM {APP_NAME}.results.{output_table};
-                                    call {APP_NAME}.api.drop_result_table('{output_table}');
-                                end;
-                            """)
-            except Exception as e:
-                msg = str(e).lower()
-                if "no columns returned" in msg or "columns of results could not be determined" in msg:
-                    if not update:
-                        # TODO: this doesn't handle a case when we have empty result, table doesn't exists and we need to create it.
-                        #   To handle it we need to get mapping from the export columns like `col000` with the SQL type to a real column name.
-                        _exec(f"""
-                            BEGIN
-                                IF (EXISTS (
-                                    SELECT 1
-                                    FROM {dest_database}.INFORMATION_SCHEMA.TABLES
-                                    WHERE table_schema = '{dest_schema}'
-                                    AND table_name = '{dest_table}'
-                                )) THEN
-                                    EXECUTE IMMEDIATE 'TRUNCATE TABLE {dest_fqn}';
-                                END IF;
-                            END;
-                        """)
-                else:
-                    raise e
-            if txn_id:
-                # These methods are available on all Snowflake Resources
-                artifact_info = self.resources._list_exec_async_artifacts(txn_id, headers=headers)
-                with debugging.span("fetch"):
-                    artifacts = self.resources._download_results(artifact_info, txn_id, "ABORTED")
-            return artifacts
-
-    def execute(self, model: ir.Model, task: ir.Task, format: Literal["pandas", "snowpark"] = "pandas",
-                export_to: Optional[Table] = None, update: bool = False, meta: dict[str, Any] | None = None) -> Any:
-        # Format meta as headers
-        json_meta = prepare_metadata_for_headers(meta)
-        headers = {QUERY_ATTRIBUTES_HEADER: json_meta} if json_meta else {}
-
-        self.check_graph_index(headers)
-        resources= self.resources
-
-        rules_code = ""
-        if self._last_model != model:
-            with debugging.span("compile", metamodel=model) as install_span:
-                install_span["compile_type"] = "model"
-                base = textwrap.dedent("""
-                    declare pyrel_error_attrs(err in ::std::common::UInt128, attr in ::std::common::String, v) requires true
-
-                """)
-                rules_code = base + self.compiler.compile(model, {"wide_outputs": self.wide_outputs})
-                install_span["rel"] = rules_code
-                rules_code = resources.create_models_code([("pyrel_qb_0", rules_code)])
-                self._last_model = model
-
-
-        with debugging.span("compile", metamodel=task) as compile_span:
-            compile_span["compile_type"] = "query"
-            base = textwrap.dedent("""
-                def output(:pyrel_error, err, attr, val):
-                    pyrel_error_attrs(err, attr, val)
-
-            """)
-            task_model = f.compute_model(f.logical([task]))
-            task_code, task_model = self.compiler.compile_inner(task_model, {"no_declares": True, "wide_outputs": self.wide_outputs})
-            task_code = base + task_code
-            compile_span["rel"] = task_code
-
-        full_code = textwrap.dedent(f"""
-            {rules_code}
-            {task_code}
-        """)
-
-        if self.dry_run:
-            return DataFrame()
-
-        cols, extra_cols = self._compute_cols(task, task_model)
-
-        if not export_to:
-            if format == "pandas":
-                raw_results = resources.exec_raw(self.database, self.engine, full_code, False, nowait_durable=True, headers=headers)
-                df, errs = result_helpers.format_results(raw_results, None, cols, generation=Generation.QB)  # Pass None for task parameter
-                self.report_errors(errs)
-                # Rename columns if wide outputs is enabled
-                if self.wide_outputs and len(cols) - len(extra_cols) == len(df.columns):
-                    df.columns = cols[: len(df.columns)]
-
-                return self._postprocess_df(self.config, df, extra_cols)
-            elif format == "snowpark":
-                results, raw = resources.exec_format(self.database, self.engine, full_code, cols, format=format, readonly=False, nowait_durable=True, headers=headers)
-                if raw:
-                    df, errs = result_helpers.format_results(raw, None, cols, generation=Generation.QB)  # Pass None for task parameter
-                    self.report_errors(errs)
-
-                return results
-        else:
-            assert cols
-            # The result cols should be a superset of the actual cols.
-            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:
-                result_cols = [col for col in cols if col not in extra_cols]
-            assert result_cols
-            raw = self._export(full_code, export_to, cols, result_cols, update, headers)
-            errors = []
-            if raw:
-                dataframe, errors = result_helpers.format_results(raw, None, result_cols, generation=Generation.QB)
-                self.report_errors(errors)
-            return DataFrame()
-
-    # NOTE(coey): this is added temporarily to support executing Rel for the solvers library in EA.
-    # It can be removed once this is no longer needed by the solvers library.
-    def execute_raw(self, raw_rel:str, readonly:bool=True, query_timeout_mins:int|None=None) -> DataFrame:
-        raw_results = self.resources.exec_raw(self.database, self.engine, raw_rel, readonly, nowait_durable=True, query_timeout_mins=query_timeout_mins)
-        df, errs = result_helpers.format_results(raw_results, None, generation=Generation.QB)  # Pass None for task parameter
-        self.report_errors(errs)
-        return df