Fix nullable Array columns in partitioned DataFrameScan

python/cudf_polars/cudf_polars/streaming/actor_graph/io.py

@@ -227,26 +227,31 @@ async def dataframescan_node(
 
         # Build list of IR slices to read
         ir_slices = []
-        # Partial workaround for
-        # https://github.com/pola-rs/polars/issues/23214 If a struct column
-        # has nulls and is sliced then polars exports invalid validity
-        # buffers. We can't detect this exact state because we can't know
-        # when the column is sliced.
-        copy_slice = any(
-            isinstance(dt, pl.Struct)
-            for dt in pl.datatypes.unpack_dtypes(ir.df.dtypes(), include_compound=True)
-        )
+        # Partial workarounds for sliced nested columns. Polars exports invalid
+        # validity buffers for struct columns with nulls
+        # (https://github.com/pola-rs/polars/issues/23214), and double-counts
+        # offsets for Array columns with outer nulls
+        # (https://github.com/pola-rs/polars/pull/28602).
+        dtypes = ir.df.dtypes()
+        has_struct = False
+        array_columns = []
+        for name, dtype in zip(ir.df.columns(), dtypes, strict=True):
+            has_struct = has_struct or any(
+                isinstance(dt, pl.Struct)
+                for dt in pl.datatypes.unpack_dtypes(dtype, include_compound=True)
+            )
+            if isinstance(dtype, pl.Array):
+                array_columns.append(name)
 
         for seq_num in range(local_count):
             offset = local_offset * rows_per_partition + seq_num * rows_per_partition
             if offset >= nrows:
                 break
             sliced = ir.df.slice(offset, rows_per_partition)
-            if copy_slice:
-                # OK, we have structs that might have nulls, and we're
-                # slicing. So let's copy to contiguous storage. This is
-                # hacky and doesn't handle the case where we didn't slice
-                # but the user sliced the input.
+            if has_struct or any(
+                sliced.get_column(name).null_count() > 0 for name in array_columns
+            ):
+                # Copy the affected slice to contiguous storage before Arrow export.
                 f = io.BytesIO()
                 sliced.serialize_binary(f)
                 f.seek(0)

python/cudf_polars/tests/streaming/test_dataframescan.py

@@ -77,6 +77,28 @@ def test_parallel_dataframescan(
         assert count == 1
 
 
+def test_nullable_array_dataframescan(streaming_engine_factory):
+    streaming_engine = streaming_engine_factory(
+        StreamingOptions(max_rows_per_partition=2, fallback_mode="raise"),
+    )
+    q = pl.LazyFrame(
+        {
+            "embedding": pl.Series(
+                # The outer null is in the nonzero-offset second partition.
+                [
+                    [0.0, 1.0],
+                    [2.0, None],
+                    None,
+                    [3.0, 4.0],
+                ],
+                dtype=pl.Array(pl.Float32, 2),
+            )
+        }
+    )
+
+    assert_gpu_result_equal(q, engine=streaming_engine)
+
+
 def test_dataframescan_concat(request, df, streaming_engine_factory):
     streaming_engine = streaming_engine_factory(
         StreamingOptions(max_rows_per_partition=1_000),