From 37e327961b59a1a6185d5483a97879bf1deb4934 Mon Sep 17 00:00:00 2001 From: Krzysztof Date: Sat, 28 Jan 2023 04:13:49 +0100 Subject: [PATCH] feat(bigquery): support in-memory tables --- ibis/backends/base/sql/compiler/query_builder.py | 9 +++++++-- ibis/backends/bigquery/compiler.py | 2 ++ ibis/backends/tests/test_client.py | 8 ++++---- 3 files changed, 13 insertions(+), 6 deletions(-) diff --git a/ibis/backends/base/sql/compiler/query_builder.py b/ibis/backends/base/sql/compiler/query_builder.py index b3760a0b0d3a..6b862290b520 100644 --- a/ibis/backends/base/sql/compiler/query_builder.py +++ b/ibis/backends/base/sql/compiler/query_builder.py @@ -80,8 +80,12 @@ def _format_in_memory_table(self, op): ) for row in op.data.to_frame().itertuples(index=False) ) - rows = ", ".join(f"({raw_row})" for raw_row in raw_rows) - return f"(VALUES {rows})" + if self.context.compiler.support_values_syntax_in_select: + rows = ", ".join(f"({raw_row})" for raw_row in raw_rows) + return f"(VALUES {rows})" + else: + rows = "UNION ALL ".join(f"(SELECT {raw_row})" for raw_row in raw_rows) + return f"({rows})" def _format_table(self, op): # TODO: This could probably go in a class and be significantly nicer @@ -488,6 +492,7 @@ class Compiler: difference_class = Difference cheap_in_memory_tables = False + support_values_syntax_in_select = True @classmethod def make_context(cls, params=None): diff --git a/ibis/backends/bigquery/compiler.py b/ibis/backends/bigquery/compiler.py index 54fe68c55899..fd67d89d9595 100644 --- a/ibis/backends/bigquery/compiler.py +++ b/ibis/backends/bigquery/compiler.py @@ -111,6 +111,8 @@ class BigQueryCompiler(sql_compiler.Compiler): intersect_class = BigQueryIntersection difference_class = BigQueryDifference + support_values_syntax_in_select = False + @staticmethod def _generate_setup_queries(expr, context): """Generate DDL for temporary resources.""" diff --git a/ibis/backends/tests/test_client.py b/ibis/backends/tests/test_client.py index 14091434d5c0..00689c1a198c 100644 --- a/ibis/backends/tests/test_client.py +++ b/ibis/backends/tests/test_client.py @@ -692,11 +692,11 @@ def test_deprecated_path_argument(backend_name, tmp_path): ], ) @pytest.mark.notyet( - ["bigquery", "mysql", "sqlite"], + ["mysql", "sqlite"], reason="SQLAlchemy generates incorrect code for `VALUES` projections.", raises=(sa.exc.ProgrammingError, sa.exc.OperationalError), ) -@pytest.mark.notimpl(["bigquery", "dask", "datafusion", "pandas"]) +@pytest.mark.notimpl(["dask", "datafusion", "pandas"]) def test_in_memory_table(backend, con, expr, expected): result = con.execute(expr) backend.assert_frame_equal(result, expected) @@ -707,7 +707,7 @@ def test_in_memory_table(backend, con, expr, expected): reason="SQLAlchemy generates incorrect code for `VALUES` projections.", raises=(sa.exc.ProgrammingError, sa.exc.OperationalError), ) -@pytest.mark.notimpl(["bigquery", "dask", "datafusion", "pandas"]) +@pytest.mark.notimpl(["dask", "datafusion", "pandas"]) def test_filter_memory_table(backend, con): t = ibis.memtable([(1, 2), (3, 4), (5, 6)], columns=["x", "y"]) expr = t.filter(t.x > 1) @@ -721,7 +721,7 @@ def test_filter_memory_table(backend, con): reason="SQLAlchemy generates incorrect code for `VALUES` projections.", raises=(sa.exc.ProgrammingError, sa.exc.OperationalError), ) -@pytest.mark.notimpl(["bigquery", "dask", "datafusion", "pandas"]) +@pytest.mark.notimpl(["dask", "datafusion", "pandas"]) def test_agg_memory_table(con): t = ibis.memtable([(1, 2), (3, 4), (5, 6)], columns=["x", "y"]) expr = t.x.count()