Skip to content

Fix nullable Array columns in partitioned DataFrameScan #24131

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
0guban0v wants to merge 2 commits into
base: main
Choose a base branch
from
+41 −14
Open

Fix nullable Array columns in partitioned DataFrameScan #24131

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0e44ceb
Fix nullable Array streaming scans
0guban0v Sep 11, 2026
c8bc2b6
Simplify nested dtype discovery
0guban0v Sep 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 17 additions & 12 deletions python/cudf_polars/cudf_polars/streaming/actor_graph/io.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
# 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 = any(
isinstance(dt, pl.Struct)
for dt in pl.datatypes.unpack_dtypes(ir.df.dtypes(), include_compound=True)
for dt in pl.datatypes.unpack_dtypes(dtypes, include_compound=True)
)
array_columns = tuple(
name
for name, dtype in zip(ir.df.columns(), dtypes, strict=True)
if isinstance(dtype, pl.Array)
)

@TomAugspurger TomAugspurger Sep 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should be able to avoid looping over df.columns() / dtypes twice. Roughly:

Suggested change
has_struct = any(
isinstance(dt, pl.Struct)
for dt in pl.datatypes.unpack_dtypes(ir.df.dtypes(), include_compound=True)
for dt in pl.datatypes.unpack_dtypes(dtypes, include_compound=True)
)
array_columns = tuple(
name
for name, dtype in zip(ir.df.columns(), dtypes, strict=True)
if isinstance(dtype, pl.Array)
)
has_struct = False
array_columns = []
for name, dtype in zip(ir.df.columns(), dtypes, strict=True):
has_struct = has_struct or isinstance(pl.datatypes.unpack_dtypes(dtype, include_compound=True), pl.Struct)
if isinstance(dtype, pl.Array):
array_columns.append(name)

@0guban0v 0guban0v Sep 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, that makes sense. unpack_dtypes returns a set, I'll use any(...) for Struct check while combining both checks into one loop.


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.
Comment thread
coderabbitai[bot] marked this conversation as resolved.

@TomAugspurger TomAugspurger Sep 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's probably worth keeping the "doesn't handle the case where ..." caveat, assuming that's still true.

@0guban0v 0guban0v Sep 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I checked this with pre-sliced, nonzero-offset Struct and Array inputs. Even when streaming creates single partition, this path still serializes and deserializes the slice, which compacts input correctly. Old caveat is not true, I'd prefer not to restore it.

f = io.BytesIO()
sliced.serialize_binary(f)
f.seek(0)
Expand Down
22 changes: 22 additions & 0 deletions python/cudf_polars/tests/streaming/test_dataframescan.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down
Loading