feat(bigquery): implement ops.StructColumn, ops.ArrayColumn

ibis-project · Jan 6, 2023 · 2bbf73c · 2bbf73c
1 parent bd1c637
commit 2bbf73c
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 2 deletions.
diff --git a/ibis/backends/bigquery/registry.py b/ibis/backends/bigquery/registry.py
@@ -88,10 +88,22 @@ def _struct_field(translator, op):
     return f"{arg}.`{op.field}`"
 
 
+def _struct_column(translator, op):
+    cols = (
+        f'{translator.translate(value)} AS {name}'
+        for name, value, in zip(op.names, op.values)
+    )
+    return "STRUCT({})".format(", ".join(cols))
+
+
 def _array_concat(translator, op):
     return "ARRAY_CONCAT({})".format(", ".join(map(translator.translate, op.args)))
 
 
+def _array_column(translator, op):
+    return "[{}]".format(", ".join(map(translator.translate, op.cols)))
+
+
 def _array_index(translator, op):
     # SAFE_OFFSET returns NULL if out of bounds
     arg = translator.translate(op.arg)
@@ -595,8 +607,10 @@ def _clip(t, op):
     ops.GroupConcat: reduction("STRING_AGG"),
     ops.Cast: _cast,
     ops.StructField: _struct_field,
+    ops.StructColumn: _struct_column,
     ops.ArrayCollect: _array_agg,
     ops.ArrayConcat: _array_concat,
+    ops.ArrayColumn: _array_column,
     ops.ArrayIndex: _array_index,
     ops.ArrayLength: unary("ARRAY_LENGTH"),
     ops.ArrayRepeat: _array_repeat,

diff --git a/ibis/backends/tests/test_array.py b/ibis/backends/tests/test_array.py
@@ -22,7 +22,7 @@
 ]
 
 
-@pytest.mark.notimpl(["bigquery", "datafusion", "snowflake"])
+@pytest.mark.notimpl(["datafusion", "snowflake"])
 def test_array_column(backend, alltypes, df):
     expr = ibis.array([alltypes['double_col'], alltypes['double_col']])
     assert isinstance(expr, ir.ArrayColumn)

diff --git a/ibis/backends/tests/test_struct.py b/ibis/backends/tests/test_struct.py
@@ -87,7 +87,7 @@ def test_null_literal(con, field):
     tm.assert_series_equal(result, expected)
 
 
-@pytest.mark.notimpl(["bigquery", "dask", "pandas", "postgres", "snowflake"])
+@pytest.mark.notimpl(["dask", "pandas", "postgres", "snowflake"])
 def test_struct_column(alltypes, df):
     t = alltypes
     expr = ibis.struct(dict(a=t.string_col, b=1, c=t.bigint_col)).name("s")