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

v0/relationalai/semantics/sql/executor/duck_db.py

@@ -0,0 +1,52 @@
+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 v0.relationalai.semantics.sql import Compiler
+from v0.relationalai.semantics.sql.executor.result_helpers import format_duckdb_columns
+from v0.relationalai.semantics.metamodel import ir, executor as e, factory as f
+
+class DuckDBExecutor(e.Executor):
+
+    def __init__(self, skip_denormalization: bool = False) -> None:
+        super().__init__()
+        self.compiler = Compiler(skip_denormalization)
+
+    def execute(self, model: ir.Model, task: ir.Task, format:Literal["pandas", "snowpark"]="pandas") -> Union[DataFrame, Any]:
+        """ 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:
+            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()
+
+    @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)

v0/relationalai/semantics/sql/executor/result_helpers.py

@@ -0,0 +1,64 @@
+from typing import Sequence, Optional
+
+import pandas as pd
+import numpy as np
+import datetime
+import pyarrow as pa
+from snowflake.connector.cursor import ResultMetadata
+from snowflake.connector.constants import FIELD_ID_TO_NAME
+
+from v0.relationalai.semantics.metamodel import ir
+from v0.relationalai.clients.result_helpers import Int128Dtype
+from v0.relationalai.semantics.metamodel.types import Int64, Int128, DateTime
+
+
+def format_columns(result_frame:pd.DataFrame, result_metadata:Sequence[ResultMetadata],
+                   original_columns_metadata: dict[str, Optional[ir.Type]]) -> pd.DataFrame:
+    for i, col in enumerate(result_frame.columns):
+        col_name = col.lower()
+        col_metadata = result_metadata[i]
+        col_type = original_columns_metadata[col_name] if col_name in original_columns_metadata else None
+        if col_type == DateTime:
+            result_frame[col] = result_frame[col].dt.floor("ms")
+        elif col_type and isinstance(col_type, ir.DecimalType):
+            pa_dtype = pa.decimal128(precision=col_type.precision, scale=col_type.scale)
+            result_frame[col] = result_frame[col].astype(pd.ArrowDtype(pa_dtype))
+        elif result_frame[col].apply(lambda x: isinstance(x, (datetime.date, datetime.datetime))).any():
+            result_frame[col] = pd.to_datetime(result_frame[col], errors='coerce')
+        elif FIELD_ID_TO_NAME[col_metadata.type_code] == "FIXED" and col_metadata.scale == 0:
+            series = result_frame[col]
+
+            # Prefer explicit type from original metadata if present
+            if col_type == Int64:
+                result_frame[col] = series.astype("Int64")
+            elif col_type == Int128:
+                result_frame[col] = series.astype(Int128Dtype())
+            elif col_metadata.precision:
+                result_frame[col] = _cast_integer_column(series, col_metadata.precision)
+        # Handle Snowflake VARIANT columns that are actually numeric
+        elif FIELD_ID_TO_NAME[col_metadata.type_code] == "VARIANT":
+            series = result_frame[col]
+
+            if col_type == Int64:
+                result_frame[col] = series.astype("Int64")
+            elif col_type == Int128:
+                result_frame[col] = series.astype(Int128Dtype())
+
+        # SQL may return None for nulls; replace with np.nan unless the column is FIXED, because
+        # both Int64 and Int128 already deal with nulls using pd.NA
+        if not FIELD_ID_TO_NAME[col_metadata.type_code] == "FIXED":
+            result_frame[col] = result_frame[col].replace({None: np.nan})
+    return result_frame
+
+def format_duckdb_columns(result_frame: pd.DataFrame, arrow_schema: pa.Schema) -> pd.DataFrame:
+    for field in arrow_schema:
+        col = field.name
+        pa_type = field.type
+        if pa.types.is_decimal(pa_type):
+            if pa_type.scale == 0:  # integer-like
+                result_frame[col] = _cast_integer_column(result_frame[col], pa_type.precision)
+    return result_frame
+
+def _cast_integer_column(series: pd.Series, precision: int) -> pd.Series:
+    """Cast a numeric pandas Series to Int64 or Int128 depending on precision."""
+    return series.astype("Int64") if precision <= 19 else series.astype(Int128Dtype())

v0/relationalai/semantics/sql/executor/snowflake.py

@@ -0,0 +1,145 @@
+from __future__ import annotations
+
+import re
+import uuid
+
+import v0.relationalai as rai
+import pandas as pd
+
+from typing import Any, Union, Optional, Literal, TYPE_CHECKING
+from snowflake.snowpark import Session
+
+from .. import Compiler
+from v0.relationalai import debugging
+from v0.relationalai.semantics.metamodel import ir, executor as e, factory as f
+from v0.relationalai.clients.result_helpers import sort_data_frame_result
+from v0.relationalai.semantics.sql.executor.result_helpers import format_columns
+from v0.relationalai.semantics.metamodel.visitor import collect_by_type
+from v0.relationalai.semantics.metamodel.typer import typer
+from v0.relationalai.tools.constants import USE_DIRECT_ACCESS
+
+if TYPE_CHECKING:
+    from v0.relationalai.semantics.snowflake import Table
+
+
+class SnowflakeExecutor(e.Executor):
+    """Executes SQL using the RAI client."""
+
+    def __init__(
+            self,
+            database: str,
+            schema: str,
+            dry_run: bool = False,
+            skip_denormalization: bool = False,
+            connection: Session | None = None,
+            config: rai.Config | None = None,
+    ) -> None:
+        super().__init__()
+        self.database = database
+        self.schema = schema
+        self.dry_run = dry_run
+        self._last_model = None
+        self._last_model_sql = None
+        self.config = config or rai.Config()
+        self.compiler = Compiler(skip_denormalization)
+        self.connection = connection
+        self._resources = None
+        self._provider = None
+
+    @property
+    def resources(self):
+        if not self._resources:
+            with debugging.span("create_session"):
+                self.dry_run |= bool(self.config.get("compiler.dry_run", False))
+                resource_class = rai.clients.snowflake.Resources
+                if self.config.get("use_direct_access", USE_DIRECT_ACCESS):
+                    resource_class = rai.clients.snowflake.DirectAccessResources
+                self._resources = resource_class(dry_run=self.dry_run, config=self.config, generation=rai.Generation.QB,
+                                                 connection=self.connection)
+        return self._resources
+
+    @property
+    def provider(self):
+        if not self._provider:
+            self._provider = rai.clients.snowflake.Provider(resources=self.resources)
+        return self._provider
+
+    def execute(self, model: ir.Model, task: ir.Task, format:Literal["pandas", "snowpark"]="pandas",
+                export_to: Optional[Table] = None,
+                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)
+        default_dynamic_table_target_lag = (
+            self.resources.config.get("reasoner.rule.sql.default_dynamic_table_target_lag", None))
+
+        options = {"warehouse": warehouse, "default_dynamic_table_target_lag": default_dynamic_table_target_lag}
+
+        if self._last_model != model:
+            with debugging.span("compile", metamodel=model) as model_span:
+                model_span["compile_type"] = "model"
+                model_sql, _ = self.compiler.compile(model, options)
+                model_span["sql"] = model_sql
+                self._last_model = model
+                self._last_model_sql = model_sql
+
+        with debugging.span("compile", metamodel=task) as compile_span:
+            compile_span["compile_type"] = "query"
+            # compile into sql and keep the new_task, which is the task model after rewrites,
+            # as it may contain results of type inference, which is useful for determining
+            # how to format the outputs
+            query_options = {**options, "query_compilation": True}
+            query_sql, new_task = self.compiler.compile(f.compute_model(f.logical([task])), query_options)
+            compile_span["sql"] = query_sql
+
+        if self.dry_run:
+            return pd.DataFrame()
+
+        if format != "pandas":
+            raise ValueError(f"Unsupported format: {format}")
+
+        _replace_pattern = re.compile(r"[ /:-]")
+
+        def sanitize_name(value: str) -> str:
+            return _replace_pattern.sub("_", value)
+
+        database = sanitize_name(self.database)
+        schema = sanitize_name(self.schema)
+        unique_id = sanitize_name(str(uuid.uuid4()).lower())
+        db_name = f"{database}_{unique_id}"
+        db_query = f"CREATE OR REPLACE DATABASE {db_name};"
+        schema_query = f"CREATE OR REPLACE SCHEMA {db_name}.{schema};"
+        use_schema_query = f"USE SCHEMA {db_name}.{schema};"
+
+        full_model_sql = f"{db_query}\n{schema_query}\n{use_schema_query}\n{self._last_model_sql}\n{query_sql}"
+
+        try:
+            result = self.provider.resources._session.connection.execute_string(full_model_sql) # type: ignore
+
+            # Assuming that `task` is a single SQL query per model, and we have it always at the end of the generated SQL.
+            rows = result[-1].fetchall()
+            result_metadata = result[-1].description
+
+            df = pd.DataFrame(rows, columns=[col.name for col in result_metadata])
+            if df.empty:
+                # return empty df without column names if it's empty
+                return df.iloc[:, 0:0]
+
+            df = format_columns(df, result_metadata, self._collect_columns_metadata(new_task))
+            return sort_data_frame_result(df)
+
+        finally:
+            self.provider.sql(f"DROP DATABASE IF EXISTS {db_name};")
+
+    def _collect_columns_metadata(self, task: ir.Task) -> dict[str, Optional[ir.Type]]:
+        if not task:
+            return {}
+
+        outputs = collect_by_type(ir.Output, task)
+
+        assert len(outputs) == 1
+        return {
+            alias.lower(): typer.to_base_primitive(var.type)
+            for alias, var in outputs[0].aliases
+            if alias and isinstance(var, ir.Var)
+        }

v0/relationalai/semantics/sql/rewrite/denormalize.py

@@ -0,0 +1,222 @@
+from __future__ import annotations
+from dataclasses import dataclass, field
+from typing import Tuple, Union
+from typing import cast
+
+from v0.relationalai.semantics.metamodel import ir, compiler as c, factory as f, types, visitor
+from v0.relationalai.semantics.metamodel.util import FrozenOrderedSet, OrderedSet, ordered_set, group_by, split_by
+
+class Denormalize(c.Pass):
+    """
+    Pass to denormalize relations, grouping relations that have the same entity as key.
+
+    General algorithm:
+
+    1. when handling the Model, go to its relations and find the ones that have common
+    entity keys. They will be denormalized (joined together as a single "table").
+    2. when handling lookups and updates that refer to one of these relations that were
+    denormalized, change the reference to be to the new relation and fill the blanks with
+    Nones, which represent NULLs.
+    3. finally, when handling Logical nodes, try to group together lookups and updates that
+    are refering to the same denormalized relation. This may remove some of the NULLs.
+    """
+    def rewrite(self, model: ir.Model, options:dict={}) -> ir.Model:
+        return OldDenormalize().walk(model)
+
+@dataclass
+class OldDenormalize(visitor.Rewriter):
+
+    denormalized: dict[ir.Relation, ir.Relation] = field(default_factory=dict, init=False, hash=False, compare=False)
+
+    def reset(self):
+        self.context = dict()
+        self.denormalized = dict()
+
+    def handle_model(self, model: ir.Model, parent: None):
+        # denormalize some relations in the model, set it in the context, and then push
+        # that context to the root traversal so that references to the denormalized
+        # relations are adjusted.
+
+        # denormalize the relations
+        relations, self.denormalized = self._denormalize_relations(model.relations)
+
+        # rewrite tasks to refer to denormalized relations
+        return ir.Model(
+            self.walk_set(model.engines, model),
+            relations,
+            self.walk_set(model.types, model),
+            self.walk(model.root, model),
+        )
+
+    def handle_update(self, node: ir.Update, parent: ir.Node):
+        # if the relation to be updated was denormalized, denormalize the reference
+
+        if node.relation not in self.denormalized:
+            return super().handle_update(node, parent)
+
+        denormalized_relation = self.denormalized[node.relation]
+        args = self._denormalize_reference(
+            denormalized_relation, node.relation, node.args
+        )
+
+        return ir.Update(
+            node.engine,
+            denormalized_relation,
+            args,
+            node.effect,
+            node.annotations
+        )
+
+    def handle_lookup(self, node: ir.Lookup, parent: ir.Node):
+        # if the relation to be looked up was denormalized, denormalize the reference
+
+        if node.relation not in self.denormalized:
+            return super().handle_lookup(node, parent)
+
+        denormalized_relation = self.denormalized[node.relation]
+        args = self._denormalize_reference(
+            denormalized_relation, node.relation, node.args
+        )
+
+        return ir.Lookup(
+            node.engine,
+            denormalized_relation,
+            args
+        )
+
+    def handle_logical(self, node: ir.Logical, parent: ir.Node):
+        # denormalize the body first
+        body = self.walk_list(node.body, node)
+
+        # Unwrap single lookups from the logical
+        # for example we may have this IR to say that `adult` may not have a `name`
+        # Logical ⇑[name=None]
+        #     name(adult, name)
+        updated_body = [
+            t.body[0] if isinstance(t, ir.Logical) and len(t.body) == 1 else t
+            for t in body
+        ]
+
+        # function to compute a key to group Lookups or Updates that can be merged
+        def task_key(t: ir.Task):
+            assert isinstance(t, (ir.Lookup, ir.Update))
+            return (t.kind, t.relation, t.args[0])
+
+        # now group references
+        tasks, new_body = split_by(updated_body, lambda t: isinstance(t, (ir.Lookup, ir.Update)))
+        for _, grouped_tasks in group_by(tasks, task_key).items():
+            grouped_tasks = cast(OrderedSet[Union[ir.Lookup, ir.Update]], grouped_tasks)
+
+            # get some of the tasks as a representative
+            some = grouped_tasks.some()
+
+            if len(grouped_tasks) == 1:
+                new_body.append(some)
+                continue
+
+            # now join the arguments of the multiple references
+            # TODO: this will fail if we have multiple tasks on the same arg
+            new_args = []
+            i = 0
+            for _ in some.relation.fields:
+                # note that we are still using None to indicate we don't care about the column
+                v = None
+                for t in grouped_tasks:
+                    if t.args[i] is not None:
+                        v = t.args[i]
+                        break
+                new_args.append(v)
+                i += 1
+
+            # finally, create a rewritten task for the whole group
+            if isinstance(some, ir.Lookup):
+                new_body.append(ir.Lookup(some.engine, some.relation, tuple(new_args)))
+            else:
+                new_body.append(ir.Update(some.engine, some.relation, tuple(new_args), some.effect, some.annotations)) # TODO maybe merge annos?
+        return ir.Logical(
+            node.engine,
+            node.hoisted,
+            tuple(new_body)
+        )
+
+
+    def _denormalize_reference(self, denormalized_relation, relation, args) -> Tuple[ir.Value, ...]:
+        """ Adjust arguments to a denormalized reference.
+
+            A reference to `relation` was used with `args`. But `relation` was denormalized
+            into `denormalized_relation`. This method returns a new tuple of adjusted args
+            to account for the denormalization, filling Nones as appropriate.
+        """
+
+        # the first arg is always the entity
+        new_args = [args[0]]
+        # for the rest of the denormalized relation columns, lookup only if the field came
+        # from this relation, otherwise use None
+        i = 1
+        for fld in denormalized_relation.fields[1:]:
+            if fld in relation.fields:
+                new_args.append(args[i])
+                i += 1
+            else:
+                new_args.append(None)
+        return tuple(new_args)
+
+    def _denormalize_relations(self, relations: FrozenOrderedSet[ir.Relation]) -> \
+        Tuple[FrozenOrderedSet[ir.Relation], dict[ir.Relation, ir.Relation]]:
+
+        """ 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 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.
+        """
+        new_relations = ordered_set()
+        denormalized: dict[ir.Relation, ir.Relation] = dict()
+
+        # entities are non-builtin types that have a unary relation
+        entity_types: OrderedSet[ir.Type] = ordered_set()
+        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) and not r.name == 'Error':
+                e = r.fields[0].type
+                entity_types.add(e)
+                entity_relations[e] = r
+
+        # filter out relations that don't have an entity as a key
+        with_entity, unaffected = split_by(
+            relations,
+            lambda r:
+                r.fields[0].type in entity_types # the first element is an entity
+                # TODO - we need to account for wide relations that have multis
+        )
+        new_relations.update(unaffected)
+
+        # group relations by entity, they will be denormalized
+        by_entity = group_by(with_entity, lambda r: cast(ir.ScalarType, r.fields[0].type))
+        for e_type, group in by_entity.items():
+            r = entity_relations[e_type]
+            fields = []
+            # the first column of the denormalized relation is the "id"
+            fields.append(r.fields[0])
+            for relation in group:
+                # add all fields except the first entity field
+                # TODO - we may have multiple relations with the same name
+                fields.extend(relation.fields[1:])
+
+            # create the new denormalized relation, add it to the new_relations set
+            denormalized_relation = f.relation(
+                e_type.name,
+                fields
+            )
+            new_relations.add(denormalized_relation)
+
+            # make sure the original relations point to the new relation
+            for g in group:
+                denormalized[g] = denormalized_relation
+
+        return new_relations.frozen(), denormalized

v0/relationalai/semantics/sql/rewrite/double_negation.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from v0.relationalai.semantics.metamodel import ir, compiler as c, factory as f, visitor
+
+
+class DoubleNegation(c.Pass):
+    """
+        Pass to drop double negations.
+
+        Examples:
+        not_(not_(Person.age > 18, Person.id == 2)) -> (Person.age > 18, Person.id == 2)
+        not_(not_(not_(Person.age > 18, Person.id == 2))) -> not_(Person.age > 18, Person.id == 2)
+    """
+
+    def rewrite(self, model: ir.Model, options: dict = {}) -> ir.Model:
+        return DropDoubleNegationsVisitor().walk(model)
+
+
+@dataclass
+class DropDoubleNegationsVisitor(visitor.Rewriter):
+    """
+    Visitor drops double negations when they are followed one by one in IR Model.
+    The rewrite occurs only when `ir.Not`'s task is an `ir.Logical` with single `ir.Not` in the body.
+
+    Example:
+        Not
+            Logical
+                Not
+                    Logical
+                        age(person, age)
+                        age > 18
+
+        Logical
+            age(person, age)
+            age > 18
+    """
+
+    def handle_not(self, node: ir.Not, parent: ir.Node):
+        if isinstance(node.task, ir.Logical) and len(node.task.body) == 1 and isinstance(node.task.body[0], ir.Not):
+            new_task = self.walk(node.task.body[0].task, node.task.body[0])
+            # dropping redundant Logical if it has only 1 element in his body
+            if isinstance(new_task, ir.Logical) and len(new_task.body) == 1:
+                return new_task.body[0]
+            else:
+                return new_task
+        else:
+            new_task = self.walk(node.task, node)
+            return node if node.task is new_task else f.not_(new_task)

v0/relationalai/semantics/sql/rewrite/recursive_union.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+
+from collections import defaultdict
+from dataclasses import dataclass, field
+from typing import cast, Optional
+
+from v0.relationalai.semantics.metamodel import ir, compiler as c, factory as f
+from v0.relationalai.semantics.metamodel.util import OrderedSet, ordered_set
+from v0.relationalai.semantics.metamodel.visitor import Visitor
+
+
+class RecursiveUnion(c.Pass):
+    """
+        Pass to rewrite recursive relations as unions of `ir.Logical` nodes.
+
+        This pass identifies recursive relations by detecting `ir.Logical` nodes that both read from and write to the same `ir.Relation`.
+        Once a recursive relation is found, all logicals that write to that relation are grouped together.
+        These groups are then wrapped in an `ir.Union`, and the model's root task is reconstructed accordingly.
+    """
+
+    def rewrite(self, model: ir.Model, options:dict={}) -> ir.Model:
+        rpv = TopLogicalReadWriteVisitor()
+        model.accept(rpv)
+
+        reads: dict[ir.Logical, OrderedSet[int]] = rpv.reads()
+        writes: dict[ir.Logical, OrderedSet[int]] = rpv.writes()
+
+        # Step 1: Identify recursive relation IDs (read & write in same logical)
+        recursive_rel_ids = {
+            rel_id
+            for logical, write_ids in writes.items()
+            for rel_id in write_ids & reads.get(logical, OrderedSet())
+        }
+
+        # Step 2: Group ir.Logical nodes that write to any recursive relation
+        recursive_groups: dict[int, OrderedSet[ir.Logical]] = defaultdict(OrderedSet)
+        recursive_logicals: OrderedSet[ir.Logical] = OrderedSet()
+
+        for logical, write_ids in writes.items():
+            for rel_id in write_ids:
+                if rel_id in recursive_rel_ids:
+                    recursive_groups[rel_id].add(logical)
+                    recursive_logicals.add(logical)
+
+        if recursive_logicals:
+
+            # Step 3: Construct the new logical body
+            root_logical = cast(ir.Logical, model.root)
+            new_body = [logical for logical in root_logical.body if logical not in recursive_logicals]
+
+            # Step 4: Add unions for each recursive 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)
+
+        return model
+
+@dataclass
+class TopLogicalReadWriteVisitor(Visitor):
+    """
+    Compute the set of reads and writes relation ids for top level Logical nodes.
+    Skip unions, because we are using this pass to find recursive rule and union them.
+
+    Note that reads are Lookups and writes are Updates. We don't consider Output a write
+    because it is not targeting a relation.
+
+    This visitor can be called from ir.Model.
+    """
+
+    _reads: dict[ir.Logical, OrderedSet[int]] = field(default_factory=dict)
+    _writes: dict[ir.Logical, OrderedSet[int]] = field(default_factory=dict)
+
+    def reads(self):
+        return self._reads
+
+    def writes(self):
+        return self._writes
+
+    _stack: list[ir.Logical] = field(default_factory=list)
+
+    def visit_logical(self, node: ir.Logical, parent: Optional[ir.Node]):
+        # We only track reads/writes for top-level Logical nodes and also skip top level program Logical node.
+        if not isinstance(parent, ir.Model):
+            self._stack.append(node)
+        super().visit_logical(node, parent)
+        if self._stack:
+            self._stack.pop()
+
+    def visit_lookup(self, node: ir.Lookup, parent: Optional[ir.Node]):
+        for logical in self._stack:
+            if logical not in self._reads:
+                self._reads[logical] = ordered_set()
+            self._reads[logical].add(node.relation.id)
+        return super().visit_lookup(node, parent)
+
+    def visit_update(self, node: ir.Update, parent: Optional[ir.Node]):
+        for logical in self._stack:
+            if logical not in self.writes():
+                self._writes[logical] = ordered_set()
+            self._writes[logical].add(node.relation.id)
+        return super().visit_update(node, parent)
+
+    def visit_union(self, node: ir.Union, parent: Optional[ir.Node]):
+        # Skip unions, because we are using this pass to find recursive rule and union them.
+        pass