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

relationalai/semantics/sql/compiler.py

@@ -66,6 +66,22 @@ class RelationInfo:
     table_selects: list[sql.Select] = field(default_factory=list)
     dynamic_table_selects: list[sql.Select] = field(default_factory=list)
 
+@dataclass
+class ImportSpec:
+    value: str
+    module: Optional[str] = None  # e.g., "scipy.special"
+
+    def render(self) -> str:
+        return f"from {self.module} import {self.value}" if self.module else f"import {self.value}"
+
+
+@dataclass
+class UDFConfig:
+    handler: str
+    code: str
+    imports: list[ImportSpec] = field(default_factory=list)
+    packages: list[str] = field(default_factory=list)
+
 @dataclass
 class ModelToSQL:
     """ Generates SQL from an IR Model, assuming the compiler rewrites were done. """
@@ -89,9 +105,9 @@ class ModelToSQL:
         return sql.Program(self._sort_dependencies(self._union_output_selects(self._generate_statements(model))))
 
     def _generate_statements(self, model: ir.Model) -> list[sql.Node]:
-        relations = self._get_relations(model)
+        table_relations, used_builtins = self._get_relations(model)
 
-        self._register_relation_args(relations)
+        self._register_relation_args(table_relations)
         self._register_external_relations(model)
 
         statements: list[sql.Node] = []
@@ -139,11 +155,15 @@ class ModelToSQL:
                 )
 
         # 4. Create physical tables for explicitly declared table relations
-        for relation in relations:
+        for relation in table_relations:
             info = self.relation_infos.get(relation)
             if info is None or info.table_selects:
                 statements.append(self._create_table(relation))
 
+        #5. Create Snowflake user-defined functions
+        if not self._is_duck_db:
+            statements.extend(self._create_user_defined_functions(used_builtins))
+
         return statements
 
     #--------------------------------------------------
@@ -204,6 +224,62 @@ class ModelToSQL:
             )
         )
 
+    def _create_user_defined_functions(self, relations: list[ir.Relation]) -> list[sql.CreateFunction]:
+        # Central UDF metadata configuration
+        udf_relations: dict[str, UDFConfig] = {
+            builtins.acot.name: UDFConfig(
+                handler="compute",
+                imports=[ImportSpec("math")],
+                code="""def compute(x): return math.atan(1 / x) if x != 0 else math.copysign(math.pi / 2, x)"""
+            ),
+            builtins.erf.name: UDFConfig(
+                handler="compute",
+                imports=[ImportSpec("math")],
+                code="""def compute(x): return math.erf(x)"""
+            ),
+            builtins.erfinv.name: UDFConfig(
+                handler="compute",
+                imports=[ImportSpec("erfinv", module="scipy.special")],
+                packages=["'scipy'"],
+                code="""def compute(x): return erfinv(x)"""
+            )
+        }
+
+        statements: list[sql.CreateFunction] = []
+
+        for r in relations:
+            meta = udf_relations.get(r.name)
+            if not meta:
+                continue
+
+            # Split relation fields into inputs and return type
+            # We expect a single return argument per builtin relation
+            return_type = None
+            input_columns: list[sql.Column] = []
+            for f in r.fields:
+                if f.input:
+                    input_columns.append(sql.Column(self._var_name(r.id, f), self._convert_type(f.type)))
+                else:
+                    return_type = self._convert_type(f.type)
+
+            # Build a full code block (imports + code)
+            imports_code = "\n".join(imp.render() for imp in meta.imports)
+            python_block = "\n".join(part for part in (imports_code, meta.code) if part)
+
+            assert return_type, f"No return type found for relation '{r.name}'"
+            statements.append(
+                sql.CreateFunction(
+                    name=r.name,
+                    inputs=input_columns,
+                    return_type=return_type,
+                    handler=meta.handler,
+                    body=python_block,
+                    packages=meta.packages
+                )
+            )
+
+        return statements
+
     def _create_statement(self, task: ir.Logical):
 
         # TODO - improve the typing info to avoid these casts
