fix replace nulls tests

karthikeyann · Dec 12, 2023 · 6fb609a · 6fb609a
1 parent b841fed
commit 6fb609a
Showing 1 changed file with 10 additions and 4 deletions.
diff --git a/cpp/tests/groupby/replace_nulls_tests.cpp b/cpp/tests/groupby/replace_nulls_tests.cpp
@@ -22,6 +22,7 @@
 
 #include <cudf/groupby.hpp>
 #include <cudf/replace.hpp>
+#include <cudf/sorting.hpp>
 #include <cudf/table/table.hpp>
 #include <cudf/table/table_view.hpp>
 #include <cudf/types.hpp>
@@ -44,10 +45,15 @@ void TestReplaceNullsGroupbySingle(K const& key,
 {
   cudf::groupby::groupby gb_obj(cudf::table_view({key}));
   std::vector<cudf::replace_policy> policies{policy};
-  auto p = gb_obj.replace_nulls(cudf::table_view({input}), policies);
-
-  CUDF_TEST_EXPECT_TABLES_EQUAL(*p.first, cudf::table_view({expected_key}));
-  CUDF_TEST_EXPECT_TABLES_EQUAL(*p.second, cudf::table_view({expected_val}));
+  auto result = gb_obj.replace_nulls(cudf::table_view({input}), policies);
+
+  cudf::test::print(result.first->get_column(0));
+  cudf::test::print(expected_key);
+  auto const sort_order  = cudf::sorted_order(result.first->view());
+  auto const sorted_keys = cudf::gather(result.first->view(), *sort_order);
+  auto const sorted_vals = cudf::gather(result.second->view(), *sort_order);
+  CUDF_TEST_EXPECT_TABLES_EQUAL(*sorted_keys, cudf::table_view({expected_key}));
+  CUDF_TEST_EXPECT_TABLES_EQUAL(*sorted_vals, cudf::table_view({expected_val}));
 }
 
 TYPED_TEST(GroupbyReplaceNullsFixedWidthTest, PrecedingFill)