relationalai 0.13.5__py3-none-any.whl → 1.0.0a1__py3-none-any.whl

v0/relationalai/semantics/std/re.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+from v0.relationalai.semantics.internal import internal as i
+from v0.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())

v0/relationalai/semantics/std/std.py

@@ -0,0 +1,14 @@
+from decimal import Decimal as PyDecimal
+import datetime as dt
+from typing import Any, Union
+
+from v0.relationalai.semantics.internal import internal as b
+
+_String = Union[b.Producer, str]
+_Integer = Union[b.Producer, int]
+_Date = Union[b.Producer, dt.date]
+_DateTime = Union[b.Producer, dt.datetime]
+_Number = Union[b.Producer, float, int, PyDecimal]
+
+def _make_expr(op: str, *args: Any) -> b.Expression:
+    return b.Expression(b.Relationship.builtins[op], *args)

v0/relationalai/semantics/std/strings.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+from typing import Union, Sequence
+
+from v0.relationalai.semantics.internal import internal as b
+from .std import _String, _Integer, _make_expr
+
+# TODO(coey) can we allow multiple args to this, and convert each to string and concatenate? similar to e.g. julia `string("foo", 1, 'b')`. then we don't need a separate concat function (and concat only works on strings, which is inconvenient).
+def string(s: Union[b.Producer, str, float, int]) -> b.Expression:
+    return _make_expr("string", s, b.String.ref("res"))
+
+def concat(s0: _String, s1: _String, *args: _String) -> b.Expression:
+    res = _make_expr("concat", s0, s1, b.String.ref("res0"))
+    for i, s in enumerate(args):
+        res = _make_expr("concat", res, s, b.String.ref(f"res{i+1}"))
+    return res
+
+def len(s: _String) -> b.Expression:
+    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)
+
+def endswith(s0: _String, s1: _String) -> b.Expression:
+    return _make_expr("ends_with", s0, s1)
+
+def contains(s0: _String, s1: _String) -> b.Expression:
+    return _make_expr("contains", s0, s1)
+
+def substring(s: _String, start: _Integer, stop: _Integer) -> b.Expression:
+    # unlike Python, Rel's range is 1..stop inclusive
+    return _make_expr("substring", s, start+1, stop, b.String.ref("res"))
+
+def like(s: _String, pattern: _String) -> b.Expression:
+    return _make_expr("like_match", s, pattern)
+
+def lower(s: _String) -> b.Expression:
+    return _make_expr("lower", s, b.String.ref("res"))
+
+def upper(s: _String) -> b.Expression:
+    return _make_expr("upper", s, b.String.ref("res"))
+
+def strip(s: _String) -> b.Expression:
+    return _make_expr("strip", s, b.String.ref("res"))
+
+def levenshtein(s: _String, t: _String) -> b.Expression:
+    return _make_expr("levenshtein", s, t, b.Int64.ref("res"))
+
+def join(strs: Sequence[_String], separator: _String) -> b.Expression:
+    return _make_expr("join", b.TupleArg(strs), separator, b.String.ref("res"))
+
+def replace(s: _String, old: _String, new: _String) -> b.Expression:
+    return _make_expr("replace", s, old, new, b.String.ref("res"))
+
+def split(s: _String, separator: _String) -> tuple[b.Producer, b.Producer]:
+    exp = _make_expr("split", separator, s, b.Int64.ref("idx"), b.String.ref("res"))
+    # indexes are 0-based
+    return exp._arg_ref(2) - 1, exp._arg_ref(3)
+
+def split_part(s: _String, separator: _String, idx: _Integer) -> b.Expression:
+    return _make_expr("split_part", separator, s, idx + 1, b.String.ref("res"))
+
+def regex_match(s:_String, regex: _String) -> b.Expression:
+    return _make_expr("regex_match", regex, s)

v0/relationalai/semantics/tests/test_snapshot_abstract.py

