ref: unify ops

Author: mesejo

python/xorq/caching/__init__.py

@@ -179,7 +179,7 @@ def rename_remote_table(node, _, **kwargs):
                     name=name,
                     schema=node.schema,
                     source=node.source,
-                    remote_expr=node.remote_expr,
+                    parent=node.parent,
                     namespace=node.namespace,
                 )
                 return rt

python/xorq/common/utils/dask_normalize/dask_normalize_expr.py

@@ -246,7 +246,7 @@ def normalize_letsql_databasetable(dt):
         return dask.base.normalize_token(
             (
                 "normalize_letsql_databasetable",
-                dt.input_expr,
+                dt.parent,
                 # we need to "stabilize" the name of the tables in the unbound expr
                 rename_unbound_static(dt.unbound_expr.op()).to_expr(),
                 dt.make_connection,
@@ -256,7 +256,7 @@ def normalize_letsql_databasetable(dt):
         return dask.base.normalize_token(
             (
                 "normalize_letsql_databasetable",
-                dt.input_expr,
+                dt.parent,
                 dt.udxf.exchange_f,
                 dt.make_connection,
             )
@@ -372,7 +372,7 @@ def normalize_remote_table(dt):
 
     return normalize_seq_with_caller(
         ("schema", dt.schema),
-        ("expr", dt.remote_expr),
+        ("expr", dt.parent),
         # only thing that matters is the type of source its going into
         ("source", dt.source.name),
         caller="normalize_remote_table",
@@ -497,7 +497,7 @@ def opaque_node_replacer(node, kwargs):
         case rel.RemoteTable():
             new_node = api.table(
                 node.schema,
-                name=dask.base.tokenize(node.remote_expr),
+                name=dask.base.tokenize(node.parent),
             ).op()
         case rel.FlightUDXF() | rel.FlightExpr():
             new_node = api.table(

python/xorq/common/utils/graph_utils.py

@@ -34,13 +34,13 @@ def gen_children_of(node: Node) -> Tuple[Node, ...]:
             gen = () if rel_node is None else (to_node(rel_node),)
 
         case rel.RemoteTable():
-            gen = (to_node(node.remote_expr),)
+            gen = (to_node(node.parent),)
 
         case rel.CachedNode():
             gen = (to_node(node.parent),)
 
         case rel.FlightExpr() | rel.FlightUDXF():
-            gen = (to_node(node.input_expr),)
+            gen = (to_node(node.parent),)
 
         case udf.ExprScalarUDF():
             gen = (to_node(node.computed_kwargs_expr),)
@@ -95,13 +95,13 @@ def process_node(op, _kwargs):
         op = replacer(op, _kwargs)
         match op:
             case rel.RemoteTable():
-                remote_expr = op.remote_expr.op().replace(process_node).to_expr()
+                remote_expr = op.parent.op().replace(process_node).to_expr()
                 return do_recreate(op, _kwargs, remote_expr=remote_expr)
             case rel.CachedNode():
                 parent = op.parent.op().replace(process_node).to_expr()
                 return do_recreate(op, _kwargs, parent=parent)
             case rel.FlightExpr() | rel.FlightUDXF():
-                input_expr = op.input_expr.op().replace(process_node).to_expr()
+                input_expr = op.parent.op().replace(process_node).to_expr()
                 return do_recreate(op, _kwargs, input_expr=input_expr)
             case udf.ExprScalarUDF():
                 computed_kwargs_expr = (

python/xorq/common/utils/lineage_utils.py

@@ -172,7 +172,7 @@ def _(node: rel.CachedNode, cfg: Dict[str, Any] | None = None) -> str:
 def _(node: rel.FlightExpr, cfg: Dict[str, Any] | None = None) -> str:
     palette: ColorScheme = (cfg or {}).get("palette", default_palette)
     col = palette.get("flight")
-    return f"{col}FlightExpr ({node.input_expr})[/]"
+    return f"{col}FlightExpr ({node.parent})[/]"
 
 
 @format_node.register

python/xorq/expr/relations.py

@@ -33,7 +33,7 @@ def replace_cache_table(node, kwargs):
     if isinstance(node, CachedNode):
         return node.parent.op().replace(replace_cache_table)
     elif isinstance(node, RemoteTable):
-        return node.remote_expr.op().replace(replace_cache_table)
+        return node.parent.op().replace(replace_cache_table)
     else:
         return node
 
@@ -141,8 +141,11 @@ def __dask_tokenize__(self):
         )
 
 
-class CachedNode(ops.DatabaseTable):
+class DatabaseTableView(ops.DatabaseTable):
     parent: Any = None
+
+
+class CachedNode(DatabaseTableView):
     storage: Any = None
 
 
@@ -154,26 +157,24 @@ class CachedNode(ops.DatabaseTable):
 )
 
 
-class RemoteTable(ops.DatabaseTable):
-    remote_expr: Expr = None
-
+class RemoteTable(DatabaseTableView):
     @classmethod
     def from_expr(cls, con, expr, name=None):
         name = name or gen_name()
         return cls(
             name=name,
             schema=expr.schema(),
             source=con,
-            remote_expr=expr,
+            parent=expr,
         )
 
 
 def into_backend(expr, con, name=None):
     return RemoteTable.from_expr(con=con, expr=expr, name=name).to_expr()
 
 
-class FlightExpr(ops.DatabaseTable):
-    input_expr: Expr = None
+class FlightExpr(DatabaseTableView):
+    parent: Expr = None
     unbound_expr: Expr = None
     make_server: Callable = None
     make_connection: Callable = None
@@ -211,7 +212,7 @@ def roundtrip_cloudpickle(obj):
             name=name or gen_name(),
             schema=unbound_expr.schema(),
             source=input_expr._find_backend(),
-            input_expr=input_expr,
+            parent=input_expr,
             unbound_expr=roundtrip_cloudpickle(unbound_expr),
             make_server=make_server or FlightServer,
             make_connection=make_connection or xo_connect,
@@ -228,7 +229,7 @@ def to_rbr(self, do_instrument_reader=None):
             do_instrument_reader = self.do_instrument_reader
 
         def inner(flight_exchange):
-            rbr_in = flight_exchange.input_expr.to_pyarrow_batches()
+            rbr_in = flight_exchange.parent.to_pyarrow_batches()
             if do_instrument_reader:
                 rbr_in = instrument_reader(rbr_in, "input: ")
             with flight_exchange.make_server() as server:
@@ -316,8 +317,8 @@ def flight_expr(
     )
 
 
-class FlightUDXF(ops.DatabaseTable):
-    input_expr: Expr = None
+class FlightUDXF(DatabaseTableView):
+    parent: Expr = None
     # FIXME: fix circular import issue so we can possibly pass an instance of AbstractExchanger
     udxf: type = None
     make_server: Callable = None
@@ -360,7 +361,7 @@ def make_mtls_server():
             name=name or gen_name(),
             schema=schema,
             source=input_expr._find_backend(),
-            input_expr=input_expr,
+            parent=input_expr,
             udxf=udxf,
             make_server=make_server or make_mtls_server,
             make_connection=make_connection or xo_connect,
@@ -374,7 +375,7 @@ def to_rbr(self, do_instrument_reader=None):
             do_instrument_reader = self.do_instrument_reader
 
         def inner(flight_udxf):
-            rbr_in = flight_udxf.input_expr.to_pyarrow_batches()
+            rbr_in = flight_udxf.parent.to_pyarrow_batches()
             if do_instrument_reader:
                 rbr_in = instrument_reader(rbr_in, "input: ")
             with flight_udxf.make_server() as server:
@@ -603,7 +604,7 @@ def register_and_transform_remote_tables(expr):
         )
     batches_table = {}
     for arg, count in counts.items():
-        ex = arg.remote_expr
+        ex = arg.parent
         batches = ex.to_pyarrow_batches()
         schema = ex.as_table().schema().to_pyarrow()
         replicas = SafeTee.tee(batches, count)

python/xorq/ibis_yaml/sql.py

@@ -58,15 +58,15 @@ def get_name(node):
 
 def find_tables(expr: ir.Expr) -> Tuple[Dict[str, QueryInfo], Dict[str, QueryInfo]]:
     def get_remote_table_backend(node):
-        return node.remote_expr._find_backend()
+        return node.parent._find_backend()
 
     grouped = toolz.groupby(type, walk_nodes((RemoteTable, Read), expr))
     remote_tables: Dict[str, QueryInfo] = {
         node.name: {
             "engine": backend.name,
             "profile_name": backend._profile.hash_name,
-            "relations": find_relations(node.remote_expr),
-            "sql": to_sql(node.remote_expr).strip(),
+            "relations": find_relations(node.parent),
+            "sql": to_sql(node.parent).strip(),
             "options": {},
         }
         for node in grouped.get(RemoteTable, ())

python/xorq/ibis_yaml/translate.py

@@ -510,7 +510,7 @@ def _cached_node_from_yaml(yaml_dict: dict, context: any) -> ibis.Expr:
 @translate_to_yaml.register(RemoteTable)
 def _remotetable_to_yaml(op: RemoteTable, context: any) -> dict:
     profile_name = op.source._profile.hash_name
-    remote_expr_yaml = translate_to_yaml(op.remote_expr, context)
+    remote_expr_yaml = translate_to_yaml(op.parent, context)
     schema_id = context.schema_registry.register_schema(op.schema)
     # TODO: change profile to profile_name
     return freeze(

python/xorq/ibis_yaml/udf.py

@@ -326,7 +326,7 @@ def _aggudf_from_yaml(yaml_dict: dict, compiler: any) -> any:
 
 @translate_to_yaml.register(FlightExpr)
 def flight_expr_to_yaml(op: FlightExpr, context: any) -> dict:
-    input_expr_yaml = translate_to_yaml(op.input_expr, context)
+    input_expr_yaml = translate_to_yaml(op.parent, context)
     unbound_expr_yaml = translate_to_yaml(op.unbound_expr, context)
 
     schema_id = context.schema_registry.register_schema(op.schema)
@@ -379,7 +379,7 @@ def flight_expr_from_yaml(yaml_dict: Dict, context: Any) -> Any:
 
 @translate_to_yaml.register(FlightUDXF)
 def flight_udxf_to_yaml(op: FlightUDXF, context: any) -> dict:
-    input_expr_yaml = translate_to_yaml(op.input_expr, context)
+    input_expr_yaml = translate_to_yaml(op.parent, context)
     schema_id = context.schema_registry.register_schema(op.schema)
     udxf_pickle = serialize_callable(op.udxf)
     make_server_pickle = serialize_callable(op.make_server)

python/xorq/vendor/ibis/expr/types/core.py

@@ -801,7 +801,7 @@ def uncached_one(self):
 
             parent = self.expr.op().parent
             if isinstance(parent.op(), RemoteTable):
-                return parent.op().remote_expr
+                return parent.op().parent
             else:
                 return parent
         else: