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

relationalai/semantics/rel/builtins.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 from relationalai.semantics.metamodel import types, factory as f
 from relationalai.semantics.metamodel.util import OrderedSet
+from relationalai.semantics.metamodel import builtins
 
 # Rel Annotations as IR Relations (to be used in IR Annotations)
 
@@ -28,7 +29,10 @@ inner_loop_non_stratified_annotation = f.annotation(inner_loop_non_stratified, [
 
 # collect all supported builtin rel annotations
 builtin_annotations = OrderedSet.from_iterable([
-    arrow, no_diagnostics, no_inline, function, inner_loop, inner_loop_non_stratified
+    arrow, no_diagnostics, no_inline, function, inner_loop, inner_loop_non_stratified,
+    # track annotations on relations do not currently propagate into Rel
+    # TODO: from Thiago, ensure annotation goes from the Logical into the proper declaration
+    builtins.track,
 ])
 
 builtin_annotation_names = OrderedSet.from_iterable([a.name for a in builtin_annotations])

relationalai/semantics/rel/compiler.py

@@ -6,13 +6,13 @@ from decimal import Decimal as PyDecimal
 from relationalai.semantics.metamodel import ir, compiler as c, builtins as bt, types, visitor, helpers, factory as f
 from relationalai.semantics.metamodel.typer import Checker, InferTypes
 from relationalai.semantics.metamodel.typer.typer import to_base_primitive, to_type, _NON_PARAMETRIC_PRIMITIVES
-from relationalai.semantics.metamodel.rewrite import (Flatten, ExtractKeys, FDConstraints, Splinter,
-                                                      ExtractNestedLogicals, DNFUnionSplitter, DischargeConstraints)
 from relationalai.semantics.metamodel.visitor import ReadWriteVisitor
 from relationalai.semantics.metamodel.util import OrderedSet, group_by, NameCache, ordered_set
 
 from relationalai.semantics.rel import rel, rel_utils as u, builtins as rel_bt
-from relationalai.semantics.rel.rewrite import CDC, QuantifyVars, ExtractCommon
+
+from ..metamodel.rewrite import (Flatten, ExtractNestedLogicals, DNFUnionSplitter, DischargeConstraints)
+from ..lqp.rewrite import CDC, ExtractCommon, ExtractKeys, FDConstraints, QuantifyVars, Splinter
 
 import math
 

relationalai/semantics/rel/executor.py

@@ -18,7 +18,8 @@ from relationalai.clients.snowflake import APP_NAME
 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 USE_DIRECT_ACCESS, Generation
+from relationalai.tools.constants import USE_DIRECT_ACCESS, Generation, QUERY_ATTRIBUTES_HEADER
+from relationalai.tools.query_utils import prepare_metadata_for_headers
 
 
 class RelExecutor(e.Executor):
@@ -60,7 +61,7 @@ class RelExecutor(e.Executor):
                         atexit.register(self._resources.delete_graph, self.database, True)
         return self._resources
 
-    def check_graph_index(self):
+    def check_graph_index(self, headers: dict[str, Any] | None = None):
         # Has to happen first, so self.dry_run is populated.
         resources = self.resources
 
@@ -84,7 +85,7 @@ class RelExecutor(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, sources, model, self.engine, engine_size, program_span_id, headers)
 
     def report_errors(self, problems: list[dict[str, Any]], abort_on_error=True):
         from relationalai import errors
@@ -143,7 +144,7 @@ class RelExecutor(e.Executor):
         elif len(all_errors) > 1:
             raise errors.RAIExceptionSet(all_errors)
 
-    def _export(self, raw_code: str, dest_fqn: str, actual_cols: list[str], declared_cols: list[str], update:bool):
+    def _export(self, raw_code: str, dest_fqn: str, actual_cols: list[str], declared_cols: list[str], update:bool, headers: dict[str, Any] | None = None):
         _exec = self.resources._exec
         output_table = "out" + str(uuid.uuid4()).replace("-", "_")
         txn_id = None
@@ -153,7 +154,7 @@ class RelExecutor(e.Executor):
         with debugging.span("transaction"):
             try:
                 with debugging.span("exec_format") as span:
-                    res = _exec(f"call {APP_NAME}.api.exec_into_table(?, ?, ?, ?, ?, ?);", [self.database, self.engine, raw_code, output_table, False, True])
+                    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
 
@@ -229,14 +230,18 @@ class RelExecutor(e.Executor):
                 else:
                     raise e
             if txn_id:
-                artifact_info = self.resources._list_exec_async_artifacts(txn_id)
+                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",
-                result_cols: list[str] | None = None, export_to: Optional[str] = None, update: bool = False) -> Any:
-        self.check_graph_index()
+                result_cols: list[str] | None = None, export_to: Optional[str] = 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 = ""
@@ -277,12 +282,16 @@ class RelExecutor(e.Executor):
 
         if not export_to:
             if format == "pandas":
-                raw_results = resources.exec_raw(self.database, self.engine, full_code, False, nowait_durable=True)
+                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)
+                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)
@@ -296,7 +305,7 @@ class RelExecutor(e.Executor):
             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)
+            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)

relationalai/semantics/sql/compiler.py

@@ -11,7 +11,7 @@ from decimal import Decimal as PyDecimal
 
 import math
 
