List lexicographic comparator

cpp/benchmarks/CMakeLists.txt

@@ -165,7 +165,7 @@ ConfigureNVBench(SEARCH_NVBENCH search/contains.cpp)
 # ##################################################################################################
 # * sort benchmark --------------------------------------------------------------------------------
 ConfigureBench(SORT_BENCH sort/rank.cpp sort/sort.cpp sort/sort_strings.cpp)
-ConfigureNVBench(SORT_NVBENCH sort/sort_structs.cpp)
+ConfigureNVBench(SORT_NVBENCH sort/sort_lists.cpp sort/sort_structs.cpp)
 
 # ##################################################################################################
 # * quantiles benchmark

cpp/benchmarks/sort/sort_lists.cpp

@@ -0,0 +1,49 @@
+/*
+ * Copyright (c) 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.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <benchmarks/common/generate_input.hpp>
+#include <benchmarks/fixture/rmm_pool_raii.hpp>
+
+#include <cudf/detail/sorting.hpp>
+
+#include <nvbench/nvbench.cuh>
+
+void nvbench_sort_lists(nvbench::state& state)
+{
+  cudf::rmm_pool_raii pool_raii;
+
+  const size_t size_bytes(state.get_int64("size_bytes"));
+  const cudf::size_type depth{static_cast<cudf::size_type>(state.get_int64("depth"))};
+  auto const null_frequency{state.get_float64("null_frequency")};
+
+  data_profile table_profile;
+  table_profile.set_distribution_params(cudf::type_id::LIST, distribution_id::UNIFORM, 0, 5);
+  table_profile.set_list_depth(depth);
+  table_profile.set_null_probability(null_frequency);
+  auto const table =
+    create_random_table({cudf::type_id::LIST}, table_size_bytes{size_bytes}, table_profile);
+
+  state.exec(nvbench::exec_tag::sync, [&](nvbench::launch& launch) {
+    rmm::cuda_stream_view stream_view{launch.get_stream()};
+    cudf::detail::sorted_order(*table, {}, {}, stream_view, rmm::mr::get_current_device_resource());
+  });
+}
+
+NVBENCH_BENCH(nvbench_sort_lists)
+  .set_name("sort_list")
+  .add_int64_power_of_two_axis("size_bytes", {10, 18, 24, 28})
+  .add_int64_axis("depth", {1, 4})
+  .add_float64_axis("null_frequency", {0, 0.2});

cpp/include/cudf/lists/detail/dremel.hpp

@@ -28,15 +28,11 @@ namespace cudf::detail {
  * @see the `dremel_data` struct for more info.
  */
 struct dremel_device_view {
-  // TODO: These elements are default initializable to support default
-  // initialization of the object. This is currently exploited to create views
-  // that will never actually be used. We should consider whether this
-  // represents a serious issue that should be worked around more robustly.
-  size_type const* offsets{};
-  uint8_t const* rep_levels{};
-  uint8_t const* def_levels{};
-  size_type const leaf_data_size{};
-  uint8_t const max_def_level{};
+  size_type const* offsets;
+  uint8_t const* rep_levels;
+  uint8_t const* def_levels;
+  size_type const leaf_data_size;
+  uint8_t const max_def_level;
 };
 
 /**

cpp/include/cudf/table/table_view.hpp

@@ -166,6 +166,14 @@ class table_view_base {
    */
   table_view_base& operator=(table_view_base&&) = default;
 };
