From f135b2b9a547a4eb694cd37334fbadcf7732f646 Mon Sep 17 00:00:00 2001
From: Alessandro Bellina <abellina@nvidia.com>
Date: Thu, 24 Feb 2022 14:37:45 +0000
Subject: [PATCH] Allow up to INT_MAX elements as a result of fused_concatenate

---
 cpp/src/copying/concatenate.cu         | 2 +-
 cpp/tests/copying/concatenate_tests.cu | 6 +++---
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/cpp/src/copying/concatenate.cu b/cpp/src/copying/concatenate.cu
index 1a240fe3ea8..82e189b5a36 100644
--- a/cpp/src/copying/concatenate.cu
+++ b/cpp/src/copying/concatenate.cu
@@ -222,7 +222,7 @@ std::unique_ptr<column> fused_concatenate(host_span<column_view const> views,
   auto const& d_offsets   = std::get<2>(device_views);
   auto const output_size  = std::get<3>(device_views);
 
-  CUDF_EXPECTS(output_size < static_cast<std::size_t>(std::numeric_limits<size_type>::max()),
+  CUDF_EXPECTS(output_size <= static_cast<std::size_t>(std::numeric_limits<size_type>::max()),
                "Total number of concatenated rows exceeds size_type range");
 
   // Allocate output
diff --git a/cpp/tests/copying/concatenate_tests.cu b/cpp/tests/copying/concatenate_tests.cu
index c5cc6a8c852..156b8741cf3 100644
--- a/cpp/tests/copying/concatenate_tests.cu
+++ b/cpp/tests/copying/concatenate_tests.cu
@@ -342,9 +342,9 @@ TEST_F(OverflowTest, OverflowTest)
   using namespace cudf;
   // should concatenate up to size_type::max - 1 rows.
   {
-    // 5 x size + size_last adds to size_type::max - 1
+    // 5 x size + size_last adds to size_type::max
     constexpr auto size      = static_cast<size_type>(static_cast<uint32_t>(250) * 1024 * 1024);
-    constexpr auto size_last = static_cast<size_type>(836763647 - 1);
+    constexpr auto size_last = static_cast<size_type>(836763647);
 
     auto many_chars      = cudf::make_fixed_width_column(data_type{type_id::INT8}, size);
     auto many_chars_last = cudf::make_fixed_width_column(data_type{type_id::INT8}, size_last);
@@ -354,7 +354,7 @@ TEST_F(OverflowTest, OverflowTest)
     std::vector<cudf::table_view> table_views_to_concat({tbl, tbl, tbl, tbl, tbl, tbl_last});
     std::unique_ptr<cudf::table> concatenated_tables = cudf::concatenate(table_views_to_concat);
     EXPECT_NO_THROW(rmm::cuda_stream_default.synchronize());
-    ASSERT_EQ(concatenated_tables->num_rows(), std::numeric_limits<size_type>::max() - 1);
+    ASSERT_EQ(concatenated_tables->num_rows(), std::numeric_limits<size_type>::max());
   }
 
   // primitive column