@@ -0,0 +1,143 @@
+import logging
+import os
+import uuid
+import sys
+import datetime
+from abc import abstractmethod, ABC
+
+from v0.relationalai.semantics.tests.logging import Capturer
+from v0.relationalai.semantics.tests.utils import reset_state
+from v0.relationalai.semantics.internal import internal
+from v0.relationalai.clients.result_helpers import sort_data_frame_result
+from v0.relationalai.clients.util import IdentityParser
+from v0.relationalai.clients.snowflake import Provider as SFProvider
+from v0.relationalai import Provider
+from typing import cast, Dict, Union
+from pathlib import Path
+
+class AbstractSnapshotTest(ABC):
+
+    provider:Provider = cast(SFProvider, Provider()) # type: ignore
+
+    def run_snapshot_test(self, snapshot, script_path, db_schema=None, use_sql=False, use_lqp=True, use_rel=False,
+                          use_direct_access=False, e2e=False, use_csv=True, e2e_only=False):
+        # Resolve use_lqp
+        use_lqp = use_lqp and (not use_rel) # use_rel overrides because use_lqp is default.
+
+        """Run a snapshot test.
+
+        This method assumes content is already in the correct string format.
+        Subclasses should handle any necessary compilation or conversion.
+        """
+        logger = logging.getLogger("pyrellogger")
+        handler = Capturer()
+        logger.addHandler(handler)
+
+        # reset state across test runs
+        reset_state()
+
+        # Track which modules are already loaded
+        before_modules = set(sys.modules.keys())
+        unique_name = f"md{str(uuid.uuid4())[-12:]}"
+        export_fqn = None
+        dry_run = not e2e
+        try:
+            overrides = {
+                'dry_run': dry_run,
+                'model_suffix': "" if not e2e else f"_{unique_name}",
+                'use_sql': use_sql,
+                'reasoner.rule.use_lqp': use_lqp,
+                'keep_model': False,
+                # fix the current time to keep snapshots stable
+                'datetime_now': datetime.datetime.fromisoformat("2025-12-01T12:00:00+00:00"),
+            }
+            if use_direct_access:
+                # for direct access we can not use user/password authentication
+                # so we override the config to use key pair authentication
+                override_config: Dict[str, Union[str, bool]] = {
+                    'use_direct_access': True,
+                }
+                if os.getenv("AUTHENTICATOR"):
+                    override_config["authenticator"] = os.getenv("AUTHENTICATOR", "")
+                if os.getenv("PRIVATE_KEY_FILE"):
+                    override_config["private_key_file"] = os.getenv("PRIVATE_KEY_FILE", "")
+                if os.getenv("SF_TEST_ACCOUNT_KEY_PASSPHRASE"):
+                    override_config["private_key_file_pwd"] = os.getenv("SF_TEST_ACCOUNT_KEY_PASSPHRASE", "")
+                overrides['config'] = override_config
+
+            with internal.with_overrides(**overrides):
+                with open(script_path, "r") as f:
+                    # Compile with the correct filename
+                    code = compile(f.read(), script_path, "exec")
+                    # Execute the compiled code
+                    if db_schema:
+                        export_fqn = f"{db_schema}.{unique_name}"
+                    exec_globals = {"EXPORT_TABLE": export_fqn}
+                    exec(code, exec_globals)
+        finally:
+            # Remove any modules that were imported during exec()
+            after_modules = set(sys.modules.keys())
+            new_modules = after_modules - before_modules
+            for mod in new_modules:
+                sys.modules.pop(mod, None)
+            # Remove handler
+            logger.removeHandler(handler)
+            # check if script created an export table
+            export_exists = not dry_run and self.is_export_exists(export_fqn)
+            # match snapshots
+            self.assert_match_snapshots(script_path, snapshot, handler, export_exists, export_fqn, use_sql=use_sql,
+                                        use_lqp=use_lqp, e2e=e2e, use_csv=use_csv, e2e_only=e2e_only)
+            # cleanup resources created during a test run
+            self.cleanup(export_exists, export_fqn)
+
+    def assert_match_snapshots(self, script_path, snapshot, handler, export_exists=False, export_fqn=None, use_sql=False,
+                               use_lqp=True, e2e=False, use_csv=True, e2e_only=False):
+        if not e2e_only:
+            self.assert_match_internal_results_snapshots(snapshot, handler, use_sql, use_lqp)
+        if e2e:
+            self.assert_match_results_snapshots(script_path, snapshot, handler)
+        if export_fqn and export_exists:
+            result_df = sort_data_frame_result(self.provider.sql(f"SELECT * FROM {export_fqn}", format="pandas"))
+            if use_csv:
+                snapshot.assert_match(result_df.to_csv(index=False).strip('"').strip('\n"'), "export_result.csv")
+            else:
+                snapshot.assert_match(result_df.to_string(), "export_result.df")
+
+    def assert_match_internal_results_snapshots(self, snapshot, handler, use_sql, use_lqp):
+        # default: do nothing
+        return
+
+    @abstractmethod
+    def assert_match_results_snapshots(self, script_path, snapshot, handler):
+        pass
+
+    def cleanup(self, export_exists=False, export_fqn=None):
+        if export_fqn and export_exists :
+            self.provider.sql(f"DROP TABLE IF EXISTS {export_fqn}")
+
+    def is_export_exists(self, export_fqn=None):
+        if export_fqn is None:
+            return False
+        database, schema, table, _ = IdentityParser(export_fqn, require_all_parts=True).to_list()
+        return self.provider.sql(f"""
+            SELECT EXISTS(
+                SELECT 1
+                FROM {database}.INFORMATION_SCHEMA.TABLES
+                WHERE TABLE_SCHEMA = '{schema}'
+                AND TABLE_NAME = '{table}'
+            ) AS TABLE_EXISTS;
+        """)[0]["TABLE_EXISTS"]
+
+    @staticmethod
+    def parent_tests(file: str):
+        return AbstractSnapshotTest.discover_test_scripts(Path(file).parent)
+
+    @staticmethod
+    def discover_test_scripts(path: Path):
+        scripts_dir = path / "tests"
+        return sorted(scripts_dir.glob("**/*.py"))
+
+    @staticmethod
+    def get_id(p):
+        p_name = p.parent.name
+        return p.stem if p_name == "tests" else f"{p_name}_{p.stem}"