@@ -1186,7 +1262,7 @@ class ModelToSQL:
                     # SF returns values in `""` and to avoid this, we need to cast it to `TEXT` type
                     part_expr = f"cast({table_sql_var}.value as TEXT)"
                     index_expr = f"({table_sql_var}.index + 1)" # SF is 0-based internally, adjust to it back
-                assert isinstance(index, ir.Var) and isinstance(part, ir.Var), "Third and forth arguments (index, part) must be variables"
+                assert isinstance(index, ir.Var) and isinstance(part, ir.Var), "Third and fourth arguments (index, part) must be variables"
                 builtin_vars[part] = part_expr
                 builtin_vars[index] = index_expr
             elif relation == builtins.range:
@@ -1205,7 +1281,7 @@ class ModelToSQL:
                     table_expr = f"LATERAL FLATTEN(input => ARRAY_GENERATE_RANGE({start}, ({stop} + 1), {step}))"
                     expr = f"{table_sql_var}.value"
                 table_expressions[table_sql_var] = table_expr
-                assert isinstance(result, ir.Var), "Forth argument (result) must be a variable"
+                assert isinstance(result, ir.Var), "Fourth argument (result) must be a variable"
                 builtin_vars[result] = f"{expr}"
             elif relation == builtins.cast:
                 assert len(args) == 3, f"Expected 3 args for `cast`, got {len(args)}: {args}"
@@ -1214,7 +1290,7 @@ class ModelToSQL:
                 assert isinstance(result, ir.Var), "Third argument (result) must be a variable"
 
                 builtin_vars[result] = original_raw
-            elif relation in (builtins.isnan, builtins.isinf):
+            elif relation in {builtins.isnan, builtins.isinf}:
                 arg_expr = self._var_to_expr(args[0], reference, resolve_builtin_var, var_to_construct)
                 expr = "cast('NaN' AS DOUBLE)" if relation == builtins.isnan else "cast('Infinity' AS DOUBLE)"
                 wheres.append(sql.Terminal(f"{arg_expr} = {expr}"))
@@ -1351,8 +1427,8 @@ class ModelToSQL:
                     builtin_vars[rhs] = lhs
                     date_period_var_type[rhs] = relation.name
                 elif relation in builtins.date_builtins:
-                    if relation in (builtins.date_add, builtins.date_subtract, builtins.datetime_add,
-                                    builtins.datetime_subtract):
+                    if relation in {builtins.date_add, builtins.date_subtract, builtins.datetime_add,
+                                    builtins.datetime_subtract}:
                         assert len(args) == 3, f"Expected 3 args for {relation}, got {len(args)}: {args}"
                         assert isinstance(rhs, ir.Var), f"Period variable must be `ir.Var`, got: {rhs}"
                         period = date_period_var_type[rhs]
@@ -1361,10 +1437,10 @@ class ModelToSQL:
                         left = self._var_to_expr(lhs, reference, resolve_builtin_var, var_to_construct)
 
                         if self._is_duck_db:
-                            op = "+" if relation in (builtins.date_add, builtins.datetime_add) else "-"
+                            op = "+" if relation in {builtins.date_add, builtins.datetime_add} else "-"
                             expr = f"({left} {op} {period_val} * interval 1 {period})"
                         else:
-                            sign = 1 if relation in (builtins.date_add, builtins.datetime_add) else -1
+                            sign = 1 if relation in {builtins.date_add, builtins.datetime_add} else -1
                             expr = f"dateadd({period}, ({sign} * {period_val}), {left})"
 
                         result_var = args[2]
@@ -1439,23 +1515,19 @@ class ModelToSQL:
                     rel_name = relation.name
                     left = self._var_to_expr(lhs, reference, resolve_builtin_var, var_to_construct)
                     if relation in builtins.math_unary_builtins:
