Skip to content

Commit

Permalink
feat(datafusion): struct literals
Browse files Browse the repository at this point in the history
  • Loading branch information
cpcloud authored and jcrist committed Jul 31, 2024
1 parent 2330b0c commit a63cee9
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 1 deletion.
6 changes: 6 additions & 0 deletions ibis/backends/sql/compilers/datafusion.py
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,12 @@ def visit_NonNullLiteral(self, op, *, value, dtype):
return sg.exp.HexString(this=value.hex())
elif dtype.is_uuid():
return sge.convert(str(value))
elif dtype.is_struct():
args = []
for name, field_value in value.items():
args.append(sge.convert(name))
args.append(field_value)
return self.f.named_struct(*args)
else:
return None

Expand Down
8 changes: 7 additions & 1 deletion ibis/backends/tests/test_struct.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@
pytestmark = [
pytest.mark.never(["mysql", "sqlite", "mssql"], reason="No struct support"),
pytest.mark.notyet(["impala"]),
pytest.mark.notimpl(["datafusion", "druid", "oracle", "exasol"]),
pytest.mark.notimpl(["druid", "oracle", "exasol"]),
]


Expand Down Expand Up @@ -76,6 +76,7 @@ def test_all_fields(struct, struct_df):


@pytest.mark.notimpl(["postgres", "risingwave"])
@pytest.mark.notyet(["datafusion"], raises=Exception, reason="unsupported syntax")
@pytest.mark.parametrize("field", ["a", "b", "c"])
def test_literal(backend, con, field):
query = _STRUCT_LITERAL[field]
Expand All @@ -87,6 +88,7 @@ def test_literal(backend, con, field):


@pytest.mark.notimpl(["postgres"])
@pytest.mark.notyet(["datafusion"], raises=Exception, reason="unsupported syntax")
@pytest.mark.parametrize("field", ["a", "b", "c"])
@pytest.mark.notyet(
["clickhouse"], reason="clickhouse doesn't support nullable nested types"
Expand All @@ -112,6 +114,7 @@ def test_struct_column(alltypes, df):


@pytest.mark.notimpl(["postgres", "risingwave", "polars"])
@pytest.mark.notyet(["datafusion"], raises=Exception, reason="unsupported syntax")
@pytest.mark.notyet(
["flink"], reason="flink doesn't support creating struct columns from collect"
)
Expand Down Expand Up @@ -144,6 +147,7 @@ def test_collect_into_struct(alltypes):
reason="struct literals not implemented",
raises=PsycoPg2InternalError,
)
@pytest.mark.notyet(["datafusion"], raises=Exception, reason="unsupported syntax")
@pytest.mark.notimpl(["flink"], raises=Py4JJavaError, reason="not implemented in ibis")
def test_field_access_after_case(con):
s = ibis.struct({"a": 3})
Expand Down Expand Up @@ -201,6 +205,7 @@ def test_field_access_after_case(con):
raises=AssertionError,
reason="snowflake doesn't have strongly typed structs",
)
@pytest.mark.notyet(["datafusion"], raises=Exception, reason="unsupported syntax")
def test_keyword_fields(con, nullable):
schema = ibis.schema(
{
Expand Down Expand Up @@ -260,6 +265,7 @@ def test_keyword_fields(con, nullable):
raises=Py4JJavaError,
reason="fails to parse due to an unsupported operation; flink docs say the syntax is supported",
)
@pytest.mark.notyet(["datafusion"], raises=Exception, reason="unsupported syntax")
def test_isin_struct(con):
needle1 = ibis.struct({"x": 1, "y": 2})
needle2 = ibis.struct({"x": 2, "y": 3})
Expand Down

0 comments on commit a63cee9

Please sign in to comment.