relationalai 1.0.0a1__py3-none-any.whl → 1.0.0a2__py3-none-any.whl

v0/relationalai/semantics/lqp/model2lqp.py

@@ -396,6 +396,41 @@ def _translate_ascending_rank(ctx: TranslationCtx, limit: int, result_var: lqp.V
             terms=terms,
         )
 
+def _rename_shadowed_abstraction_vars(
+    ctx: TranslationCtx,
+    aggr: ir.Aggregate,
+    abstr_args: list[Tuple[lqp.Var, lqp.Type]],
+    body_conjs: list[lqp.Formula]
+) -> list[Tuple[lqp.Var, lqp.Type]]:
+    """
+    Rename abstraction variables that shadow group-by variables.
+
+    This can happen when the same variable appears in both aggr.group and as an input
+    to the aggregation, e.g., min(Person.age).per(Person.age). The group-by variables
+    are in the outer scope, while the abstraction parameters are in the inner scope,
+    so we need different names to avoid shadowing.
+    """
+    # Get the LQP names of group-by variables
+    group_var_names = set()
+    for group_var in aggr.group:
+        lqp_var = _translate_var(ctx, group_var)
+        group_var_names.add(lqp_var.name)
+
+    # Rename any abstraction parameters that conflict with group-by variables
+    renamed_abstr_args = []
+    for var, typ in abstr_args:
+        if var.name in group_var_names:
+            # This variable shadows a group-by variable, so rename it
+            fresh_var = gen_unique_var(ctx, var.name)
+            # Add an equality constraint: fresh_var == var
+            # var is a free variable referring to the outer scope group-by variable
+            body_conjs.append(mk_primitive("rel_primitive_eq", [fresh_var, var]))
+            renamed_abstr_args.append((fresh_var, typ))
+        else:
+            renamed_abstr_args.append((var, typ))
+
+    return renamed_abstr_args
+
 def _translate_aggregate(ctx: TranslationCtx, aggr: ir.Aggregate, body: lqp.Formula) -> Union[lqp.Reduce, lqp.Formula]:
     # TODO: handle this properly
     aggr_name = aggr.aggregation.name
@@ -432,6 +467,9 @@ def _translate_aggregate(ctx: TranslationCtx, aggr: ir.Aggregate, body: lqp.Form
     body_conjs.extend(projected_eqs)
     abstr_args: list[Tuple[lqp.Var, lqp.Type]] = projected_args + input_args
 
+    # Rename abstraction variables that shadow group-by variables
+    abstr_args = _rename_shadowed_abstraction_vars(ctx, aggr, abstr_args, body_conjs)
+
     if aggr_name == "count":
         assert len(output_terms) == 1, "Count and avg expect a single output variable"
         assert isinstance(meta_output_terms[0], ir.Var)
@@ -441,7 +479,6 @@ def _translate_aggregate(ctx: TranslationCtx, aggr: ir.Aggregate, body: lqp.Form
         one_var, eq = constant_to_var(ctx, to_lqp_value(1, meta_output_terms[0].type), "one")
         body_conjs.append(eq)
         abstr_args.append((one_var, typ))
-        body = mk_and(body_conjs)
 
     # Average needs to wrap the reduce in Exists(Conjunction(Reduce, div))
     if aggr_name == "avg":
@@ -454,7 +491,6 @@ def _translate_aggregate(ctx: TranslationCtx, aggr: ir.Aggregate, body: lqp.Form
         one_var, eq = constant_to_var(ctx, to_lqp_value(1, types.Int64), "one")
         body_conjs.append(eq)
         abstr_args.append((one_var, count_type))
-        body = mk_and(body_conjs)
 
         # The average will produce two output variables: sum and count.
         sum_result = gen_unique_var(ctx, "sum")
@@ -462,6 +498,7 @@ def _translate_aggregate(ctx: TranslationCtx, aggr: ir.Aggregate, body: lqp.Form
 
         # Second to last is the variable we're summing over.
         (sum_var, sum_type) = abstr_args[-2]
+        body = mk_and(body_conjs)
 
         result = lqp.Reduce(
             op=lqp_avg_op(ctx, aggr.aggregation.name, sum_var.name, sum_type),
@@ -494,6 +531,7 @@ def _translate_aggregate(ctx: TranslationCtx, aggr: ir.Aggregate, body: lqp.Form
     # `input_args` hold the types of the input arguments, but they may have been modified
     # if we're dealing with a count, so we use `abstr_args` to find the type.
     (aggr_arg, aggr_arg_type) = abstr_args[-1]
+    body = mk_and(body_conjs)
 
     # Group-bys do not need to be handled at all, since they are introduced outside already
     reduce = lqp.Reduce(
@@ -668,11 +706,11 @@ def to_lqp_value(value: ir.PyValue, value_type: ir.Type) -> lqp.Value:
         val = value
     elif typ.type_name == lqp.TypeName.STRING and isinstance(value, str):
         val = value
-    elif typ.type_name == lqp.TypeName.DECIMAL and isinstance(value, PyDecimal):
+    elif typ.type_name == lqp.TypeName.DECIMAL and isinstance(value, (int, float, PyDecimal)):
         precision = typ.parameters[0].value
         scale = typ.parameters[1].value
         assert isinstance(precision, int) and isinstance(scale, int)
-        val = lqp.DecimalValue(precision=precision, scale=scale, value=value, meta=None)
+        val = lqp.DecimalValue(precision=precision, scale=scale, value=PyDecimal(value), meta=None)
     elif typ.type_name == lqp.TypeName.DATE and isinstance(value, date):
         val = lqp.DateValue(value=value, meta=None)
     elif typ.type_name == lqp.TypeName.DATETIME and isinstance(value, datetime):

v0/relationalai/semantics/lqp/passes.py

@@ -390,7 +390,7 @@ class EliminateData(Pass):
                     [
                         f.logical(
                             [
-                                f.lookup(rel_builtins.eq, [f.literal(val), var])
+                                f.lookup(rel_builtins.eq, [f.literal(val, var.type), var])
                                 for (val, var) in zip(row, node.vars)
                             ],
                         )

v0/relationalai/semantics/lqp/rewrite/cdc.py

@@ -200,7 +200,7 @@ class CDC(Pass):
         Get the relation that represents this property var in this wide_cdc_relation. If the
         relation is not yet available in the context, this method will create and register it.
         """
-        relation_name = wide_cdc_relation.name.lower().replace(".", "_")
+        relation_name = helpers.sanitize(wide_cdc_relation.name).replace("-", "_")
         key = (relation_name, property.name)
         if key not in ctx.cdc_relations:
             # the property relation is overloaded for all properties of the same wide cdc relation, so they have

v0/relationalai/semantics/lqp/rewrite/extract_keys.py

@@ -335,9 +335,45 @@ class ExtractKeysRewriter(Rewriter):
         partitions:dict[ir.Var, OrderedSet[ir.Task]] = defaultdict(OrderedSet)
         dependencies:dict[ir.Var, OrderedSet[ir.Var]] = defaultdict(OrderedSet)
 
-        worklist = list(vars)
-        while worklist:
-            var = worklist.pop()
+        def dfs_collect_deps(task, deps):
+            if isinstance(task, ir.Lookup):
+                args = helpers.vars(task.args)
+                for i, v in enumerate(args):
+                    # v depends on all previous vars
+                    for j in range(i):
+                        deps[v].add(args[j])
+                    # for ternary+ lookups, a var also depends on the next vars
+                    if i > 0 and len(args) >= 3:
+                        for j in range(i+1, len(args)):
+                            deps[v].add(args[j])
+            elif isinstance(task, ir.Construct):
+                vars = helpers.vars(task.values)
+                for val_var in vars:
+                    deps[task.id_var].add(val_var)
+            elif isinstance(task, ir.Logical):
+                for child in task.body:
+                    dfs_collect_deps(child, deps)
+            elif isinstance(task, (ir.Match, ir.Union)):
+                for child in task.tasks:
+                    dfs_collect_deps(child, deps)
+
+        for task in tasks:
+            dfs_collect_deps(task, dependencies)
+
+        def dfs_transitive_deps(var, visited):
+            for dep_var in dependencies[var]:
+                if dep_var not in visited:
+                    visited.add(dep_var)
+                    dfs_transitive_deps(dep_var, visited)
+
+        transitive_deps = defaultdict(OrderedSet)
+        for var in list(dependencies.keys()):
+            visited = OrderedSet()
+            dfs_transitive_deps(var, visited)
+            transitive_deps[var] = visited
+        dependencies = transitive_deps
+
+        for var in vars:
             extended_vars = OrderedSet[ir.Var]()
             extended_vars.add(var)
 
@@ -347,28 +383,33 @@ class ExtractKeysRewriter(Rewriter):
                 for task in tasks:
                     if task in partitions[var]:
                         continue
-                    # Already added this task to this partition
+
+                    if isinstance(task, (ir.Logical, ir.Match, ir.Union)):
+                        hoisted = helpers.hoisted_vars(task.hoisted)
+                        if var in hoisted:
+                            partitions[var].add(task)
+                            there_is_progress = True
+                    elif isinstance(task, ir.Construct):
+                        if task.id_var == var:
+                            partitions[var].add(task)
+                            there_is_progress = True
                     elif isinstance(task, ir.Lookup):
                         args = helpers.vars(task.args)
                         if len(args) == 1 and args[0] in extended_vars:
                             partitions[var].add(task)
-                        # TODO: hack to have dot_joins work
+                            there_is_progress = True
+                        # NOTE: heuristics to have dot_joins work
                         elif len(args) >= 3 and args[-2] in extended_vars:
                             partitions[var].add(task)
                             extended_vars.add(args[-1])
-                            dependencies[var].add(args[-1])
                             there_is_progress = True
                         elif len(args) > 1 and args[-1] in extended_vars:
                             partitions[var].add(task)
                             for arg in args[:-1]:
                                 extended_vars.add(arg)
-                                dependencies[var].add(arg)
-                            there_is_progress = True
-                    elif isinstance(task, ir.Logical):
-                        hoisted = helpers.hoisted_vars(task.hoisted)
-                        if var in hoisted:
-                            partitions[var].add(task)
                             there_is_progress = True
+                    else:
+                        assert False, f"invalid node kind {type(task)}"
 
         return partitions, dependencies
 

v0/relationalai/semantics/metamodel/builtins.py

@@ -443,12 +443,14 @@ datetime_second = f.relation("datetime_second", [f.input_field("a", types.DateTi
 datetime_weekday = f.relation("datetime_weekday", [f.input_field("a", types.DateTime), f.input_field("b", types.String), f.field("c", types.Int64)])
 
 # Other
-range = f.relation("range", [
-    f.input_field("start", types.Int64),
-    f.input_field("stop", types.Int64),
-    f.input_field("step", types.Int64),
-    f.field("result", types.Int64),
-])
+range = f.relation(
+    "range",
+    [f.input_field("start", types.Number), f.input_field("stop", types.Number), f.input_field("step", types.Number), f.field("result", types.Number)],
+    overloads=[
+        f.relation("range", [f.input_field("start", types.Int64), f.input_field("stop", types.Int64), f.input_field("step", types.Int64), f.field("result", types.Int64)]),
+        f.relation("range", [f.input_field("start", types.Int128), f.input_field("stop", types.Int128), f.input_field("step", types.Int128), f.field("result", types.Int128)]),
+    ],
+)
 
 hash = f.relation("hash", [f.input_field("args", types.AnyList), f.field("hash", types.Hash)])
 

v0/relationalai/semantics/metamodel/rewrite/flatten.py

@@ -124,10 +124,10 @@ class Flatten(Pass):
                 output
     """
 
-    def __init__(self, handle_outputs: bool=True):
+    def __init__(self, use_sql: bool=False):
         super().__init__()
         self.name_cache = NameCache(start_from_one=True)
-        self._handle_outputs = handle_outputs
+        self._use_sql = use_sql
 
 
     #--------------------------------------------------
@@ -181,7 +181,12 @@ class Flatten(Pass):
     def handle(self, task: ir.Task, ctx: Context) -> Flatten.HandleResult:
         if isinstance(task, ir.Logical):
             return self.handle_logical(task, ctx)
-        elif isinstance(task, ir.Union):
+        elif isinstance(task, ir.Union) and (task.hoisted or self._use_sql):
+            # Only flatten Unions which hoist variables. If there are no hoisted variables,
+            # then the Union acts as a filter, and it can be inefficient to flatten it.
+            #
+            # However, for the SQL backend, we always need to flatten Unions for correct SQL
+            # generation.
             return self.handle_union(task, ctx)
         elif isinstance(task, ir.Match):
             return self.handle_match(task, ctx)
@@ -238,7 +243,7 @@ class Flatten(Pass):
         # If there are outputs, flatten each into its own top-level rule, along with its
         # dependencies.
         if groups["outputs"]:
-            if not self._handle_outputs:
+            if self._use_sql:
                 ctx.rewrite_ctx.top_level.append(ir.Logical(task.engine, task.hoisted, tuple(body), task.annotations))
                 return Flatten.HandleResult(None)
 

v0/relationalai/semantics/reasoners/graph/core.py

@@ -6222,9 +6222,9 @@ class Graph():
     def _distance_reversed_non_weighted(self):
         """Lazily define and cache the self._distance_reversed_non_weighted relationship, a non-public helper."""
         _distance_reversed_non_weighted_rel = self._model.Relationship(f"{{node_u:{self._NodeConceptStr}}} and {{node_v:{self._NodeConceptStr}}} have a reversed distance of {{d:Integer}}")
-        node_u, node_v, node_n, d1 = self.Node.ref(), self.Node.ref(), self.Node.ref(), Integer.ref()
+        node_u, node_v, node_n = self.Node.ref(), self.Node.ref(), self.Node.ref()
         node_u, node_v, d = union(
-            where(node_u == node_v, d1 == 0).select(node_u, node_v, d1), # Base case.
+            where(node_u == node_v, d1 := 0).select(node_u, node_v, d1), # Base case.
             where(self._edge(node_v, node_n),
                   d2 := _distance_reversed_non_weighted_rel(node_u, node_n, Integer) + 1).select(node_u, node_v, d2) # Recursive case.
         )
@@ -6326,13 +6326,12 @@ class Graph():
         _is_connected_rel.annotate(annotations.track("graphs", "is_connected"))
 
         where(
-            self._num_nodes(0) |
-            count(self._reachable_from_min_node(self.Node.ref())) == self._num_nodes(Integer.ref())
-        ).define(_is_connected_rel(True))
-
-        where(
-            not_(_is_connected_rel(True))
-        ).define(_is_connected_rel(False))
+            union(
+                self._num_nodes(0),
+                count(self._reachable_from_min_node(self.Node.ref())) == self._num_nodes(Integer.ref())
+            )
+        ).define(_is_connected_rel(True)) \
+        | define(_is_connected_rel(False))
 
         return _is_connected_rel
 

v0/relationalai/semantics/sql/compiler.py

@@ -33,7 +33,7 @@ class Compiler(c.Compiler):
             ExtractNestedLogicals(),  # before InferTypes to avoid extracting casts
             InferTypes(),
             DNFUnionSplitter(),
-            Flatten(handle_outputs=False),
+            Flatten(use_sql=True),
             rewrite.RecursiveUnion(),
             rewrite.DoubleNegation(),
             rewrite.SortOutputQuery()
@@ -1264,7 +1264,7 @@ class ModelToSQL:
                 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:
+            elif relation == builtins.range or relation in builtins.range.overloads:
                 assert len(args) == 4, f"Expected 4 args for `range`, got {len(args)}: {args}"
                 start_raw, stop_raw, step_raw, result = args
                 start = self._var_to_expr(start_raw, reference, resolve_builtin_var, var_to_construct)