feat(build system): consolidate YAML/SQL serialization

examples/yaml_roundrip.py

@@ -0,0 +1,27 @@
+import xorq as xo
+from xorq.common.utils.defer_utils import deferred_read_parquet
+from xorq.expr.relations import into_backend
+from xorq.ibis_yaml.compiler import BuildManager
+
+
+pg = xo.postgres.connect_examples()
+db = xo.duckdb.connect()
+
+batting = pg.table("batting")
+
+backend = xo.duckdb.connect()
+awards_players = deferred_read_parquet(
+    backend,
+    xo.config.options.pins.get_path("awards_players"),
+    table_name="award_players",
+)
+left = batting.filter(batting.yearID == 2015)
+right = awards_players.filter(awards_players.lgID == "NL").drop("yearID", "lgID")
+expr = left.join(
+    into_backend(right, pg, "pg-filtered-table"), ["playerID"], how="semi"
+)[["yearID", "stint"]]
+
+build_manager = BuildManager("builds")
+expr_hash = build_manager.compile_expr(expr)
+
+roundtrip_expr = build_manager.load_expr(expr_hash)

pyproject.toml

@@ -28,6 +28,9 @@ dependencies = [
     "sqlglot==25.20.2",
     "toolz>=0.11",
     "typing-extensions>=4.3.0",
+    "hypothesis>=6.124.9",
+    "pyyaml>=6.0.2",
+    "cloudpickle>=3.1.1",
 ]
 requires-python = ">=3.10"
 authors = [

python/letsql/common/utils/graph_utils.py

@@ -0,0 +1,52 @@
+import xorq.expr.relations as rel
+
+
+def walk_nodes(node_types, expr):
+    def process_node(op):
+        match op:
+            case rel.RemoteTable():
+                yield op
+                yield from walk_nodes(
+                    node_types,
+                    op.remote_expr,
+                )
+            case rel.CachedNode():
+                yield op
+                yield from walk_nodes(
+                    node_types,
+                    op.parent,
+                )
+            case _:
+                yield from op.find(node_types)
+
+    def inner(rest, seen):
+        if not rest:
+            return seen
+        op = rest.pop()
+        seen.add(op)
+        new = process_node(op)
+        rest.update(set(new).difference(seen))
+        return inner(rest, seen)
+
+    rest = process_node(expr.op())
+    return inner(set(rest), set())
+
+
+def find_all_sources(expr):
+    import xorq.vendor.ibis.expr.operations as ops
+
+    node_types = (
+        ops.DatabaseTable,
+        ops.SQLQueryResult,
+        rel.CachedNode,
+        rel.Read,
+        rel.RemoteTable,
+        # ExprScalarUDF has an expr we need to get to
+        # FlightOperator has a dynamically generated connection: it should be passed a Profile instead
+    )
+    nodes = walk_nodes(node_types, expr)
+    sources = tuple(
+        source
+        for (source, _) in set((node.source, node.source._profile) for node in nodes)
+    )
+    return sources

python/xorq/common/utils/graph_utils.py

@@ -28,7 +28,8 @@ def inner(rest, seen):
         rest.update(set(new).difference(seen))
         return inner(rest, seen)
 
-    rest = process_node(expr.op())
+    initial_op = expr.op() if hasattr(expr, "op") else expr
+    rest = process_node(initial_op)
     return inner(set(rest), set())
 
 

python/xorq/ibis_yaml/common.py

@@ -0,0 +1,77 @@
+import functools
+from typing import Any
+
+import attr
+
+import xorq.vendor.ibis.expr.datatypes as dt
+import xorq.vendor.ibis.expr.types as ir
+from xorq.ibis_yaml.utils import freeze
+from xorq.vendor.ibis.common.collections import FrozenOrderedDict
+
+
+FROM_YAML_HANDLERS: dict[str, Any] = {}
+
+
+class SchemaRegistry:
+    def __init__(self):
+        self.schemas = {}
+        self.counter = 0
+
+    def register_schema(self, schema):
+        frozen_schema = freeze(
+            {name: _translate_type(dtype) for name, dtype in schema.items()}
+        )
+
+        for schema_id, existing_schema in self.schemas.items():
+            if existing_schema == frozen_schema:
+                return schema_id
+
+        schema_id = f"schema_{self.counter}"
+        self.schemas[schema_id] = frozen_schema
+        self.counter += 1
+        return schema_id
+
+    def _register_expr_schema(self, expr: ir.Expr) -> str:
+        if hasattr(expr, "schema"):
+            schema = expr.schema()
+            return self.register_schema(schema)
+        return None
+
+
+@attr.s(frozen=True)
+class TranslationContext:
+    schema_registry: SchemaRegistry = attr.ib(factory=SchemaRegistry)
+    profiles: FrozenOrderedDict = attr.ib(factory=FrozenOrderedDict)
+    definitions: FrozenOrderedDict = attr.ib(factory=lambda: freeze({"schemas": {}}))
+
+    def update_definitions(self, new_definitions: FrozenOrderedDict):
+        return attr.evolve(self, definitions=new_definitions)
+
+
+def register_from_yaml_handler(*op_names: str):
+    def decorator(func):
+        for name in op_names:
+            FROM_YAML_HANDLERS[name] = func
+        return func
+
+    return decorator
+
+
+@functools.cache
+@functools.singledispatch
+def translate_from_yaml(yaml_dict: dict, context: TranslationContext) -> Any:
+    op_type = yaml_dict["op"]
+    if op_type not in FROM_YAML_HANDLERS:
+        raise NotImplementedError(f"No handler for operation {op_type}")
+    return FROM_YAML_HANDLERS[op_type](yaml_dict, context)
+
+
+@functools.cache
+@functools.singledispatch
+def translate_to_yaml(op: Any, context: TranslationContext) -> dict:
+    raise NotImplementedError(f"No translation rule for {type(op)}")
+
+
+@functools.singledispatch
+def _translate_type(dtype: dt.DataType) -> dict:
+    return freeze({"name": type(dtype).__name__, "nullable": dtype.nullable})