-from relationalai.semantics.metamodel.rewrite import (Flatten, ExtractNestedLogicals, FDConstraints, DNFUnionSplitter,
+from relationalai.semantics.metamodel.rewrite import (Flatten, ExtractNestedLogicals, DNFUnionSplitter,
                                                       DischargeConstraints)
 from relationalai.semantics.metamodel.visitor import ReadWriteVisitor
 from relationalai.util.graph import topological_sort
@@ -28,7 +28,6 @@ from relationalai.semantics.sql import sql, rewrite
 class Compiler(c.Compiler):
     def __init__(self, skip_denormalization:bool=False):
         rewrites = [
-            FDConstraints(),
             DischargeConstraints(),
             Checker(),
             ExtractNestedLogicals(),  # before InferTypes to avoid extracting casts
@@ -2495,4 +2494,4 @@ class DerivedRelationsVisitor(v.Visitor):
 
     def visit_relation(self, node: ir.Relation, parent: Optional[ir.Node]):
         if self._is_derived and from_cdc_annotation in node.annotations:
-            self._is_derived = False
+            self._is_derived = False

relationalai/semantics/sql/executor/duck_db.py

@@ -8,7 +8,7 @@ from scipy.special import erfinv as special_erfinv
 
 from relationalai.semantics.sql import Compiler
 from relationalai.semantics.sql.executor.result_helpers import format_duckdb_columns
-from relationalai.semantics.metamodel import ir, executor as e
+from relationalai.semantics.metamodel import ir, executor as e, factory as f
 
 class DuckDBExecutor(e.Executor):
 
@@ -29,8 +29,12 @@ class DuckDBExecutor(e.Executor):
         connection.create_function("erfinv", self.erfinv)
 
         try:
-            sql, _ = self.compiler.compile(model, {"is_duck_db": True})
-            arrow_table = connection.query(sql).fetch_arrow_table()
+            model_sql, _ = self.compiler.compile(model, {"is_duck_db": True})
+            query_options = {"is_duck_db": True, "query_compilation": True}
+            query_sql, _ = self.compiler.compile(f.compute_model(f.logical([task])), query_options)
+
+            full_sql = model_sql + "\n" + query_sql
+            arrow_table = connection.query(full_sql).fetch_arrow_table()
             return format_duckdb_columns(arrow_table.to_pandas(), arrow_table.schema)
         finally:
             connection.close()
@@ -45,4 +49,4 @@ class DuckDBExecutor(e.Executor):
 
     @staticmethod
     def acot(x: float) -> float:
-        return math.atan(1 / x) if x != 0 else math.copysign(math.pi / 2, x)
+        return math.atan(1 / x) if x != 0 else math.copysign(math.pi / 2, x)

relationalai/semantics/sql/executor/snowflake.py

@@ -62,7 +62,7 @@ class SnowflakeExecutor(e.Executor):
 
     def execute(self, model: ir.Model, task: ir.Task, format:Literal["pandas", "snowpark"]="pandas",
                 result_cols: Optional[list[str]] = None, export_to: Optional[str] = None,
-                update: bool = False) -> Union[pd.DataFrame, Any]:
+                update: bool = False, meta: dict[str, Any] | None = None) -> Union[pd.DataFrame, Any]:
         """ Execute the SQL query directly. """
 
         warehouse = self.resources.config.get("warehouse", None)

relationalai/tools/cli.py

@@ -16,8 +16,10 @@ from InquirerPy.base.control import Choice
 from relationalai.clients.util import IdentityParser
 from .cli_controls import divider, Spinner
 from . import cli_controls as controls
-from relationalai.clients import azure
-from typing import Sequence, cast, Any, List
+from typing import Sequence, cast, Any, List, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from relationalai.clients import azure
 from relationalai.errors import RAIException
 from relationalai.loaders.types import LoadType, UnsupportedTypeError
 from relationalai.loaders.csv import CSVLoader
@@ -717,7 +719,10 @@ def config_check(all_profiles:bool=False):
 @cli.command(help="Print version info")
 def version():
     from .. import __version__
-    from railib import __version__ as railib_version
+    try:
+        from railib import __version__ as railib_version
+    except Exception:
+        railib_version = None
 
     table = Table(show_header=False, border_style="dim", header_style="bold", box=rich_box.SIMPLE)
     def print_version(name, version, latest=None):
@@ -728,7 +733,8 @@ def version():
 
     divider()
     print_version("RelationalAI", __version__, latest_version("relationalai"))
-    print_version("Rai-sdk", railib_version, latest_version("rai-sdk"))
+    if railib_version is not None:
+        print_version("Rai-sdk", railib_version, latest_version("rai-sdk"))
     print_version("Python", sys.version.split()[0])
 
     app_version = None
@@ -1150,9 +1156,14 @@ def import_source_flow(provider: ResourcesBase) -> Sequence[ImportSource]:
 
     if isinstance(provider, snowflake.Resources):
         return snowflake_import_source_flow(provider)
-    elif isinstance(provider, azure.Resources):
-        return azure_import_source_flow(provider)
     else:
+        # Lazy import for azure to avoid optional dependency issues
+        try:
+            from relationalai.clients import azure
+            if isinstance(provider, azure.Resources):
+                return azure_import_source_flow(provider)
+        except ImportError:
+            pass
         raise Exception(f"No import source flow available for {provider_type.__module__}.{provider_type.__name__}")
 
 def snowflake_import_source_flow(provider: snowflake.Resources) -> Sequence[ImportSource]: