From 065da5d23e9eca3c6b6473b199ecb590aa5fa8b0 Mon Sep 17 00:00:00 2001
From: Nick Crews <nicholas.b.crews@gmail.com>
Date: Sun, 30 Jun 2024 10:26:40 -0800
Subject: [PATCH] fix(polars): add workaround to compile Array<Array> correctly

---
 ibis/backends/polars/compiler.py  | 7 +++++++
 ibis/backends/tests/test_array.py | 9 +--------
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/ibis/backends/polars/compiler.py b/ibis/backends/polars/compiler.py
index 4d9a497191b4..3b0aec010d11 100644
--- a/ibis/backends/polars/compiler.py
+++ b/ibis/backends/polars/compiler.py
@@ -976,6 +976,13 @@ def array_concat(op, **kw):
 @translate.register(ops.Array)
 def array_column(op, **kw):
     cols = [translate(col, **kw) for col in op.exprs]
+    # Workaround for https://github.com/pola-rs/polars/issues/17294
+    # pl.concat_list(Iterable[T]) results in pl.List[T], EXCEPT when T is a
+    # pl.List, in which case pl.concat_list(Iterable[pl.List[T]]) results in pl.List[T].
+    # If polars ever supports a more consistent array constructor,
+    # we should switch to that.
+    if op.dtype.is_array():
+        cols = [c.implode() for c in cols]
     return pl.concat_list(cols)
 
 
diff --git a/ibis/backends/tests/test_array.py b/ibis/backends/tests/test_array.py
index 8b55c189e9d6..4103d7cf817a 100644
--- a/ibis/backends/tests/test_array.py
+++ b/ibis/backends/tests/test_array.py
@@ -1352,14 +1352,7 @@ def test_unnest_range(con):
             [[1], ibis.literal([2])],
             [[1], [2]],
             id="array",
-            marks=[
-                pytest.mark.notyet(["bigquery"], raises=GoogleBadRequest),
-                pytest.mark.broken(
-                    ["polars"],
-                    reason="expression input not supported with nested arrays",
-                    raises=TypeError,
-                ),
-            ],
+            marks=[pytest.mark.notyet(["bigquery"], raises=GoogleBadRequest)],
         ),
     ],
 )