v0/relationalai/semantics/tests/test_snapshot_base.py

@@ -0,0 +1,9 @@
+from v0.relationalai.semantics.tests.test_snapshot_abstract import AbstractSnapshotTest
+
+class SnapshotTestBase(AbstractSnapshotTest):
+    """Base class for tests with snapshots."""
+
+    def assert_match_results_snapshots(self, script_path, snapshot, handler):
+        for i, msg in enumerate(handler.results, start=1):
+            if "results" in msg:
+                snapshot.assert_match(msg["results"].to_csv(index=False).strip('"').strip('\n"'), f"result_{i}.csv")

v0/relationalai/semantics/tests/utils.py

@@ -0,0 +1,46 @@
+from itertools import count
+from v0.relationalai.semantics.metamodel import ir, types, util
+from v0.relationalai.semantics.internal import internal as b
+from v0.relationalai.semantics.snowflake import Table
+
+def reset_state():
+    """
+    Reset global state for consistent test snapshots.
+
+    When we execute a pyrel program we accumulate some state, such as custom decimals, and
+    we increase the counter for object ids. This function resets those counters and other
+    state to a known baseline to ensure that test snapshots are consistent if we run the
+    test alone vs in a suite of tests.
+    """
+    # reset the global id counters
+    ir._global_id = count(10000)
+    b._global_id = count(10000)
+    b._global_roots = util.ordered_set()
+
+    # Used to generate error source IDs
+    b.errors.ModelError.error_locations.clear()
+
+    # reset the ErrorConcept state, so that Error.new() always generates the same IR
+    b.ErrorConcept._error_props.clear()
+    b.ErrorConcept._relation = None
+    b.ErrorConcept._overloads.clear()
+
+    # cleanup relationships declared on global concepts
+    concepts_to_cleaanup = [b.Concept.builtins["Error"], b.Concept.builtins["Any"]]
+    keep = ["shape", "builtins", "globals"]
+    for concept in concepts_to_cleaanup:
+        for attr in list(concept._relationships.keys()):
+            if attr not in keep:
+                del concept._relationships[attr]
+        concept.globals.clear()
+
+    # caches of custom decimals
+    for k in list(types._decimal_types.keys()):
+        if types._decimal_types[k] != types.Decimal:
+            del types._decimal_types[k]
+    for k in list(b.Concept.builtins):
+        if k.startswith("Decimal(") and b.Concept.builtins[k] is not b.Decimal:
+            del b.Concept.builtins[k]
+
+    # clear any cached Table sources
+    Table._used_sources.clear()

v0/relationalai/std/__init__.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+import functools
+from typing import Any, Callable
+
+from v0.relationalai.errors import AsBoolForNonFilter
+from v0.relationalai.metamodel import Builtins
+from ..dsl import alias, rel, Vars, create_vars, create_var, tag
+from .. import dsl, metamodel as m
+from . import graphs, aggregates, strings, math, dates
+
+#--------------------------------------------------
+# Tagging some base rel functions
+#--------------------------------------------------
+
+tag(rel.decimal, Builtins.SingleValued)
+tag(rel.float, Builtins.SingleValued)
+tag(rel.int, Builtins.SingleValued)
+tag(rel.sqrt, Builtins.SingleValued)
+tag(rel.concat, Builtins.SingleValued)
+
+#--------------------------------------------------
+# Utilities
+#--------------------------------------------------
+
+def as_bool(expr: dsl.Expression) -> dsl.Expression:
+    if expr._expr.entity.isa(Builtins.Filter):
+        # add the filter to apply
+        prev_filter = expr._expr.entity.value
+        filter = dsl.build.property_named("filter", [])
+        expr._expr.append(filter, m.Var(value=prev_filter))
+        # add a result var
+        expr._var = m.Var(Builtins.Unknown)
+        prop = dsl.build.property_named("result", [])
+        expr._expr.append(prop, expr._var)
+        # use pyrel_bool_filter to apply the filter and get a bool
+        expr._expr.entity = m.Var(value=Builtins.BoolFilter)
+    else:
+        raise AsBoolForNonFilter()
+    return expr
+
+def as_rows(data:list[tuple|dict|int|float|str]) -> dsl.Rows:
+    return dsl.Rows(dsl.get_graph(), data)
+
+def _wrapped(relation: Callable, *args: Any) -> Callable:
+    def f(*args: Any):
+        res = create_var()
+        relation(*args, res)
+        return res
+    return f
+
+def minimum(arg1: Any, arg2: Any, *args: Any) -> dsl.Expression:
+    return functools.reduce(_wrapped(rel.minimum), [arg1, arg2, *args])
+
+def maximum(arg1: Any, arg2: Any, *args: Any) -> dsl.Expression:
+    return functools.reduce(_wrapped(rel.maximum), [arg1, arg2, *args])
+
+__all__ = [
+    "aggregates",
+    "alias",
+    "dates",
+    "graphs",
+    "math",
+    "minimum",
+    "maximum",
+    "rel",
+    "strings",
+    "Vars",
+    "create_vars",
+    "create_var"
+]