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

relationalai/semantics/internal/snowflake.py

@@ -0,0 +1,325 @@
+from __future__ import annotations
+from collections import defaultdict
+from dataclasses import dataclass
+import json
+import re
+import textwrap
+from relationalai.clients.util import IdentityParser
+from relationalai.semantics.metamodel import types
+from relationalai.semantics.metamodel.util import OrderedSet, ordered_set
+from relationalai.semantics.rel.rel_utils import sanitize_identifier
+from . import internal as b, annotations as anns
+from relationalai import debugging
+from relationalai.errors import UnsupportedColumnTypesWarning
+from snowflake.snowpark.context import get_active_session
+from typing import ClassVar, Optional
+
+#--------------------------------------------------
+# Iceberg Configuration
+#--------------------------------------------------
+@dataclass
+class IcebergConfig:
+    """Configuration for exporting to Iceberg tables."""
+    external_volume: str | None = None
+    default: ClassVar[Optional["IcebergConfig"]]
+
+IcebergConfig.default = IcebergConfig()
+#--------------------------------------------------
+# Helpers
+#--------------------------------------------------
+
+def schema_dict_to_fields(schema: dict[str, str|b.Concept]) -> list[b.Field]:
+    fields = []
+    for name, type in schema.items():
+        if isinstance(type, str):
+            fields.append(b.Field(name=name, type_str=type))
+        elif isinstance(type, b.Concept):
+            fields.append(b.Field(name=name, type_str=type._name, type=type))
+    return fields
+
+#--------------------------------------------------
+# Constants
+#--------------------------------------------------
+
+# List of supported Snowflake column types. In V0 this comes from the app.get_supported_column_types() function.
+# We can't use that function in V1 since it would add a dependency on the app, so we hardcode the list here.
+SUPPORTED_SNOWFLAKE_TYPES = [
+    'CHAR', 'STRING', 'VARCHAR', 'BINARY', 'NUMBER', 'FLOAT', 'REAL',
+    'BOOLEAN', 'DATE', 'FIXED', 'TEXT', 'TIME', 'TIMESTAMP_LTZ',
+    'TIMESTAMP_NTZ', 'TIMESTAMP_TZ'
+]
+
+#--------------------------------------------------
+# Globals
+#--------------------------------------------------
+
+_session = None
+def get_session():
+    global _session
+    if not _session:
+        try:
+            _session = get_active_session()
+        except Exception:
+            from relationalai import Resources
+            from relationalai.clients.resources.snowflake import Resources as SnowflakeResources
+            # TODO: we need a better way to handle global config
+
+            # using the resource constructor to differentiate between direct access and
+            # regular resource as this impacts the construction of the session.
+            resources = Resources()
+            assert isinstance(resources, SnowflakeResources), "Can only retrieve session from Snowflake resources"
+            _session = resources.get_sf_session()
+    return _session
+
+def is_direct_access_enabled():
+    from relationalai import Resources
+    resources = Resources()
+    if hasattr(resources, "is_direct_access_enabled"):
+        return resources.is_direct_access_enabled()
+    return False
+
+
+#--------------------------------------------------
+# Snowflake sources
+#--------------------------------------------------
+
+SFTypes = {
+    "TEXT": "String",
+    "FIXED": "Number",
+    "DATE": "Date",
+    "TIME": "DateTime",
+    "TIMESTAMP": "DateTime",
+    "TIMESTAMP_LTZ": "DateTime",
+    "TIMESTAMP_TZ": "DateTime",
+    "TIMESTAMP_NTZ": "DateTime",
+    "FLOAT": "Float",
+    "REAL": "Float",
+    "BOOLEAN": "Bool",
+}
+
+SF_ID_REGEX = re.compile(r'^[A-Za-z_][A-Za-z0-9_$]*$')
+def quoted(ident:str):
+    if SF_ID_REGEX.match(ident) or ident[0] == '"':
+        return ident
+    return f'"{ident}"'
+
+@dataclass
+class TableInfo:
+    source:Table|None
+    fields:list[b.Field]
+    raw_columns:list[dict]
+
+class SchemaInfo:
+    def __init__(self, database:str, schema:str) -> None:
+        self.database = database
+        self.schema = schema
+        self.tables = defaultdict(lambda: TableInfo(None, [], []))
+        self.fetched = set()
+
+    def fetch(self):
+        session = get_session()
+        table_names = [name for name in self.tables.keys() if name.upper() not in self.fetched]
+        self.fetched.update([x.upper() for x in table_names])
+        name_lookup = {x.upper(): x for x in table_names}
+        tables = ", ".join([f"'{x.upper()}', '{x}'" for x in table_names])
+        supported_types_list = ", ".join([f"'{t}'" for t in SUPPORTED_SNOWFLAKE_TYPES])
+
+        query = textwrap.dedent(f"""
+            begin
+                SHOW COLUMNS IN SCHEMA {quoted(self.database)}.{quoted(self.schema)};
+                let r resultset := (select "table_name", "column_name", "data_type",
+                                    CASE
+                                        WHEN PARSE_JSON("data_type"):"type" in ({supported_types_list}) THEN TRUE
+                                        ELSE FALSE
+                                    END as "supported_type"
+                                    from table(result_scan(-1)) as t
+                                    where "table_name" in ({tables}));
+                return table(r);
+            end;
+        """)
+        with debugging.span("fetch_schema", sql=query):
+            columns = session.sql(query).collect()
+
+        # Collect unsupported columns for warning
+        unsupported_columns = {}
+
+        for row in columns:
+            table_name, column_name, data_type, supported_type = row
+            table_name = name_lookup.get(table_name, table_name)
+
+            # Skip columns with unsupported types
+            if not supported_type:
+                if table_name not in unsupported_columns:
+                    unsupported_columns[table_name] = {}
+                unsupported_columns[table_name][column_name] = data_type
+                continue
+
+            info = self.tables[table_name]
+            sf_type_info = json.loads(data_type)
+            typ = sf_type_info.get("type")
+            if typ in SUPPORTED_SNOWFLAKE_TYPES:
+                if typ == "FIXED":
+                    type_str = self.sf_numeric_to_type_str(column_name, sf_type_info)
+                else:
+                    type_str = SFTypes[typ]
+                info.fields.append(b.Field(name=column_name, type_str=type_str))
+                info.raw_columns.append(row.as_dict())
+            else:
+                # TODO - raise a warning
+                pass
+
+        # Show warning for unsupported columns
+        if unsupported_columns:
+            UnsupportedColumnTypesWarning(unsupported_columns)
+
+    def sf_numeric_to_type_str(self, column_name, sf_type_info):
+        """
+        Computes the appropriate type to use for this column. This code reflects exactly
+        the logic currently used by RAI's CDC implementation to ensure we map to the exact
+        same number types.
+        """
+        if "scale" not in sf_type_info or "precision" not in sf_type_info:
+            raise ValueError(f"Invalid definition for column '{column_name}': 'scale' or 'precision' missing")
+
+        precision = sf_type_info.get("precision")
+        scale = sf_type_info.get("scale")
+
+        if scale > precision or scale < 0 or scale > 37:
+            raise ValueError(f"Invalid numeric scale '{scale}' for column '{column_name}'")
+
+        if scale == 0:
+            # Integers (load_csv only supports these two (and not 8/16/32 bit ints)
+            # RAI int sizes are 8, 16, 32, 64, or 128; 64 by default
+            # SF precision is expressed in base 10 digits, max 38
+            # Max positive int value: 127, 32767, 2147483647, 9223372036854775807, 170141183460469231731687303715884105727
+            # Base 10 digits fully covered: 2, 4, 9, 18, 38
+            bits = types.digits_to_bits(precision)
+            return "Int128" if bits == 128 else "Int64"
+        else:
+            return f"Decimal({precision},{scale})"
+
+class Table():
+    _schemas:dict[tuple[str, str], SchemaInfo] = {}
+    _used_sources:OrderedSet[Table] = ordered_set()
+
+    def __init__(self, fqn:str, cols:list[str]|None=None, schema:dict[str, str|b.Concept]|None=None, config: IcebergConfig|None=None) -> None:
+        self._fqn = fqn
+        parser = IdentityParser(fqn, require_all_parts=True)
+        self._database, self._schema, self._table, self._fqn = parser.to_list()
+        self._inited = False
+        self._concept = b.Concept(self._fqn, extends=[b.Concept.builtins["RowId"]])
+        b.Concept.globals[self._fqn.lower()] = self._concept
+        self._ref = self._concept.ref("row_id")
+        self._cols = {}
+        self._col_names = cols
+        self._iceberg_config = config
+        self._is_iceberg = config is not None
+        info = self._schemas.get((self._database, self._schema))
+        if not info:
+            info = self._schemas[(self._database, self._schema)] = SchemaInfo(self._database, self._schema)
+        if schema:
+            info.fetched.add(self._table)
+            info.tables[self._table] = TableInfo(self, schema_dict_to_fields(schema), [])
+        info.tables[self._table].source = self
+
+    def _lazy_init(self):
+        if self._inited:
+            return
+        self._inited = True
+        schema_info = self._schemas[(self._database, self._schema)]
+        if self._table not in schema_info.fetched:
+            schema_info.fetch()
+        table_info = schema_info.tables[self._table]
+        CDC_name = self._fqn.lower() if '"' not in self._fqn else self._fqn.replace("\"", "_")
+        self._rel = b.Relationship(CDC_name, fields=[b.Field(name="RowId", type_str=self._fqn, type=self._concept)] + table_info.fields)
+        self._rel.annotate(anns.external).annotate(anns.from_cdc)
+
+    def __getattr__(self, name: str):
+        if name.startswith("_"):
+            return super().__getattribute__(name)
+        if self._col_names is not None and name not in self._col_names:
+            raise ValueError(f"Column {name} not found in {self._fqn}")
+        if name.lower() not in self._cols:
+            self._cols[name.lower()] = Column(self, name)
+        return self._cols[name.lower()]
+
+    def into(self, concept:b.Concept, keys:list[Column]=[]):
+        self._lazy_init()
+        key_dict = {sanitize_identifier(k._column_name.lower()): k for k in keys}
+
+        # if no keys are defined but the concept has a reference scheme, try to match
+        # reference schemes with column fields by name
+        if not keys and concept._reference_schemes:
+            # do not modify the default argument, create a new array
+            keys = []
+            sanitized_fields = [sanitize_identifier(n.lower()) for n in self._rel._field_names]
+            for scheme in concept._reference_schemes:
+                # if all keys in a scheme map to columns, then use it
+                if all(k._name in sanitized_fields for k in scheme):
+                    for k in scheme:
+                        key = Column(self, k._name)
+                        keys.append(key)
+                        key_dict[k._name] = key
+                    break
+        # TODO: this is correct, but the rel backend does the wrong thing with it
+        # me = concept.new(**key_dict)
+        # items = [me]
+        # for field in self._rel._fields[1:]:
+        #     field_rel = getattr(concept, field.name.lower())
+        #     if sanitize_identifier(field.name.lower()) not in key_dict:
+        #         items.append(field_rel(me, getattr(self, field.name)))
+        # b.where(me).define(*items)
+
+        # this is much less efficient than above
+
+        if keys:
+            me = concept.new(**key_dict)
+            b.define(me)
+        else:
+            me = self._rel._field_refs[0]
+            b.where(self).define(concept(me))
+        # All the fields are treated as properties
+        for field in self._rel._fields[1:]:
+            field_name = sanitize_identifier(field.name.lower())
+            if field_name not in key_dict:
+                r = b.Property(
+                    f"{{{concept}}} has {{{field_name}:{field.type_str}}}",
+                    parent=concept,
+                    short_name=field_name,
+                    model=self._rel._model
+                )
+                setattr(concept, field_name, r)
+                relationship = getattr(concept, field_name)
+                table_name = getattr(self, field.name)
+                update = relationship(me, table_name)
+                b.define(update)
+
+    def _to_type(self) -> b.Concept:
+        return self._concept
+
+    def _to_keys(self) -> list[b.Ref]:
+        self._lazy_init()
+        return [self._rel._field_refs[0]]
+
+    def _compile_lookup(self, compiler:b.Compiler, ctx:b.CompilerContext):
+        self._lazy_init()
+        Table._used_sources.add(self)
+        compiler.lookup(self._rel, ctx)
+        return compiler.lookup(b.RelationshipFieldRef(None, self._rel, 0), ctx)
+
+class Column(b.Producer):
+    def __init__(self, source:Table, column_name:str) -> None:
+        super().__init__(None)
+        self._column_name = column_name
+        self._source = source
+
+    def _to_keys(self) -> list[b.Ref]:
+        return self._source._to_keys()
+
+    def _compile_lookup(self, compiler:b.Compiler, ctx:b.CompilerContext):
+        compiler.lookup(self._source, ctx)
+        for ix, field in enumerate(self._source._rel._field_refs):
+            assert field._name is not None
+            if field._name.lower() == self._column_name.lower():
+                return compiler.lookup(b.RelationshipFieldRef(None, self._source._rel, ix), ctx)
+        raise ValueError(f"Column {self._column_name} not found in {self._source._fqn}")