-                        if rel_name in {builtins.acot.name, builtins.erf.name, builtins.erfinv.name}:
-                            # TODO: implement acot using atan or atan2, it needs to be handled as a union of two cases
-                            raise Exception(f"The function {rel_name} is not supported.")
-                        else:
-                            method = "ln" if rel_name == builtins.natural_log.name else rel_name
-                            sub_expr = left
-                            if rel_name == builtins.factorial.name and self._is_duck_db:
-                                # Factorial requires an integer operand in DuckDB
-                                sub_expr = f"{left}::INTEGER"
-                            elif rel_name == builtins.log10.name:
-                                # log10 is not supported, so we use log with base 10
-                                sub_expr = f"10, {left}"
-                                method = "log"
-                            expr = f"{method}({sub_expr})"
-                    elif rel_name in (builtins.minimum.name, builtins.maximum.name, builtins.trunc_div.name,
+                        method = "ln" if rel_name == builtins.natural_log.name else rel_name
+                        sub_expr = left
+                        if rel_name == builtins.factorial.name and self._is_duck_db:
+                            # Factorial requires an integer operand in DuckDB
+                            sub_expr = f"{left}::INTEGER"
+                        elif rel_name == builtins.log10.name:
+                            # log10 is not supported, so we use log with base 10
+                            sub_expr = f"10, {left}"
+                            method = "log"
+                        expr = f"{method}({sub_expr})"
+                    elif rel_name in {builtins.minimum.name, builtins.maximum.name, builtins.trunc_div.name,
                                       builtins.power.name, builtins.mod.name, builtins.pow.name,
-                                      builtins.log.name):
+                                      builtins.log.name}:
                         assert len(args) == 3, f"Expected 3 args for {relation}, got {len(args)}: {args}"
 
                         result_var = args[2]
@@ -1480,7 +1552,7 @@ class ModelToSQL:
                         f"but got `{type(result_var).__name__}`: {result_var}"
                     )
                     builtin_vars[result_var] = expr
-                elif relation in (builtins.parse_int64, builtins.parse_int128) and isinstance(rhs, ir.Var):
+                elif relation in {builtins.parse_int64, builtins.parse_int128} and isinstance(rhs, ir.Var):
                     builtin_vars[rhs] = self._var_to_expr(lhs, reference, resolve_builtin_var, var_to_construct, False)
                 elif helpers.is_from_cast(lookup) and isinstance(rhs, ir.Var):
                     # For the `from cast` relations we keep the raw var, and we will ground it later.
@@ -2071,21 +2143,27 @@ class ModelToSQL:
             return f"DECIMAL({base_type.precision},{base_type.scale})"
         raise Exception(f"Unknown built-in type: {t}")
 
-    def _get_relations(self, model: ir.Model) -> list[ir.Relation]:
+    def _get_relations(self, model: ir.Model) -> Tuple[list[ir.Relation], list[ir.Relation]]:
+        rw = ReadWriteVisitor()
+        model.accept(rw)
+
+        root = cast(ir.Logical, model.root)
+
+        # For query compilation exclude read-only tables because we do not need to declare `CREATE TABLE` statements
+        used_relations = rw.writes(root) if self._query_compilation else rw.writes(root) | rw.reads(root)
+
         # Filter only relations that require table creation
