Avoid leaving multi-rank execution in a bad state in cudf-polars tests (NVIDIA#23340)

If a cudf-polars query raises an exception during execution, it is possible that it can leave a dangling collective that then tears down one process.

This is racy because it depends on tasks being cancelled in Python and then dropping C++ objects in a particular order.

The one test in the cudf-polars tests suite that could do this is fixed on main, so backport the relevant changes (NVIDIA#23235), and un-xfail the test.

Authors:
  - Lawrence Mitchell (https://github.com/wence-)
  - Mads R. B. Kristensen (https://github.com/madsbk)

Approvers:
  - Richard (Rick) Zamora (https://github.com/rjzamora)
  - Matthew Murray (https://github.com/Matt711)
  - Bradley Dice (https://github.com/bdice)
  - Vyas Ramasubramani (https://github.com/vyasr)

URL: NVIDIA#23340

Author: wence-

cpp/include/cudf/join/join.hpp

@@ -248,6 +248,10 @@ full_join(cudf::table_view const& left_keys,
  *
  * The cross join returns the cartesian product of rows from each table.
  *
+ * The result has `left.num_rows() * right.num_rows()` rows and
+ * `left.num_columns() + right.num_columns()` columns. Either operand may have
+ * zero columns and still contribute its row count.
+ *
  * @note Warning: This function can easily cause out-of-memory errors. The size of the output is
  * equal to `left.num_rows() * right.num_rows()`. Use with caution.
  *
@@ -256,8 +260,9 @@ full_join(cudf::table_view const& left_keys,
  * Right b: {3, 4, 5}
  * Result: { a: {0, 0, 0, 1, 1, 1, 2, 2, 2}, b: {3, 4, 5, 3, 4, 5, 3, 4, 5} }
  * @endcode
-
- * @throw cudf::logic_error if the number of columns in either `left` or `right` table is 0
+ *
+ * @throw std::overflow_error if `left.num_rows() * right.num_rows()` exceeds the maximum
+ * number of rows a column can hold.
  *
  * @param left  The left table
  * @param right The right table

cpp/src/join/cross_join.cu

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -19,7 +19,10 @@
 
 #include <rmm/cuda_stream_view.hpp>
 
+#include <cstdint>
+#include <limits>
 #include <memory>
+#include <stdexcept>
 
 namespace cudf {
 namespace detail {
@@ -33,9 +36,6 @@ std::unique_ptr<cudf::table> cross_join(cudf::table_view const& left,
                                         rmm::cuda_stream_view stream,
                                         rmm::device_async_resource_ref mr)
 {
-  CUDF_EXPECTS(0 != left.num_columns(), "Left table is empty");
-  CUDF_EXPECTS(0 != right.num_columns(), "Right table is empty");
-
   // If left or right table has no rows, return an empty table with all columns
   if ((0 == left.num_rows()) || (0 == right.num_rows())) {
     auto empty_left_columns  = empty_like(left)->release();
@@ -46,6 +46,12 @@ std::unique_ptr<cudf::table> cross_join(cudf::table_view const& left,
     return std::make_unique<table>(std::move(empty_left_columns));
   }
 
+  auto const out_num_rows =
+    static_cast<int64_t>(left.num_rows()) * static_cast<int64_t>(right.num_rows());
+  CUDF_EXPECTS(out_num_rows <= std::numeric_limits<size_type>::max(),
+               "Cross join result exceeds the column size limit",
+               std::overflow_error);
+
   // Repeat left table
   auto left_repeated = detail::repeat(left, right.num_rows(), stream, mr);
 
@@ -59,7 +65,8 @@ std::unique_ptr<cudf::table> cross_join(cudf::table_view const& left,
             right_tiled_columns.end(),
             std::back_inserter(left_repeated_columns));
 
-  return std::make_unique<table>(std::move(left_repeated_columns));
+  return std::make_unique<table>(std::move(left_repeated_columns),
+                                 static_cast<size_type>(out_num_rows));
 }
 }  // namespace detail
 

cpp/tests/join/cross_join_tests.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -14,6 +14,9 @@
 #include <cudf/table/table.hpp>
 #include <cudf/table/table_view.hpp>
 
+#include <limits>
+#include <stdexcept>
+
 template <typename T, typename SourceT = T>
 using column_wrapper = cudf::test::fixed_width_column_wrapper<T, SourceT>;
 
@@ -67,26 +70,47 @@ TYPED_TEST(CrossJoinTypeTests, CrossJoin)
   CUDF_TEST_EXPECT_TABLES_EQUAL(join_table->view(), table_expect);
 }
 
-class CrossJoinInvalidInputs : public cudf::test::BaseFixture {};
+class CrossJoinZeroColumnOperand : public cudf::test::BaseFixture {};
 
-TEST_F(CrossJoinInvalidInputs, EmptyTable)
+TEST_F(CrossJoinZeroColumnOperand, PreservesRowCount)
 {
   auto b_0 = column_wrapper<int32_t>{10, 20, 20, 50};
   auto b_1 = column_wrapper<float>{5.0, .7, .7, .7};
   auto b_2 = column_wrapper<int8_t>{90, 75, 62, 41};
   auto b_3 = cudf::test::strings_column_wrapper({"quick", "words", "result", ""});
 
-  auto column_a = std::vector<std::unique_ptr<cudf::column>>{};
-  auto table_a  = cudf::table(std::move(column_a));
-  auto table_b  = cudf::table_view{{b_0, b_1, b_2, b_3}};
-
-  //
-  //  table_a has no columns, table_b has columns
-  //  Let's check different permutations of passing table
-  //  with no columns to verify that exceptions are thrown
-  //
-  EXPECT_THROW(cudf::cross_join(table_a, table_b), cudf::logic_error);
-  EXPECT_THROW(cudf::cross_join(table_b, table_a), cudf::logic_error);
+  // A zero-column operand with a non-zero row count is valid: it contributes
+  // only to the output row count, not to the output columns.
+  auto table_a = cudf::table_view{std::vector<cudf::column_view>{}, 3};
+  auto table_b = cudf::table_view{{b_0, b_1, b_2, b_3}};
+
+  auto join_ab = cudf::cross_join(table_a, table_b);
+  EXPECT_EQ(join_ab->num_columns(), table_b.num_columns());
+  EXPECT_EQ(join_ab->num_rows(), table_a.num_rows() * table_b.num_rows());
+
+  auto join_ba = cudf::cross_join(table_b, table_a);
+  EXPECT_EQ(join_ba->num_columns(), table_b.num_columns());
+  EXPECT_EQ(join_ba->num_rows(), table_a.num_rows() * table_b.num_rows());
+
+  // Both operands zero-column: result is a zero-column table whose row count is
+  // the product of the operands' row counts.
+  auto table_c = cudf::table_view{std::vector<cudf::column_view>{}, 5};
+  auto join_ac = cudf::cross_join(table_a, table_c);
+  EXPECT_EQ(join_ac->num_columns(), 0);
+  EXPECT_EQ(join_ac->num_rows(), table_a.num_rows() * table_c.num_rows());
+}
+
+TEST_F(CrossJoinZeroColumnOperand, OverflowThrows)
+{
+  // A cross join whose row-count product exceeds size_type::max() throws
+  // std::overflow_error. Using zero-column operands makes the check happen
+  // before any output is allocated.
+  auto table_a =
+    cudf::table_view{std::vector<cudf::column_view>{}, std::numeric_limits<cudf::size_type>::max()};
+  auto table_b = cudf::table_view{std::vector<cudf::column_view>{}, 2};
+
+  EXPECT_THROW(cudf::cross_join(table_a, table_b), std::overflow_error);
+  EXPECT_THROW(cudf::cross_join(table_b, table_a), std::overflow_error);
 }
 
 class CrossJoinEmptyResult : public cudf::test::BaseFixture {};

python/cudf/cudf/core/join/join.py

@@ -371,10 +371,20 @@ def perform_merge(self) -> DataFrame:
 
         if self.how == "cross":
             lib_table = plc.join.cross_join(
-                plc.Table([col.plc_column for col in self.lhs._columns]),
-                plc.Table([col.plc_column for col in self.rhs._columns]),
+                plc.Table(
+                    [col.plc_column for col in self.lhs._columns],
+                    num_rows=len(self.lhs),
+                ),
+                plc.Table(
+                    [col.plc_column for col in self.rhs._columns],
+                    num_rows=len(self.rhs),
+                ),
             )
             columns = lib_table.columns()
+            # A cross join of two column-less operands is a column-less table
+            # whose row count is the product of the operands' row counts.
+            # ``columns`` cannot carry that count, so remember it explicitly.
+            cross_num_rows = lib_table.num_rows()
             num_left_cols = len(self.lhs._column_names)
             left_result = DataFrame._from_data(
                 {
@@ -444,6 +454,13 @@ def perform_merge(self) -> DataFrame:
         result = DataFrame._from_data(
             *self._merge_results(left_result, right_result)
         )
+        if self.how == "cross" and result._num_columns == 0:
+            from cudf.core.index import RangeIndex
+
+            # Both operands were column-less, so the cross product is a
+            # column-less table whose row count an empty ColumnAccessor cannot
+            # carry. Rebuild the result with the row count captured above.
+            result = DataFrame._from_data({}, index=RangeIndex(cross_num_rows))
 
         if self.sort:
             result = self._sort_result(result)

python/cudf/cudf/tests/reshape/test_merge.py

@@ -1577,6 +1577,24 @@ def test_merge_invalid_input(param):
         cudf.merge(left["a"], param)
 
 
+@pytest.mark.parametrize("left_cols", [True, False])
+@pytest.mark.parametrize("right_cols", [True, False])
+def test_cross_merge_zero_column_operand(left_cols, right_cols):
+    # A cross merge where one or both operands have no columns must still
+    # produce ``len(left) * len(right)`` rows (matching pandas), rather than
+    # dropping the row count of the column-less operand.
+    ldata = {"x": [1, 2, 3]} if left_cols else {}
+    rdata = {"y": [10, 20, 30, 40]} if right_cols else {}
+    pleft = pd.DataFrame(ldata, index=range(3))
+    pright = pd.DataFrame(rdata, index=range(4))
+    gleft = cudf.DataFrame(ldata, index=range(3))
+    gright = cudf.DataFrame(rdata, index=range(4))
+
+    expected = pleft.merge(pright, how="cross")
+    result = gleft.merge(gright, how="cross")
+    assert_eq(result, expected)
+
+
 def test_merge_natural_join_key_order_matches_left_frame():
     # A merge without ``on`` joins on the common columns in left-frame
     # column order, like pandas. Building the key list from an unordered

python/cudf_polars/tests/streaming/test_select.py

@@ -22,7 +22,6 @@
     assert_gpu_result_equal,
 )
 from cudf_polars.testing.engine_utils import warns_on_spmd
-from cudf_polars.utils.versions import POLARS_VERSION_LT_141
 
 
 @pytest.fixture
@@ -181,13 +180,6 @@ def test_select_mean_with_decimals(engine):
     assert_gpu_result_equal(q, engine=engine)
 
 
-@pytest.mark.xfail(
-    condition=not POLARS_VERSION_LT_141,
-    reason=(
-        "len() row count lost in zero-column streaming chunks "
-        "(https://github.com/rapidsai/cudf/issues/21428)"
-    ),
-)
 def test_select_with_len(streaming_engine_factory):
     engine = streaming_engine_factory(
         StreamingOptions(max_rows_per_partition=3, fallback_mode="warn"),