Fix cudf::shift to handle offset greater than column size

cpp/include/cudf/copying.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018-2021, NVIDIA CORPORATION.
+ * Copyright (c) 2018-2022, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -371,7 +371,7 @@ std::unique_ptr<column> copy_range(
  * @param fill_value Fill value for indeterminable outputs.
  * @param mr         Device memory resource used to allocate the returned result's device memory
  *
- * @throw cudf::logic_error if @p input dtype is not fixed-with.
+ * @throw cudf::logic_error if @p input dtype is neither fixed-width nor string type
  * @throw cudf::logic_error if @p fill_value dtype does not match @p input dtype.
  */
 std::unique_ptr<column> shift(

cpp/src/copying/shift.cu

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2021, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2022, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -70,7 +70,7 @@ struct shift_functor {
                    std::unique_ptr<column>>
   operator()(Args&&...)
   {
-    CUDF_FAIL("shift does not support non-fixed-width types.");
+    CUDF_FAIL("shift only supports fixed-width or string types.");
   }
 
   template <typename T, typename... Args>
@@ -125,6 +125,7 @@ struct shift_functor {
 
     // avoid assigning elements we know to be invalid.
     if (not scalar_is_valid) {
+      if (std::abs(offset) > size) { return output; }
       if (offset > 0) {
         index_begin = thrust::make_counting_iterator<size_type>(offset);
         data        = data + offset;

cpp/src/strings/copying/shift.cu

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021, NVIDIA CORPORATION.
+ * Copyright (c) 2021-2022, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -94,6 +94,9 @@ std::unique_ptr<column> shift(strings_column_view const& input,
 {
   auto d_fill_str = static_cast<string_scalar const&>(fill_value).value(stream);
 
+  // adjust offset when greater than the size of the input
+  if (std::abs(offset) > input.size()) { offset = input.size(); }
+
   // output offsets column is the same size as the input
   auto const input_offsets =
     cudf::detail::slice(

cpp/tests/copying/shift_tests.cpp

@@ -20,16 +20,14 @@
 #include <cudf_test/cudf_gtest.hpp>
 #include <cudf_test/type_lists.hpp>
 
-#include <cudf/column/column.hpp>
 #include <cudf/copying.hpp>
 #include <cudf/scalar/scalar.hpp>
+#include <cudf/utilities/traits.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
 
-#include <functional>
 #include <limits>
 #include <memory>
-#include <type_traits>
 
 using cudf::test::fixed_width_column_wrapper;
 using TestTypes = cudf::test::Types<int32_t>;
@@ -72,28 +70,12 @@ constexpr auto lowest()
 }
 
 template <typename T>
-struct ShiftTest : public cudf::test::BaseFixture {
+struct ShiftTestsTyped : public cudf::test::BaseFixture {
 };
 
-TYPED_TEST_SUITE(ShiftTest, cudf::test::FixedWidthTypes);
+TYPED_TEST_SUITE(ShiftTestsTyped, cudf::test::FixedWidthTypes);
 
-TYPED_TEST(ShiftTest, OneColumnEmpty)
-{
-  using T = TypeParam;
-
-  std::vector<T> vals{};
-  std::vector<bool> mask{};
-
-  auto input    = fixed_width_column_wrapper<T>{};
-  auto expected = fixed_width_column_wrapper<T>(vals.begin(), vals.end(), mask.begin());
-
-  auto fill   = make_scalar<T>();
-  auto actual = cudf::shift(input, 5, *fill);
-
-  CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, *actual);
-}
-
-TYPED_TEST(ShiftTest, TwoColumnsEmpty)
+TYPED_TEST(ShiftTestsTyped, ColumnEmpty)
 {
   using T = TypeParam;
 
@@ -109,7 +91,7 @@ TYPED_TEST(ShiftTest, TwoColumnsEmpty)
   CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, *actual);
 }
 
-TYPED_TEST(ShiftTest, OneColumn)
+TYPED_TEST(ShiftTestsTyped, NonNullColumn)
 {
   using T = TypeParam;
 
@@ -134,7 +116,7 @@ TYPED_TEST(ShiftTest, OneColumn)
   CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, *actual);
 }
 
-TYPED_TEST(ShiftTest, OneColumnNegativeShift)
+TYPED_TEST(ShiftTestsTyped, NegativeShift)
 {
   using T = TypeParam;
 
@@ -159,7 +141,7 @@ TYPED_TEST(ShiftTest, OneColumnNegativeShift)
   CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, *actual);
 }
 
-TYPED_TEST(ShiftTest, OneColumnNullFill)
+TYPED_TEST(ShiftTestsTyped, NullScalar)
 {
   using T = TypeParam;
 
@@ -186,7 +168,7 @@ TYPED_TEST(ShiftTest, OneColumnNullFill)
   CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, *actual);
 }
 
-TYPED_TEST(ShiftTest, TwoColumnsNullableInput)
+TYPED_TEST(ShiftTestsTyped, NullableColumn)
 {
   using T = TypeParam;
 
@@ -199,25 +181,21 @@ TYPED_TEST(ShiftTest, TwoColumnsNullableInput)
   CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, *actual);
 }
 
-TYPED_TEST(ShiftTest, MismatchFillValueDtypes)
+TYPED_TEST(ShiftTestsTyped, MismatchFillValueDtypes)
 {
   using T = TypeParam;
 
-  if (std::is_same_v<T, int>) { return; }
-
   auto input = fixed_width_column_wrapper<T>{};
 
-  auto fill = make_scalar<int>();
+  auto fill = cudf::string_scalar("");
 
-  std::unique_ptr<cudf::column> output;
-
-  EXPECT_THROW(output = cudf::shift(input, 5, *fill), cudf::logic_error);
+  EXPECT_THROW(cudf::shift(input, 5, fill), cudf::logic_error);
 }
 
-struct ShiftTestNonFixedWidth : public cudf::test::BaseFixture {
+struct ShiftTests : public cudf::test::BaseFixture {
 };
 
-TEST_F(ShiftTestNonFixedWidth, StringsShiftTest)
+TEST_F(ShiftTests, StringsShiftTest)
 {
   auto input =
     cudf::test::strings_column_wrapper({"", "bb", "ccc", "ddddddé", ""}, {0, 1, 1, 1, 0});
@@ -243,3 +221,33 @@ TEST_F(ShiftTestNonFixedWidth, StringsShiftTest)
   auto sliced_left = cudf::test::strings_column_wrapper({"ccc", "ddddddé", "xx"});
   CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(sliced_left, *results);
 }
+
+TEST_F(ShiftTests, OffsetGreaterThanSize)
+{
+  auto const input_str =
+    cudf::test::strings_column_wrapper({"", "bb", "ccc", "ddé", ""}, {0, 1, 1, 1, 0});
+  auto results      = cudf::shift(input_str, 6, cudf::string_scalar("xx"));
+  auto expected_str = cudf::test::strings_column_wrapper({"xx", "xx", "xx", "xx", "xx"});
+  CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(expected_str, *results);
+  results = cudf::shift(input_str, -6, cudf::string_scalar("xx"));
+  CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(expected_str, *results);
+
+  results      = cudf::shift(input_str, 6, cudf::string_scalar("", false));
+  expected_str = cudf::test::strings_column_wrapper({"", "", "", "", ""}, {0, 0, 0, 0, 0});
+  CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(expected_str, *results);
+  results = cudf::shift(input_str, -6, cudf::string_scalar("", false));
+  CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(expected_str, *results);
+
+  auto const input = fixed_width_column_wrapper<int32_t>({0, 2, 3, 4, 0}, {0, 1, 1, 1, 0});
+  results          = cudf::shift(input, 6, cudf::numeric_scalar<int32_t>(9));
+  auto expected    = fixed_width_column_wrapper<int32_t>({9, 9, 9, 9, 9});
+  CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(expected, *results);
+  results = cudf::shift(input, -6, cudf::numeric_scalar<int32_t>(9));
+  CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(expected, *results);
+
+  results  = cudf::shift(input, 6, cudf::numeric_scalar<int32_t>(0, false));
+  expected = fixed_width_column_wrapper<int32_t>({0, 0, 0, 0, 0}, {0, 0, 0, 0, 0});
+  CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(expected, *results);
+  results = cudf::shift(input, -6, cudf::numeric_scalar<int32_t>(0, false));
+  CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(expected, *results);
+}

python/cudf/cudf/core/column/column.py

@@ -349,14 +349,6 @@ def _fill(
         return self
 
     def shift(self, offset: int, fill_value: ScalarLike) -> ColumnBase:
-        # libcudf currently doesn't handle case when offset > len(df)
-        # ticket to fix the bug in link below:
-        # https://github.com/rapidsai/cudf/issues/10314
-        if abs(offset) > len(self):
-            if fill_value is None:
-                return column_empty_like(self, masked=True)
-            else:
-                return full(len(self), fill_value, dtype=self.dtype)
         return libcudf.copying.shift(self, offset, fill_value)
 
     @property