relationalai/semantics/lqp/README.md

@@ -0,0 +1,34 @@
+# Logic Engine LQP Backend
+
+The logic engine runs the *Logical Query Protocol* (short *LQP*). This module includes a
+compiler from the semantic metamodel to LQP along with an executor.
+
+## Running against a local logic engine
+
+For development and testing, it is possible to run PyRel models against a local logic engine
+server process.
+
+To start your local server, please refer to the [logic engine
+docs](https://github.com/RelationalAI/raicode/tree/master/src/Server#starting-the-server).
+
+With the local server running, add this to your `raiconfig.toml`:
+
+```toml
+[profile.local]
+platform = "local"
+engine = "local"
+host = "localhost"
+port = 8010
+```
+
+Then set `active_profile = "local"` at the top of the file.
+
+**Known limitations:**
+
+Local execution does not support running against Snowflake source tables.
+
+At the moment, locally created databases cannot be cleaned up by the client. Eventually you
+will need to clear your local pager directory.
+
+At the moment, local execution is only supported for fast-path transactions, i.e. those
+which complete in less than 5 seconds. Polling support will be added soon.

relationalai/semantics/lqp/builtins.py

@@ -0,0 +1,16 @@
+from relationalai.semantics.metamodel import factory as f
+from relationalai.semantics.metamodel.util import FrozenOrderedSet
+from relationalai.semantics.metamodel import builtins
+
+# Indicates a relation is short-lived, thus, backends should not optimize for incremental
+# maintenance.
+adhoc = f.relation("adhoc", [])
+adhoc_annotation = f.annotation(adhoc, [])
+
+# We only want to emit attributes for a known set of annotations.
+annotations_to_emit = FrozenOrderedSet([
+    adhoc.name,
+    builtins.function.name,
+    builtins.track.name,
+    builtins.recursion_config.name,
+])

relationalai/semantics/lqp/compiler.py

@@ -0,0 +1,22 @@
+#--------------------------------------------------
+# Compiler
+#--------------------------------------------------
+
+from relationalai.semantics.lqp import ir as lqp, model2lqp
+from relationalai.semantics.lqp.utils import TranslationCtx, UniqueNames
+from relationalai.semantics.lqp.passes import lqp_passes
+from relationalai.semantics.metamodel import ir, compiler as c
+
+from typing import Optional
+
+class Compiler(c.Compiler):
+    def __init__(self):
+        super().__init__(lqp_passes())
+        self.def_names = UniqueNames()
+
+    def do_compile(self, model: ir.Model, options:dict={}) -> tuple[Optional[tuple], lqp.Epoch]:
+        fragment_id: bytes = options.get("fragment_id", bytes(404))
+        # Reset the var context for each compilation
+        # TODO: Change to unique var names per lookup
+        # TODO: Have our needs realigned with NameCache?
+        return model2lqp.to_lqp(model, fragment_id, TranslationCtx(self.def_names))

relationalai/semantics/lqp/constructors.py

@@ -0,0 +1,68 @@
+from typing import Tuple
+from relationalai.semantics.lqp import ir as lqp
+
+def mk_and(args: list[lqp.Formula]) -> lqp.Formula:
+    # Flatten nested conjunctions
+    if any(isinstance(arg, lqp.Conjunction) for arg in args):
+        final_args = []
+        for arg in args:
+            if isinstance(arg, lqp.Conjunction):
+                final_args.extend(arg.args)
+            else:
+                final_args.append(arg)
+        args = final_args
+
+    if len(args) == 1:
+        return args[0]
+
+    return lqp.Conjunction(args=args, meta=None)
+
+def mk_or(args: list[lqp.Formula]) -> lqp.Formula:
+    # Flatten nested disjunctions
+    if any(isinstance(arg, lqp.Disjunction) for arg in args):
+        final_args = []
+        for arg in args:
+            if isinstance(arg, lqp.Disjunction):
+                final_args.extend(arg.args)
+            else:
+                final_args.append(arg)
+        args = final_args
+
+    if len(args) == 1:
+        return args[0]
+
+    return lqp.Disjunction(args=args, meta=None)
+
+def mk_abstraction(vars: list[Tuple[lqp.Var, lqp.Type]], value: lqp.Formula) -> lqp.Abstraction:
+    return lqp.Abstraction(vars=vars, value=value, meta=None)
+
+def mk_exists(vars: list[Tuple[lqp.Var, lqp.Type]], value: lqp.Formula) -> lqp.Formula:
+    if len(vars) == 0:
+        return value
+    abstr = mk_abstraction(vars, value)
+    return lqp.Exists(body=abstr, meta=None)
+
+def mk_specialized_value(value) -> lqp.SpecializedValue:
+    return lqp.SpecializedValue(value=value, meta=None)
+
+def mk_value(value) -> lqp.Value:
+    return lqp.Value(value=value, meta=None)
+
+def mk_type(typename: lqp.TypeName, parameters: list[lqp.Value]=[]) -> lqp.Type:
+    return lqp.Type(type_name=typename, parameters=parameters, meta=None)
+
+def mk_primitive(name: str, terms: list[lqp.RelTerm]) -> lqp.Primitive:
+    return lqp.Primitive(name=name, terms=terms, meta=None)
+
+def mk_pragma(name: str, terms: list[lqp.Var]) -> lqp.Pragma:
+    return lqp.Pragma(name=name, terms=terms, meta=None)
+
+def mk_attribute(name: str, args: list[lqp.Value]) -> lqp.Attribute:
+    return lqp.Attribute(name=name, args=args, meta=None)
+
+def mk_transaction(
+    epochs: list[lqp.Epoch],
+    configure: lqp.Configure = lqp.construct_configure({}, None),
+    sync = None
+) -> lqp.Transaction:
+    return lqp.Transaction(epochs=epochs, configure=configure, sync=sync, meta=None)