Expose test APIs for converting vectors to variant's (#6808)

Summary: Part of #6595 Pull Request resolved: #6808 Reviewed By: xiaoxmeng Differential Revision: D49758968 Pulled By: mbasmanova fbshipit-source-id: 889f4d0d8c24ee370d0cc5a0b5a65046b06a965d
facebookincubator · Sep 29, 2023 · 448370f · 448370f
1 parent b1db991
commit 448370f
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 28 deletions.
diff --git a/velox/exec/tests/utils/QueryAssertions.cpp b/velox/exec/tests/utils/QueryAssertions.cpp
@@ -464,26 +464,6 @@ variant variantAt(const VectorPtr& vector, vector_size_t row) {
   return VELOX_DYNAMIC_SCALAR_TYPE_DISPATCH(variantAt, typeKind, vector, row);
 }
 
-std::vector<MaterializedRow> materialize(const RowVectorPtr& vector) {
-  auto size = vector->size();
-  std::vector<MaterializedRow> rows;
-  rows.reserve(size);
-
-  auto rowType = vector->type()->as<TypeKind::ROW>();
-
-  for (size_t i = 0; i < size; ++i) {
-    auto numColumns = rowType.size();
-    MaterializedRow row;
-    row.reserve(numColumns);
-    for (size_t j = 0; j < numColumns; ++j) {
-      row.push_back(variantAt(vector->childAt(j), i));
-    }
-    rows.push_back(row);
-  }
-
-  return rows;
-}
-
 MaterializedRow getColumns(
     const MaterializedRow& row,
     const std::vector<uint32_t>& columnIndices) {
@@ -808,6 +788,26 @@ void verifyDuckDBResult(const DuckDBQueryResult& result, std::string_view sql) {
 
 } // namespace
 
+std::vector<MaterializedRow> materialize(const RowVectorPtr& vector) {
+  auto size = vector->size();
+  std::vector<MaterializedRow> rows;
+  rows.reserve(size);
+
+  auto rowType = vector->type()->as<TypeKind::ROW>();
+
+  for (size_t i = 0; i < size; ++i) {
+    auto numColumns = rowType.size();
+    MaterializedRow row;
+    row.reserve(numColumns);
+    for (size_t j = 0; j < numColumns; ++j) {
+      row.push_back(variantAt(vector->childAt(j), i));
+    }
+    rows.push_back(row);
+  }
+
+  return rows;
+}
+
 void DuckDbQueryRunner::createTable(
     const std::string& name,
     const std::vector<RowVectorPtr>& data) {
@@ -1053,16 +1053,23 @@ bool assertEqualResults(
   return true;
 }
 
-bool assertEqualResults(
-    const MaterializedRowMultiset& expectedRows,
-    const std::vector<RowVectorPtr>& actual) {
-  MaterializedRowMultiset actualRows;
-  for (auto vector : actual) {
+MaterializedRowMultiset materialize(const std::vector<RowVectorPtr>& vectors) {
+  MaterializedRowMultiset materialized;
+  for (auto vector : vectors) {
     auto rows = materialize(vector);
     std::copy(
-        rows.begin(), rows.end(), std::inserter(actualRows, actualRows.end()));
+        rows.begin(),
+        rows.end(),
+        std::inserter(materialized, materialized.end()));
   }
-  return assertEqualResults(expectedRows, actualRows, "Unexpected results");
+  return materialized;
+}
+
+bool assertEqualResults(
+    const MaterializedRowMultiset& expectedRows,
+    const std::vector<RowVectorPtr>& actual) {
+  return assertEqualResults(
+      expectedRows, materialize(actual), "Unexpected results");
 }
 
 void assertResults(

diff --git a/velox/exec/tests/utils/QueryAssertions.h b/velox/exec/tests/utils/QueryAssertions.h
@@ -26,9 +26,15 @@ namespace facebook::velox::exec::test {
 using MaterializedRow = std::vector<velox::variant>;
 using DuckDBQueryResult = std::unique_ptr<::duckdb::MaterializedQueryResult>;
 
-// Multiset that compares floating-point values directly.
+/// Multiset that compares floating-point values directly.
 using MaterializedRowMultiset = std::multiset<MaterializedRow>;
 
+/// Converts input 'RowVector' into a list of 'MaterializedRow's.
+std::vector<MaterializedRow> materialize(const RowVectorPtr& vector);
+
+/// Converts a list of 'RowVector's into 'MaterializedRowMultiset'.
+MaterializedRowMultiset materialize(const std::vector<RowVectorPtr>& vectors);
+
 class DuckDbQueryRunner {
  public:
   DuckDbQueryRunner() : db_(nullptr) {}