+
+/**
+ * @brief The function to collect all nested columns in a given table.
+ *
+ * @param table The input table
+ * @return A vector containing all nested columns in the input table
+ */
+bool has_nested_columns(table_view const& table);
 }  // namespace detail
 
 /**

cpp/src/binaryop/compiled/struct_binary_ops.cuh

@@ -93,9 +93,17 @@ void apply_struct_binary_op(mutable_column_view& out,
       out.end<bool>(),
       device_comparison_functor{optional_iter, is_lhs_scalar, is_rhs_scalar, device_comparator});
   };
-  is_any_v<BinaryOperator, ops::LessEqual, ops::GreaterEqual>
-    ? tabulate_device_operator(table_comparator.less_equivalent(comparator_nulls, comparator))
-    : tabulate_device_operator(table_comparator.less(comparator_nulls, comparator));
+  if (cudf::detail::has_nested_columns(tlhs) || cudf::detail::has_nested_columns(trhs)) {
+    is_any_v<BinaryOperator, ops::LessEqual, ops::GreaterEqual>
+      ? tabulate_device_operator(
+          table_comparator.less_equivalent<true>(comparator_nulls, comparator))
+      : tabulate_device_operator(table_comparator.less<true>(comparator_nulls, comparator));
+  } else {
+    is_any_v<BinaryOperator, ops::LessEqual, ops::GreaterEqual>
+      ? tabulate_device_operator(
+          table_comparator.less_equivalent<false>(comparator_nulls, comparator))
+      : tabulate_device_operator(table_comparator.less<false>(comparator_nulls, comparator));
+  }
 }
 
 template <typename PhysicalEqualityComparator =

cpp/src/search/search_ordered.cu

@@ -68,28 +68,49 @@ std::unique_ptr<column> search_ordered(table_view const& haystack,
 
   auto const comparator = cudf::experimental::row::lexicographic::two_table_comparator(
     matched_haystack, matched_needles, column_order, null_precedence, stream);
-  auto const has_nulls    = has_nested_nulls(matched_haystack) or has_nested_nulls(matched_needles);
-  auto const d_comparator = comparator.less(nullate::DYNAMIC{has_nulls});
+  auto const has_nulls = has_nested_nulls(matched_haystack) or has_nested_nulls(matched_needles);
 
   auto const haystack_it = cudf::experimental::row::lhs_iterator(0);
   auto const needles_it  = cudf::experimental::row::rhs_iterator(0);
 
-  if (find_first) {
-    thrust::lower_bound(rmm::exec_policy(stream),
-                        haystack_it,
-                        haystack_it + haystack.num_rows(),
-                        needles_it,
-                        needles_it + needles.num_rows(),
-                        out_it,
-                        d_comparator);
+  if (cudf::detail::has_nested_columns(haystack) || cudf::get_nested_columns(needles).size() > 0) {
+    auto const d_comparator = comparator.less<true>(nullate::DYNAMIC{has_nulls});
+    if (find_first) {
+      thrust::lower_bound(rmm::exec_policy(stream),
+                          haystack_it,
+                          haystack_it + haystack.num_rows(),
+                          needles_it,
+                          needles_it + needles.num_rows(),
+                          out_it,
+                          d_comparator);
+    } else {
+      thrust::upper_bound(rmm::exec_policy(stream),
+                          haystack_it,
+                          haystack_it + haystack.num_rows(),
+                          needles_it,
+                          needles_it + needles.num_rows(),
+                          out_it,
+                          d_comparator);
+    }
   } else {
-    thrust::upper_bound(rmm::exec_policy(stream),
-                        haystack_it,
-                        haystack_it + haystack.num_rows(),
-                        needles_it,
-                        needles_it + needles.num_rows(),
-                        out_it,
-                        d_comparator);
+    auto const d_comparator = comparator.less<false>(nullate::DYNAMIC{has_nulls});
+    if (find_first) {
+      thrust::lower_bound(rmm::exec_policy(stream),
+                          haystack_it,
+                          haystack_it + haystack.num_rows(),
+                          needles_it,
+                          needles_it + needles.num_rows(),
+                          out_it,
+                          d_comparator);
+    } else {
+      thrust::upper_bound(rmm::exec_policy(stream),
+                          haystack_it,
+                          haystack_it + haystack.num_rows(),
+                          needles_it,
+                          needles_it + needles.num_rows(),
+                          out_it,
+                          d_comparator);
+    }
   }
   return result;
 }

cpp/src/sort/sort_impl.cuh

@@ -127,18 +127,32 @@ std::unique_ptr<column> sorted_order(table_view input,
 
   auto comp =
     experimental::row::lexicographic::self_comparator(input, column_order, null_precedence, stream);
-  auto comparator = comp.less(nullate::DYNAMIC{has_nested_nulls(input)});
-
-  if (stable) {
-    thrust::stable_sort(rmm::exec_policy(stream),
-                        mutable_indices_view.begin<size_type>(),
-                        mutable_indices_view.end<size_type>(),
-                        comparator);
+  if (cudf::detail::has_nested_columns(input)) {
+    auto comparator = comp.less<true>(nullate::DYNAMIC{has_nested_nulls(input)});
+    if (stable) {
+      thrust::stable_sort(rmm::exec_policy(stream),
+                          mutable_indices_view.begin<size_type>(),
+                          mutable_indices_view.end<size_type>(),
+                          comparator);
+    } else {
+      thrust::sort(rmm::exec_policy(stream),
+                   mutable_indices_view.begin<size_type>(),
+                   mutable_indices_view.end<size_type>(),
+                   comparator);
+    }
   } else {
-    thrust::sort(rmm::exec_policy(stream),
-                 mutable_indices_view.begin<size_type>(),
-                 mutable_indices_view.end<size_type>(),
-                 comparator);
+    auto comparator = comp.less<false>(nullate::DYNAMIC{has_nested_nulls(input)});
+    if (stable) {
+      thrust::stable_sort(rmm::exec_policy(stream),
+                          mutable_indices_view.begin<size_type>(),
+                          mutable_indices_view.end<size_type>(),
+                          comparator);
+    } else {
+      thrust::sort(rmm::exec_policy(stream),
+                   mutable_indices_view.begin<size_type>(),
+                   mutable_indices_view.end<size_type>(),
+                   comparator);
+    }
   }
   // protection for temporary d_column_order and d_null_precedence
   stream.synchronize();

cpp/src/table/row_operators.cu

@@ -255,6 +255,24 @@ auto decompose_structs(table_view table,
                          std::move(verticalized_col_depths));
 }
 
+/*
+ * This helper function generates dremel data for any list-type columns in a
+ * table. This data is necessary for lexicographic comparisons.
+ */
+auto list_lex_preprocess(table_view table, rmm::cuda_stream_view stream)
+{
+  std::vector<detail::dremel_data> dremel_data;
+  std::vector<detail::dremel_device_view> dremel_device_views;
+  for (auto const& col : table) {
+    if (col.type().id() == type_id::LIST) {
+      dremel_data.push_back(detail::get_dremel_data(col, {}, false, stream));
+      dremel_device_views.push_back(dremel_data.back());
+    }
+  }
+  auto d_dremel_device_views = detail::make_device_uvector_sync(dremel_device_views, stream);
+  return std::make_tuple(std::move(dremel_data), std::move(d_dremel_device_views));
+}
+
 using column_checker_fn_t = std::function<void(column_view const&)>;
 
 /**
@@ -264,18 +282,25 @@ using column_checker_fn_t = std::function<void(column_view const&)>;
  */
 void check_lex_compatibility(table_view const& input)
 {
-  // Basically check if there's any LIST hiding anywhere in the table
+  // Basically check if there's any LIST of STRUCT or STRUCT of LIST hiding anywhere in the table
   column_checker_fn_t check_column = [&](column_view const& c) {
-    CUDF_EXPECTS(c.type().id() != type_id::LIST,
-                 "Cannot lexicographic compare a table with a LIST column");
+    if (c.type().id() == type_id::LIST) {
+      auto const& list_col = lists_column_view(c);
+      CUDF_EXPECTS(list_col.child().type().id() != type_id::STRUCT,
+                   "Cannot lexicographic compare a table with a LIST of STRUCT column");
+      check_column(list_col.child());
+    } else if (c.type().id() == type_id::STRUCT) {
+      for (auto child = c.child_begin(); child < c.child_end(); ++child) {
+        CUDF_EXPECTS(child->type().id() != type_id::LIST,
+                     "Cannot lexicographic compare a table with a STRUCT of LIST column");
+        check_column(*child);
+      }
+    }
     if (not is_nested(c.type())) {
       CUDF_EXPECTS(is_relationally_comparable(c.type()),
                    "Cannot lexicographic compare a table with a column of type " +
                      jit::get_type_name(c.type()));
     }
-    for (auto child = c.child_begin(); child < c.child_end(); ++child) {
-      check_column(*child);
-    }
   };
   for (column_view const& c : input) {
     check_column(c);
@@ -336,8 +361,23 @@ std::shared_ptr<preprocessed_table> preprocessed_table::create(
   auto d_null_precedence = detail::make_device_uvector_async(new_null_precedence, stream);
   auto d_depths          = detail::make_device_uvector_async(verticalized_col_depths, stream);
 
-  return std::shared_ptr<preprocessed_table>(new preprocessed_table(
-    std::move(d_t), std::move(d_column_order), std::move(d_null_precedence), std::move(d_depths)));
+  if (detail::has_nested_columns(t)) > 0)
+    {
+      auto [dremel_data, d_dremel_device_view] = list_lex_preprocess(verticalized_lhs, stream);
+      return std::shared_ptr<preprocessed_table>(
+        new preprocessed_table(std::move(d_t),
+                               std::move(d_column_order),
+                               std::move(d_null_precedence),
+                               std::move(d_depths),
+                               std::move(dremel_data),
+                               std::move(d_dremel_device_view)));
+    }
+  else {
+    return std::shared_ptr<preprocessed_table>(new preprocessed_table(std::move(d_t),
+                                                                      std::move(d_column_order),
+                                                                      std::move(d_null_precedence),
+                                                                      std::move(d_depths)));
+  }
 }
 
 two_table_comparator::two_table_comparator(table_view const& left,

cpp/src/table/table_view.cpp

@@ -134,5 +134,11 @@ template bool is_relationally_comparable<table_view>(table_view const& lhs, tabl
 template bool is_relationally_comparable<mutable_table_view>(mutable_table_view const& lhs,
                                                              mutable_table_view const& rhs);
 
+bool has_nested_columns(table_view const& table)
+{
+  return std::any_of(
+    table.begin(), table.end(), [](column_view const& col) { return is_nested(col.type()); });
+}
+
 }  // namespace detail
 }  // namespace cudf

cpp/tests/rolling/collect_ops_test.cpp

@@ -2275,10 +2275,23 @@ TEST_F(CollectSetTest, ListTypeRollingWindow)
   auto const prev_column = fixed_width_column_wrapper<size_type>{1, 2, 2, 2, 2};
   auto const foll_column = fixed_width_column_wrapper<size_type>{1, 1, 1, 1, 0};
 
-  EXPECT_THROW(rolling_collect_set(input_column,
-                                   prev_column,
-                                   foll_column,
-                                   1,
-                                   *make_collect_set_aggregation<rolling_aggregation>()),
-               cudf::logic_error);
+  auto const expected = [] {
+    auto data = fixed_width_column_wrapper<int32_t>{1, 2, 3, 4, 5, 1, 2, 3, 4,  5, 6, 4, 5,
+                                                    6, 7, 8, 9, 6, 7, 8, 9, 10, 7, 8, 9, 10};
+    auto inner_offsets =
+      fixed_width_column_wrapper<int32_t>{0, 3, 5, 8, 10, 11, 13, 14, 17, 18, 21, 22, 25, 26};
+    auto outer_offsets = fixed_width_column_wrapper<size_type>{0, 2, 5, 8, 11, 13};
+
+    auto inner_list = cudf::make_lists_column(13, inner_offsets.release(), data.release(), 0, {});
+
+    return cudf::make_lists_column(5, outer_offsets.release(), std::move(inner_list), 0, {});
+  }();
+
+  auto const result = rolling_collect_set(input_column,
+                                          prev_column,
+                                          foll_column,
+                                          1,
+                                          *make_collect_set_aggregation<rolling_aggregation>());
+
+  CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(expected->view(), result->view());
 }