diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py index f69c8757e809..6786c0d2b81e 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py +++ b/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) diff --git a/python/cudf_polars/tests/streaming/test_dataframescan.py b/python/cudf_polars/tests/streaming/test_dataframescan.py index a6e86045f2b6..a4e9ddb56236 100644 --- a/python/cudf_polars/tests/streaming/test_dataframescan.py +++ b/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),