From a63cee9ddb2ab20df0e3519df017fc3e02101dfd Mon Sep 17 00:00:00 2001
From: Phillip Cloud <417981+cpcloud@users.noreply.github.com>
Date: Wed, 31 Jul 2024 06:19:15 -0400
Subject: [PATCH] feat(datafusion): struct literals

---
 ibis/backends/sql/compilers/datafusion.py | 6 ++++++
 ibis/backends/tests/test_struct.py        | 8 +++++++-
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/ibis/backends/sql/compilers/datafusion.py b/ibis/backends/sql/compilers/datafusion.py
index 0e7ba7ae2aff..ac104494ba1a 100644
--- a/ibis/backends/sql/compilers/datafusion.py
+++ b/ibis/backends/sql/compilers/datafusion.py
@@ -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
 
diff --git a/ibis/backends/tests/test_struct.py b/ibis/backends/tests/test_struct.py
index c9fb76e49fe3..e1ba89162b2a 100644
--- a/ibis/backends/tests/test_struct.py
+++ b/ibis/backends/tests/test_struct.py
@@ -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"]),
 ]
 
 
@@ -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]
@@ -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"
@@ -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"
 )
@@ -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})
@@ -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(
         {
@@ -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})