diff --git a/mplang/backends/table_impl.py b/mplang/backends/table_impl.py index 2893500d..d5eae462 100644 --- a/mplang/backends/table_impl.py +++ b/mplang/backends/table_impl.py @@ -609,7 +609,12 @@ def _unwrap(val: TableValue | pa.Table | pd.DataFrame) -> pa.Table: @table.run_sql_p.def_impl def run_sql_impl(interpreter: Interpreter, op: Operation, *args: Any) -> TableValue: - """Execute SQL query on input tables.""" + """Execute SQL query on input tables. + + Query relations from a DuckDB state other than the selected execution state + are materialized as Arrow tables before registration. This may use + substantial memory for large intermediate relations. + """ query = op.attrs["query"] dialect = op.attrs.get("dialect", "duckdb") table_names = op.attrs["table_names"] @@ -626,8 +631,6 @@ def run_sql_impl(interpreter: Interpreter, op: Operation, *args: Any) -> TableVa if isinstance(data, QueryTableSource): if state is None: state = data.state - elif state != data.state: - raise ValueError("All tables must belong to the same DuckDB connection") if state is None: conn = duckdb.connect() @@ -638,6 +641,8 @@ def run_sql_impl(interpreter: Interpreter, op: Operation, *args: Any) -> TableVa # register tables or create view for name, tbl in zip(table_names, tables, strict=True): data = tbl.unwrap() + if isinstance(data, QueryTableSource) and data.state is not state: + data = tbl.data if name in state.tables: if state.tables[name] is not data: # TODO: rename and rewrite sql?? diff --git a/tests/backends/test_table_impl.py b/tests/backends/test_table_impl.py index 7847f8c9..c3176163 100644 --- a/tests/backends/test_table_impl.py +++ b/tests/backends/test_table_impl.py @@ -15,12 +15,20 @@ """Tests for Table Runtime Implementation.""" import os +from types import SimpleNamespace +import duckdb import pyarrow as pa import pytest import mplang.edsl.typing as elt -from mplang.backends.table_impl import ParquetReader, TableValue +from mplang.backends.table_impl import ( + DuckDBState, + ParquetReader, + QueryTableSource, + TableValue, + run_sql_impl, +) from mplang.dialects import table @@ -81,6 +89,175 @@ def workload(): os.remove(path) +def test_run_sql_materializes_query_source_from_different_state(): + """Foreign DuckDB relations are materialized into the selected state.""" + local_conn = duckdb.connect() + foreign_conn = duckdb.connect() + try: + local_state = DuckDBState(local_conn) + foreign_state = DuckDBState(foreign_conn) + local = TableValue( + QueryTableSource( + local_conn.sql("SELECT 1 AS id, 'local' AS local_value"), + local_state, + ) + ) + foreign = TableValue( + QueryTableSource( + foreign_conn.sql("SELECT 1 AS id, 'foreign' AS foreign_value"), + foreign_state, + ) + ) + op = SimpleNamespace( + attrs={ + "query": ( + "SELECT local_t.id, local_value, foreign_value " + "FROM local_t JOIN foreign_t USING (id)" + ), + "dialect": "duckdb", + "table_names": ["local_t", "foreign_t"], + } + ) + + result = run_sql_impl(None, op, local, foreign) + + assert result.data == pa.table({ + "id": pa.array([1], type=pa.int32()), + "local_value": ["local"], + "foreign_value": ["foreign"], + }) + assert isinstance(local.unwrap(), QueryTableSource) + assert isinstance(foreign.unwrap(), pa.Table) + assert local_state.tables["foreign_t"] is foreign.data + finally: + local_conn.close() + foreign_conn.close() + + +def test_run_sql_keeps_same_state_query_sources_lazy(): + """Relations from the execution state keep the lazy view path.""" + conn = duckdb.connect() + try: + state = DuckDBState(conn) + left = TableValue( + QueryTableSource(conn.sql("SELECT 1 AS id, 'left' AS left_value"), state) + ) + right = TableValue( + QueryTableSource(conn.sql("SELECT 1 AS id, 'right' AS right_value"), state) + ) + op = SimpleNamespace( + attrs={ + "query": ( + "SELECT left_t.id, left_value, right_value " + "FROM left_t JOIN right_t USING (id)" + ), + "dialect": "duckdb", + "table_names": ["left_t", "right_t"], + } + ) + + result = run_sql_impl(None, op, left, right) + + assert result.data == pa.table({ + "id": pa.array([1], type=pa.int32()), + "left_value": ["left"], + "right_value": ["right"], + }) + assert isinstance(left.unwrap(), QueryTableSource) + assert isinstance(right.unwrap(), QueryTableSource) + assert state.tables["left_t"] is left.unwrap() + assert state.tables["right_t"] is right.unwrap() + finally: + conn.close() + + +def test_run_sql_reuses_materialized_foreign_table_with_same_name(): + """The registered Arrow object preserves same-name identity validation.""" + local_conn = duckdb.connect() + foreign_conn = duckdb.connect() + try: + local_state = DuckDBState(local_conn) + foreign_state = DuckDBState(foreign_conn) + local = TableValue( + QueryTableSource( + local_conn.sql("SELECT 1 AS id, 'local' AS local_value"), + local_state, + ) + ) + foreign = TableValue( + QueryTableSource( + foreign_conn.sql("SELECT 1 AS id, 'foreign' AS foreign_value"), + foreign_state, + ) + ) + op = SimpleNamespace( + attrs={ + "query": "SELECT * FROM local_t JOIN foreign_t USING (id)", + "dialect": "duckdb", + "table_names": ["local_t", "foreign_t"], + } + ) + + run_sql_impl(None, op, local, foreign) + result = run_sql_impl(None, op, local, foreign) + + assert result.data.to_pylist() == [ + { + "id": 1, + "local_value": "local", + "foreign_value": "foreign", + } + ] + finally: + local_conn.close() + foreign_conn.close() + + +def test_run_sql_rejects_different_table_with_registered_name(): + """A different object cannot replace a registered foreign table name.""" + local_conn = duckdb.connect() + foreign_conn = duckdb.connect() + try: + local_state = DuckDBState(local_conn) + foreign_state = DuckDBState(foreign_conn) + local = TableValue( + QueryTableSource( + local_conn.sql("SELECT 1 AS id, 'local' AS local_value"), + local_state, + ) + ) + foreign = TableValue( + QueryTableSource( + foreign_conn.sql("SELECT 1 AS id, 'foreign' AS foreign_value"), + foreign_state, + ) + ) + op = SimpleNamespace( + attrs={ + "query": "SELECT * FROM local_t JOIN foreign_t USING (id)", + "dialect": "duckdb", + "table_names": ["local_t", "foreign_t"], + } + ) + + run_sql_impl(None, op, local, foreign) + replacement = TableValue( + pa.table({ + "id": pa.array([1], type=pa.int32()), + "foreign_value": ["replacement"], + }) + ) + + with pytest.raises(RuntimeError) as error: + run_sql_impl(None, op, local, replacement) + + assert isinstance(error.value.__cause__, ValueError) + assert str(error.value.__cause__) == "foreign_t has been registered." + finally: + local_conn.close() + foreign_conn.close() + + def test_table_constant_dataframe(): """Test creating constant table from DataFrame.""" import pandas as pd