cudf-polars: preserve zero-column row counts

python/cudf_polars/cudf_polars/containers/dataframe.py

@@ -93,7 +93,6 @@ class DataFrame:
     table: plc.Table
     columns: list[NamedColumn]
     stream: Stream
-    _num_rows_override: int | None
 
     def __init__(
         self, columns: Iterable[Column], stream: Stream, num_rows: int | None = None
@@ -104,24 +103,25 @@ def __init__(
         self.columns = [cast("NamedColumn", c) for c in columns]
         self.dtypes = [c.dtype for c in self.columns]
         self.column_map = {c.name: c for c in self.columns}
-        self.table = plc.Table([c.obj for c in self.columns])
+        self.table = plc.Table([c.obj for c in self.columns], num_rows=num_rows)
         self.stream = stream
-        self._num_rows_override = num_rows
 
     def copy(self) -> Self:
         """Return a shallow copy of self."""
         return type(self)(
             (c.copy() for c in self.columns),
             stream=self.stream,
-            num_rows=self._num_rows_override,
+            num_rows=self.num_rows,
         )
 
     def to_polars(self) -> pl.DataFrame:
         """Convert to a polars DataFrame."""
-        if self._num_rows_override is not None and len(self.column_map) == 0:
+        if len(self.column_map) == 0:
+            # polars < 1.38 has no DataFrame(height=...) constructor and cannot
+            # represent a zero-column frame with a non-zero row count.
             if POLARS_VERSION_LT_138:  # pragma: no cover
                 return pl.DataFrame()
-            return pl.DataFrame(height=self._num_rows_override)
+            return pl.DataFrame(height=self.num_rows)
 
         # If the arrow table has empty names, from_arrow produces
         # column_$i. But here we know there is only one such column
@@ -163,8 +163,6 @@ def num_columns(self) -> int:
     @cached_property
     def num_rows(self) -> int:
         """Number of rows."""
-        if self._num_rows_override is not None:
-            return self._num_rows_override
         return self.table.num_rows()
 
     @classmethod
@@ -195,6 +193,7 @@ def from_polars(cls, df: pl.DataFrame, stream: Stream) -> Self:
                 )
             ),
             stream=stream,
+            num_rows=plc_table.num_rows(),
         )
 
     @classmethod
@@ -204,7 +203,6 @@ def from_table(
         names: Sequence[str],
         dtypes: Sequence[DataType],
         stream: Stream,
-        num_rows: int | None = None,
     ) -> Self:
         """
         Create from a pylibcudf table.
@@ -221,10 +219,6 @@ def from_table(
             CUDA stream used for device memory operations and kernel launches
             on this dataframe. The caller is responsible for ensuring that
             the data in ``table`` is valid on ``stream``.
-        num_rows
-            Optional row count override for zero-width tables. Used to
-            preserve row count when zero-width tables lose their row count
-            during conversion. See https://github.com/rapidsai/cudf/issues/21428
 
         Returns
         -------
@@ -244,7 +238,7 @@ def from_table(
                 for c, name, dtype in zip(table.columns(), names, dtypes, strict=True)
             ),
             stream=stream,
-            num_rows=num_rows,
+            num_rows=table.num_rows(),
         )
 
     @classmethod
@@ -285,6 +279,8 @@ def deserialize(
                 for c, kw in zip(table.columns(), header["columns_kwargs"], strict=True)
             ),
             stream=stream,
+            # A zero-column frame's row count is carried by the packed metadata; preserve it.
+            num_rows=table.num_rows(),
         )
 
     def serialize(
@@ -358,6 +354,7 @@ def sorted_like(
                 for c, other in zip(self.columns, like.columns, strict=True)
             ),
             stream=self.stream,
+            num_rows=self.num_rows,
         )
 
     def with_columns(
@@ -396,20 +393,31 @@ def with_columns(
         new = {c.name: c for c in columns}
         if replace_only and not self.column_names_set.issuperset(new.keys()):
             raise ValueError("Cannot replace with non-existing names")
-        return type(self)((self.column_map | new).values(), stream=stream)
+        merged = self.column_map | new
+        # Only pass num_rows for a zero-column result. For results with columns, it
+        # must remain None because HStack(should_broadcast=False) intentionally
+        # produces mismatched column lengths that its Select parent reconciles later.
+        return type(self)(
+            merged.values(),
+            stream=stream,
+            num_rows=self.num_rows if not merged else None,
+        )
 
     def discard_columns(self, names: Set[str]) -> Self:
         """Drop columns by name."""
         return type(self)(
             (column for column in self.columns if column.name not in names),
             stream=self.stream,
+            num_rows=self.num_rows,
         )
 
     def select(self, names: Sequence[str] | Mapping[str, Any]) -> Self:
         """Select columns by name returning DataFrame."""
         try:
             return type(self)(
-                (self.column_map[name] for name in names), stream=self.stream
+                (self.column_map[name] for name in names),
+                stream=self.stream,
+                num_rows=self.num_rows,
             )
         except KeyError as e:
             raise ValueError("Can't select missing names") from e
@@ -419,6 +427,7 @@ def rename_columns(self, mapping: Mapping[str, str]) -> Self:
         return type(self)(
             (c.rename(mapping.get(c.name, c.name)) for c in self.columns),
             stream=self.stream,
+            num_rows=self.num_rows,
         )
 
     def select_columns(self, names: Set[str]) -> list[Column]:

python/cudf_polars/cudf_polars/dsl/ir.py

@@ -1136,19 +1136,22 @@ def read_csv_header(
                 table, names = cls._apply_parquet_projection(
                     plc.Table(concatenated_columns), names, with_columns
                 )
-                num_rows = (
-                    cls._get_parquet_row_count_from_metadata(
-                        paths, skip_rows, n_rows, parquet_options, cached_parquet_info
+                if not names:
+                    table = plc.Table(
+                        table.columns(),
+                        num_rows=cls._get_parquet_row_count_from_metadata(
+                            paths,
+                            skip_rows,
+                            n_rows,
+                            parquet_options,
+                            cached_parquet_info,
+                        ),
                     )
-                    if not names
-                    else None
-                )
                 df = DataFrame.from_table(
                     table,
                     names=names,
                     dtypes=[schema[name] for name in names],
                     stream=stream,
-                    num_rows=num_rows,
                 )
                 if include_file_paths is not None:
                     df = Scan.add_file_paths(  # pragma: no cover
@@ -1165,19 +1168,22 @@ def read_csv_header(
                 table, col_names = cls._apply_parquet_projection(
                     tbl_w_meta.tbl, col_names, with_columns
                 )
-                num_rows = (
-                    cls._get_parquet_row_count_from_metadata(
-                        paths, skip_rows, n_rows, parquet_options, cached_parquet_info
+                if not col_names:
+                    table = plc.Table(
+                        table.columns(),
+                        num_rows=cls._get_parquet_row_count_from_metadata(
+                            paths,
+                            skip_rows,
+                            n_rows,
+                            parquet_options,
+                            cached_parquet_info,
+                        ),
                     )
-                    if not col_names
-                    else None
-                )
                 df = DataFrame.from_table(
                     table,
                     col_names,
                     [schema[name] for name in col_names],
                     stream=stream,
-                    num_rows=num_rows,
                 )
                 if include_file_paths is not None:
                     df = Scan.add_file_paths(
@@ -3251,7 +3257,11 @@ def do_evaluate(
             target_length=df.num_rows,
             stream=df.stream,
         )
-        return DataFrame(columns, stream=df.stream)
+        return DataFrame(
+            columns,
+            stream=df.stream,
+            num_rows=df.num_rows if len(schema) == 0 else None,
+        )
 
 
 class MergeSorted(IR):

python/cudf_polars/cudf_polars/testing/inject_gpu_engine.py

@@ -282,14 +282,10 @@ def pytest_report_header(config: pytest.Config) -> str:
     "tests/unit/sql/test_window_functions.py::test_window_multiple_named_windows": "TODO: https://github.com/rapidsai/cudf/pull/22048#discussion_r3238041970",
     "tests/unit/sql/test_window_functions.py::test_window_frame_validation": "TODO: https://github.com/rapidsai/cudf/pull/22048#discussion_r3238041970",
     "tests/unit/operations/test_window.py::test_over_literal_cum_sum_26800": "TODO: https://github.com/rapidsai/cudf/pull/22048#discussion_r3238041970",
-    "tests/unit/sql/test_miscellaneous.py::test_select_output_heights_20058_21084[-WHERE a == 1 OR a != 1]": "column-less SELECT (always-true WHERE) loses its row count (https://github.com/rapidsai/cudf/issues/21428)",
-    "tests/unit/sql/test_miscellaneous.py::test_select_output_heights_20058_21084[ORDER BY 1-WHERE a == 1 OR a != 1]": "column-less SELECT (always-true WHERE) loses its row count (https://github.com/rapidsai/cudf/issues/21428)",
-    "tests/unit/sql/test_miscellaneous.py::test_select_output_heights_20058_21084[ORDER BY a-WHERE a == 1 OR a != 1]": "column-less SELECT (always-true WHERE) loses its row count (https://github.com/rapidsai/cudf/issues/21428)",
     "tests/unit/operations/namespaces/array/test_array.py::test_array_idx_size_limit_eval": "polars-internal IdxSize chunking debug assertion does not apply with the GPU engine",
     "tests/unit/operations/aggregation/test_aggregations.py::test_implode_and_agg": "implode + agg returns a mismatched dtype",
     "tests/unit/operations/aggregation/test_aggregations.py::test_duration_aggs": "Unsupported libcudf reduction operator for Duration dtype",
     "tests/unit/operations/aggregation/test_aggregations.py::test_boolean_aggs": "boolean-agg mean floating-point precision mismatch",
-    "tests/unit/lazyframe/test_projections.py::test_select_len_20337": "len() over a column-less input returns 0 (https://github.com/rapidsai/cudf/issues/21428)",
     "tests/unit/io/test_scan.py::test_scan_sink_metrics_multiple_phases": "sink metrics are not reported by the GPU engine",
     "tests/unit/io/test_parquet.py::test_read_parquet_legacy_nested_maps_27159": "legacy nested-map parquet read produces a mismatched result",
     "tests/unit/datatypes/test_struct.py::test_struct_equal_missing_null_25360": "struct equality with a null raises libcudf 'Index out of bounds' (get_element)",
@@ -451,7 +447,6 @@ def pytest_report_header(config: pytest.Config) -> str:
     "tests/unit/functions/range/test_linear_space.py::test_linear_space_num_samples_expr": "https://github.com/rapidsai/cudf/issues/22072",
     "tests/unit/functions/test_concat.py::test_concat_horizontal_zero_width_height_mismatch_26876": "https://github.com/rapidsai/cudf/issues/21644",
     "tests/unit/functions/test_concat.py::test_concat_horizontally_strict": "Correct polars.exceptions.ShapeError raised but it's in a ExceptionGroup",
-    "tests/unit/interop/test_interop.py::test_0_width_df_roundtrip": "https://github.com/rapidsai/cudf/issues/21644",
     "tests/unit/operations/test_slice.py::test_slice_pushdown_literal_projection_14349": "https://github.com/rapidsai/cudf/issues/22072",
     "tests/unit/operations/test_group_by.py::test_group_by_lit_series": "Incorrect broadcasting of literals in groupby-agg",
     "tests/unit/operations/test_group_by.py::test_group_by_series_partitioned": "https://github.com/rapidsai/cudf/issues/22072",
@@ -471,21 +466,7 @@ def pytest_report_header(config: pytest.Config) -> str:
     "tests/unit/sql/test_window_functions.py::test_window_multiple_named_window": "TODO: https://github.com/rapidsai/cudf/pull/22048#discussion_r3238041970",
     "tests/unit/functions/test_concat.py::test_concat_horizontal_lazy_strict_raises_shape_error_27415": "horizontal-concat strict height-mismatch raised inside an ExceptionGroup under the streaming engine",
     "tests/unit/io/test_io_plugin.py::test_defer_validate_true": "correct SchemaError raised but wrapped in an ExceptionGroup under the streaming engine",
-    "tests/unit/io/test_scan_lines.py::test_scan_lines[False-False-True]": "len() row count lost in zero-column streaming chunks (https://github.com/rapidsai/cudf/issues/21428)",
-    "tests/unit/io/test_scan_lines.py::test_scan_lines[False-True-True]": "len() row count lost in zero-column streaming chunks (https://github.com/rapidsai/cudf/issues/21428)",
-    "tests/unit/io/test_scan_lines.py::test_scan_lines[True-False-True]": "len() row count lost in zero-column streaming chunks (https://github.com/rapidsai/cudf/issues/21428)",
-    "tests/unit/io/test_scan_lines.py::test_scan_lines[True-True-True]": "len() row count lost in zero-column streaming chunks (https://github.com/rapidsai/cudf/issues/21428)",
-    "tests/unit/lazyframe/test_lazyframe.py::test_len": "len() row count lost in zero-column streaming chunks (https://github.com/rapidsai/cudf/issues/21428)",
-    "tests/unit/lazyframe/test_projections.py::test_projection_pushdown_select_len": "len() row count lost in zero-column streaming chunks (https://github.com/rapidsai/cudf/issues/21428)",
-    "tests/unit/operations/test_scalar.py::test_scalar_len_20046": "len() row count lost in zero-column streaming chunks (https://github.com/rapidsai/cudf/issues/21428)",
     "tests/unit/operations/test_slice.py::test_hconcat_tail_unequal_heights_strict_raises_27552": "horizontal-concat strict height-mismatch raised inside an ExceptionGroup under the streaming engine",
-    "tests/unit/sql/test_group_by.py::test_group_by_empty_or_scalar_key_exprs_23397": "len() row count lost in zero-column streaming chunks (https://github.com/rapidsai/cudf/issues/21428)",
-    "tests/unit/sql/test_miscellaneous.py::test_select_output_heights_20058_21084[-WHERE 1 = 1]": "row count lost in zero-column streaming chunk; RapidsMPF cannot pack the empty table (https://github.com/rapidsai/cudf/issues/21428)",
-    "tests/unit/sql/test_miscellaneous.py::test_select_output_heights_20058_21084[-]": "row count lost in zero-column streaming chunk; RapidsMPF cannot pack the empty table (https://github.com/rapidsai/cudf/issues/21428)",
-    "tests/unit/sql/test_miscellaneous.py::test_select_output_heights_20058_21084[ORDER BY 1-WHERE 1 = 1]": "row count lost in zero-column streaming chunk; RapidsMPF cannot pack the empty table (https://github.com/rapidsai/cudf/issues/21428)",
-    "tests/unit/sql/test_miscellaneous.py::test_select_output_heights_20058_21084[ORDER BY 1-]": "row count lost in zero-column streaming chunk; RapidsMPF cannot pack the empty table (https://github.com/rapidsai/cudf/issues/21428)",
-    "tests/unit/sql/test_miscellaneous.py::test_select_output_heights_20058_21084[ORDER BY a-WHERE 1 = 1]": "row count lost in zero-column streaming chunk; RapidsMPF cannot pack the empty table (https://github.com/rapidsai/cudf/issues/21428)",
-    "tests/unit/sql/test_miscellaneous.py::test_select_output_heights_20058_21084[ORDER BY a-]": "row count lost in zero-column streaming chunk; RapidsMPF cannot pack the empty table (https://github.com/rapidsai/cudf/issues/21428)",
 }
 
 

python/cudf_polars/tests/containers/test_dataframe.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 from __future__ import annotations
@@ -125,6 +125,25 @@ def test_shallow_copy():
     assert copy.column_map["a"].is_sorted == plc.types.Sorted.NO
 
 
+def test_with_columns_preserves_zero_column_row_count():
+    stream = get_cuda_stream()
+    dtype = DataType(pl.Int8())
+    column = Column(
+        plc.column_factories.make_numeric_column(
+            dtype.plc_type, 5, plc.MaskState.ALL_VALID, stream=stream
+        ),
+        dtype=dtype,
+        name="a",
+    )
+    zero_col = DataFrame([column], stream=stream).discard_columns({"a"})
+    assert zero_col.num_columns == 0
+    assert zero_col.num_rows == 5
+    # Adding no columns to a zero-column frame must keep the row count.
+    result = zero_col.with_columns([], stream=stream)
+    assert result.num_columns == 0
+    assert result.num_rows == 5
+
+
 def test_sorted_flags_preserved_empty():
     stream = get_cuda_stream()
     df = pl.DataFrame({"a": pl.Series([], dtype=pl.Int8())})

python/cudf_polars/tests/streaming/test_select.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 from __future__ import annotations
@@ -21,7 +21,7 @@
 from cudf_polars.testing.asserts import (
     assert_gpu_result_equal,
 )
-from cudf_polars.testing.engine_utils import is_streaming_engine, warns_on_spmd
+from cudf_polars.testing.engine_utils import warns_on_spmd
 from cudf_polars.utils.versions import POLARS_VERSION_LT_141
 
 
@@ -119,20 +119,8 @@ def test_select_fill_null_with_strategy(df, streaming_engine_factory):
         (pl.col("a").min(), pl.col("b"), pl.col("c").max()),
     ],
 )
-def test_select_aggs(df, engine, aggs, request):
+def test_select_aggs(df, engine, aggs):
     # Test supported aggs (e.g. "min", "max", "mean", "n_unique")
-    if (
-        not POLARS_VERSION_LT_141
-        and is_streaming_engine(engine)
-        and len(aggs) == 1
-        and "len()" in str(aggs[0])
-    ):
-        request.applymarker(
-            pytest.mark.xfail(
-                reason="len() row count lost in zero-column streaming chunks "
-                "(https://github.com/rapidsai/cudf/issues/21428)"
-            )
-        )
     query = df.select(*aggs)
     assert_gpu_result_equal(query, engine=engine)
 

python/cudf_polars/tests/test_dataframescan.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 from __future__ import annotations
@@ -13,7 +13,6 @@
     assert_gpu_result_equal,
     assert_ir_translation_raises,
 )
-from cudf_polars.testing.engine_utils import is_streaming_engine
 from cudf_polars.utils.versions import POLARS_VERSION_LT_138
 
 
@@ -88,13 +87,7 @@ def test_dataframescan_with_decimals(engine: pl.GPUEngine):
     POLARS_VERSION_LT_138,
     reason="height parameter added in Polars 1.38",
 )
-def test_dataframescan_zero_width_with_rows(engine: pl.GPUEngine, request):
-    request.applymarker(
-        pytest.mark.xfail(
-            is_streaming_engine(engine),
-            reason="https://github.com/rapidsai/cudf/issues/21644",
-        )
-    )
+def test_dataframescan_zero_width_with_rows(engine: pl.GPUEngine):
     df = pl.LazyFrame(height=5)
     q = df.select(pl.len())
     assert_gpu_result_equal(q, engine=engine)

python/cudf_polars/tests/test_groupby.py

@@ -24,7 +24,6 @@
 from cudf_polars.utils.versions import (
     POLARS_VERSION_LT_136,
     POLARS_VERSION_LT_140,
-    POLARS_VERSION_LT_141,
 )
 
 
@@ -721,17 +720,10 @@ def test_groupby_literal_agg(engine: pl.GPUEngine):
     assert_gpu_result_equal(q, engine=engine, check_row_order=False)
 
 
-def test_groupby_empty_keys_raises(engine: pl.GPUEngine, request):
+def test_groupby_empty_keys_raises(engine: pl.GPUEngine):
     df = pl.LazyFrame({"x": [1, 2, 3]})
     q = df.group_by([]).agg(pl.len())
     if POLARS_VERSION_LT_140:
         assert_ir_translation_raises(q, engine, NotImplementedError)
     else:
-        if not POLARS_VERSION_LT_141 and is_streaming_engine(engine):
-            request.applymarker(
-                pytest.mark.xfail(
-                    reason="len() row count lost in zero-column streaming chunks "
-                    "(https://github.com/rapidsai/cudf/issues/21428)"
-                )
-            )
         assert_gpu_result_equal(q, engine=engine)