-        relations = [
-            r for r in model.relations
+        table_relations = [
+            r for r in used_relations
             if self._is_table_creation_required(r)
         ]
 
-        # Optionally exclude read-only tables
-        if self._query_compilation:
-            rw = ReadWriteVisitor()
-            model.accept(rw)
-            writable = rw.writes(cast(ir.Logical, model.root))
-            relations = [r for r in relations if r in writable]
+        used_builtins = [
+            r for r in rw.reads(root)
+            if builtins.is_builtin(r)
+        ]
 
-        return relations
+        return table_relations, used_builtins
 
     def _is_table_creation_required(self, r: ir.Relation) -> bool:
         """
@@ -2161,7 +2239,7 @@ class ModelToSQL:
 
     def _var_name(self, relation_id: int, arg: Union[ir.Var, ir.Field]):
         name = helpers.sanitize(self.relation_arg_name_cache.get_name((relation_id, arg.id), arg.name))
-        return f'"{name}"' if name.lower() in ("any", "order") else name
+        return f'"{name}"' if name.lower() in {"any", "order"} else name
 
     def _register_relation_args(self, relations: list[ir.Relation]):
         """
@@ -2308,6 +2386,7 @@ class ModelToSQL:
             3. Other statements except SELECT queries
             4. SELECT queries
         """
+        udfs = []
         create_tables = []
         need_sort: dict[str, list[Union[sql.Insert, sql.CreateView, sql.CreateDynamicTable]]] = defaultdict(list)
         updates = []
@@ -2327,12 +2406,14 @@ class ModelToSQL:
                 updates.append(statement)
             elif isinstance(statement, sql.Select):
                 selects.append(statement)
+            elif isinstance(statement, sql.CreateFunction):
+                udfs.append(statement)
             else:
                 miscellaneous_statements.append(statement)
 
         sorted_statements = self._sort_statements_dependency_graph(need_sort)
 
-        return create_tables + sorted_statements + updates + miscellaneous_statements + selects
+        return udfs + create_tables + sorted_statements + updates + miscellaneous_statements + selects
 
     @staticmethod
     def _sort_statements_dependency_graph(statements: dict[str, list[Union[sql.Insert, sql.CreateView, sql.CreateDynamicTable]]]) -> list[sql.Insert]:

relationalai/semantics/sql/executor/duck_db.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
 import duckdb
+import math
 from pandas import DataFrame
 from typing import Any, Union, Literal
+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
@@ -18,10 +20,29 @@ class DuckDBExecutor(e.Executor):
         """ Execute the SQL query directly. """
         if format != "pandas":
             raise ValueError(f"Unsupported format: {format}")
+
         connection = duckdb.connect()
+
+        # Register scalar functions
+        connection.create_function("erf", self.erf)
+        connection.create_function("acot", self.acot)
+        connection.create_function("erfinv", self.erfinv)
+
         try:
             sql, _ = self.compiler.compile(model, {"is_duck_db": True})
             arrow_table = connection.query(sql).fetch_arrow_table()
             return format_duckdb_columns(arrow_table.to_pandas(), arrow_table.schema)
         finally:
             connection.close()
+
+    @staticmethod
+    def erf(x: float) -> float:
+        return math.erf(x)
+
+    @staticmethod
+    def erfinv(x: float) -> float:
+        return special_erfinv(x)
+
+    @staticmethod
+    def acot(x: float) -> float:
+        return math.atan(1 / x) if x != 0 else math.copysign(math.pi / 2, x)

relationalai/semantics/sql/rewrite/denormalize.py

@@ -25,7 +25,6 @@ class Denormalize(c.Pass):
 
 @dataclass
 class OldDenormalize(visitor.Rewriter):
-    # TODO: use the new Rewriter when available.
 
     denormalized: dict[ir.Relation, ir.Relation] = field(default_factory=dict, init=False, hash=False, compare=False)
 
@@ -168,13 +167,12 @@ class OldDenormalize(visitor.Rewriter):
         """ Denormalize the relations that can be denormalized.
 
         Group together relations that are keyed by the same "entity". This method defines
-        entities as being types that haver a unary relation containing only that type. All
+        entities as being types that have a unary relation containing only that type. All
         relations whose first argument is this type are grouped together.
 
         Returns a tuple with 2 elements:
-        1. the new set of relations after denormalization
-        2. a dict from relations that were denormalized away to the new relation that took
-        its place.
+        1. The new set of relations after denormalization
+        2. A dict from relations that were denormalized away to the new relation that took its place.
         """
         new_relations = ordered_set()
         denormalized: dict[ir.Relation, ir.Relation] = dict()
@@ -184,7 +182,7 @@ class OldDenormalize(visitor.Rewriter):
         entity_relations: dict[ir.Type, ir.Relation] = dict()
 
         for r in relations:
-            if len(r.fields) == 1 and not types.is_builtin(r.fields[0].type):
+            if len(r.fields) == 1 and not types.is_builtin(r.fields[0].type) and not r.name == 'Error':
                 e = r.fields[0].type
                 entity_types.add(e)
                 entity_relations[e] = r

relationalai/semantics/sql/rewrite/recursive_union.py

@@ -49,9 +49,29 @@ class RecursiveUnion(c.Pass):
             new_body = [logical for logical in root_logical.body if logical not in recursive_logicals]
 
             # Step 4: Add unions for each recursive group
-            for group in recursive_groups.values():
-                if group:
-                    new_body.append(f.union(list(group)))
+            for rel_id, logical_group in recursive_groups.items():
+                split_group = ordered_set()
+
+                for logical in logical_group:
+                    # Count total ir.Update tasks in this logical
+                    update_count = sum(isinstance(t, ir.Update) for t in logical.body)
+
+                    # If there's only one, keep the original logical as-is
+                    if update_count == 1:
+                        split_group.add(logical)
+                        continue
+
+                    # Otherwise, keep only updates relevant to this relation (and non-update tasks)
+                    filtered_body = [
+                        t for t in logical.body
+                        if not isinstance(t, ir.Update) or t.relation.id == rel_id
+                    ]
+
+                    if filtered_body:
+                        split_group.add(f.logical(filtered_body))
+
+                if split_group:
+                    new_body.append(f.union(list(split_group)))
 
             return model.reconstruct(model.engines, model.relations, model.types, f.logical(new_body), model.annotations)
 

relationalai/semantics/sql/sql.py

@@ -55,6 +55,17 @@ class CreateView(Node):
     name: str
     query: Union[list[Select], CTE]
 
+@dataclass(frozen=True)
+class CreateFunction(Node):
+    name: str
+    inputs: list[Column]
+    return_type: str
+    body: str
+    language: str = "PYTHON"
+    runtime_version: str = "3.11"
+    handler: str = "compute"
+    packages: Optional[list[str]] = None
+
 @dataclass(frozen=True)
 class Insert(Node):
     table: str
@@ -286,6 +297,22 @@ class Printer(BasePrinter):
         elif isinstance(node, CreateView):
             self._print(f"CREATE VIEW {self._get_table_name(node.name)} AS ")
             self._print_query(indent, node, "UNION")
+        elif isinstance(node, CreateFunction):
+            self._print(f"CREATE OR REPLACE FUNCTION {self._get_table_name(node.name)} (")
+            self._join(node.inputs)
+            self._print_nl(")")
+            self._print_nl(f"RETURNS {node.return_type}")
+            self._print_nl(f"LANGUAGE {node.language}")
+            self._print_nl(f"RUNTIME_VERSION = '{node.runtime_version}'")
+            self._print_nl(f"HANDLER = '{node.handler}'")
+            if node.packages:
+                self._print("PACKAGES = (")
+                self._join(node.packages)
+                self._print_nl(")")
+            self._print_nl("AS ")
+            self._print_nl("$$")
+            self._print_nl(node.body)
+            self._print("$$;")
         elif isinstance(node, Insert):
             self._print(f"INSERT INTO {self._get_table_name(node.table)} ")
             if len(node.columns) > 0:

relationalai/semantics/std/__init__.py

@@ -3,7 +3,7 @@ from typing import Any
 
 from relationalai.semantics.internal import internal as i
 from .std import _Date, _DateTime, _Number, _String, _Integer, _make_expr
-from . import datetime, math, strings, decimals, integers, floats, pragmas, constraints
+from . import datetime, math, strings, decimals, integers, floats, pragmas, constraints, re
 
 def range(*args: _Integer) -> i.Expression:
     # supports range(stop), range(start, stop), range(start, stop, step)
@@ -43,6 +43,7 @@ __all__ = [
     "datetime",
     "math",
     "strings",
+    "re",
     "decimals",
     "integers",
     "floats",

relationalai/semantics/std/datetime.py

@@ -153,6 +153,10 @@ class datetime:
                           std.cast_to_int64(day), std.cast_to_int64(hour), std.cast_to_int64(minute),
                           std.cast_to_int64(second), std.cast_to_int64(millisecond), tz, b.DateTime.ref("res"))
 
+    @classmethod
+    def now(cls) -> b.Expression:
+        return _make_expr("datetime_now", b.DateTime.ref("res"))
+
     @classmethod
     def year(cls, datetime: _DateTime, tz: dt.tzinfo|_String|None = None) -> b.Expression:
         tz = _extract_tz(datetime, tz)

relationalai/semantics/std/re.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+from relationalai.semantics.internal import internal as i
+from relationalai.semantics.metamodel.util import OrderedSet
+from .std import _Integer, _String, _make_expr
+from typing import Literal, Any
+from .. import std
+
+
+def escape(regex: _String) -> i.Expression:
+    return _make_expr("escape_regex_metachars", regex, i.String.ref())
+
+class Match(i.Producer):
+
+    def __init__(self, regex: _String, string: _String, pos: _Integer = 0, _type: Literal["search", "fullmatch", "match"] = "match"):
+        super().__init__(i.find_model([regex, string, pos]))
+        self.regex = regex
+        self.string = string
+        self.pos = pos
+
+        if _type == "match":
+            self._expr = _regex_match_all(self.regex, self.string, std.cast_to_int64(self.pos + 1))
+            self._offset, self._full_match = self._expr._arg_ref(2), self._expr._arg_ref(3)
+        elif _type == "search":
+            raise NotImplementedError("`search` is not implemented")
+        elif _type == "fullmatch":
+            _exp = _regex_match_all(self.regex, self.string, std.cast_to_int64(self.pos + 1))
+            self._offset, self._full_match = _exp._arg_ref(2), _exp._arg_ref(3)
+            self._expr = self._full_match == std.strings.substring(self.string, std.cast_to_int64(self.pos), std.strings.len(self.string))
+
+    def group(self, index: _Integer = 0) -> i.Producer:
+        if index == 0:
+            return self._full_match
+        else:
+            return _make_expr("capture_group_by_index", self.regex, self.string, std.cast_to_int64(self.pos + 1), std.cast_to_int64(index), i.String.ref("res"))
+
+    def group_by_name(self, name: _String) -> i.Producer:
+        return _make_expr("capture_group_by_name", self.regex, self.string, std.cast_to_int64(self.pos + 1), name, i.String.ref("res"))
+
+    def start(self) -> i.Expression:
+        return self._offset - 1
+
+    def end(self) -> i.Expression:
+        return std.strings.len(self.group(0)) + self.start() - 1
+
+    def span(self) -> tuple[i.Producer, i.Producer]:
+        return self.start(), self.end()
+
+    def _to_keys(self) -> OrderedSet[Any]:
+        return i.find_keys(self._expr)
+
+    def _compile_lookup(self, compiler:i.Compiler, ctx:i.CompilerContext):
+        compiler.lookup(self.regex, ctx)
+        compiler.lookup(self.string, ctx)
+        compiler.lookup(self.pos, ctx)
+        return compiler.lookup(self._expr, ctx)
+
+    def __getattr__(self, name: str) -> Any:
+        return object.__getattribute__(self, name)
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        object.__setattr__(self, name, value)
+
+
+def match(regex: _String, string: _String) -> Match:
+    return Match(regex, string)
+
+def search(regex: _String, string: _String, pos: _Integer = 0) -> Match:
+    return Match(regex, string, pos, _type="search")
+
+def fullmatch(regex: _String, string: _String, pos: _Integer = 0) -> Match:
+    return Match(regex, string, pos, _type="fullmatch")
+
+def findall(regex: _String, string: _String) -> tuple[i.Producer, i.Producer]:
+    exp = _regex_match_all(regex, string)
+    ix, match = exp._arg_ref(2), exp._arg_ref(3)
+    rank = i.rank(i.asc(ix, match))
+    return rank, match
+
+def _regex_match_all(regex: _String, string: _String, pos: _Integer|None = None) -> i.Expression:
+    if pos is None:
+        pos = i.Int64.ref()
+    return _make_expr("regex_match_all", regex, string, pos, i.String.ref())

relationalai/semantics/std/strings.py

@@ -15,7 +15,7 @@ def concat(s0: _String, s1: _String, *args: _String) -> b.Expression:
     return res
 
 def len(s: _String) -> b.Expression:
-    return _make_expr("num_chars", s, b.Int128.ref("res"))
+    return _make_expr("num_chars", s, b.Int64.ref("res"))
 
 def startswith(s0: _String, s1: _String) -> b.Expression:
     return _make_expr("starts_with", s0, s1)