Add TopNRowNumberFuzzer

facebookincubator · Jan 28, 2025 · ee07a9d · ee07a9d
1 parent a837ebd
commit ee07a9d
Show file tree

Hide file tree

Showing 10 changed files with 930 additions and 0 deletions.
diff --git a/velox/exec/fuzzer/CMakeLists.txt b/velox/exec/fuzzer/CMakeLists.txt
@@ -95,6 +95,16 @@ target_link_libraries(
   velox_exec_test_lib
   velox_expression_test_utility)
 
+add_library(velox_topn_row_number_fuzzer TopNRowNumberFuzzer.cpp)
+
+target_link_libraries(
+  velox_topn_row_number_fuzzer
+  velox_fuzzer_util
+  velox_type
+  velox_vector_fuzzer
+  velox_exec_test_lib
+  velox_expression_test_utility)
+
 add_library(velox_join_fuzzer JoinFuzzer.cpp)
 
 target_link_libraries(

diff --git a/velox/exec/fuzzer/DuckQueryRunner.cpp b/velox/exec/fuzzer/DuckQueryRunner.cpp
@@ -166,6 +166,11 @@ std::optional<std::string> DuckQueryRunner::toSql(
     return toSql(rowNumberNode);
   }
 
+  if (const auto rowNumberNode =
+          std::dynamic_pointer_cast<const core::TopNRowNumberNode>(plan)) {
+    return toSql(rowNumberNode);
+  }
+
   if (const auto joinNode =
           std::dynamic_pointer_cast<const core::HashJoinNode>(plan)) {
     return toSql(joinNode);
@@ -377,4 +382,45 @@ std::optional<std::string> DuckQueryRunner::toSql(
 
   return sql.str();
 }
+
+std::optional<std::string> DuckQueryRunner::toSql(
+    const std::shared_ptr<const core::TopNRowNumberNode>& rowNumberNode) {
+  std::stringstream sql;
+  sql << "SELECT * FROM (SELECT ";
+
+  const auto& inputType = rowNumberNode->sources()[0]->outputType();
+  for (auto i = 0; i < inputType->size(); ++i) {
+    appendComma(i, sql);
+    sql << inputType->nameOf(i);
+  }
+
+  sql << ", row_number() OVER (";
+
+  const auto& partitionKeys = rowNumberNode->partitionKeys();
+  if (!partitionKeys.empty()) {
+    sql << "partition by ";
+    for (auto i = 0; i < partitionKeys.size(); ++i) {
+      appendComma(i, sql);
+      sql << partitionKeys[i]->name();
+    }
+  }
+
+  const auto& sortingKeys = rowNumberNode->sortingKeys();
+  const auto& sortingOrders = rowNumberNode->sortingOrders();
+
+  if (!sortingKeys.empty()) {
+    sql << " ORDER BY ";
+    for (auto j = 0; j < sortingKeys.size(); ++j) {
+      appendComma(j, sql);
+      sql << sortingKeys[j]->name() << " " << sortingOrders[j].toString();
+    }
+  }
+
+  sql << ") as row_number FROM tmp) ";
+
+  sql << " where row_number <= " << rowNumberNode->limit();
+
+  return sql.str();
+}
+
 } // namespace facebook::velox::exec::test
diff --git a/velox/exec/fuzzer/DuckQueryRunner.h b/velox/exec/fuzzer/DuckQueryRunner.h
@@ -74,6 +74,9 @@ class DuckQueryRunner : public ReferenceQueryRunner {
   std::optional<std::string> toSql(
       const std::shared_ptr<const core::RowNumberNode>& rowNumberNode);
 
+  std::optional<std::string> toSql(
+      const std::shared_ptr<const core::TopNRowNumberNode>& rowNumberNode);
+
   std::unordered_set<std::string> aggregateFunctionNames_;
 };
 

diff --git a/velox/exec/fuzzer/PrestoQueryRunner.cpp b/velox/exec/fuzzer/PrestoQueryRunner.cpp
@@ -183,6 +183,11 @@ std::optional<std::string> PrestoQueryRunner::toSql(
     return toSql(rowNumberNode);
   }
 
+  if (const auto rowNumberNode =
+          std::dynamic_pointer_cast<const core::TopNRowNumberNode>(plan)) {
+    return toSql(rowNumberNode);
+  }
+
   if (auto tableWriteNode =
           std::dynamic_pointer_cast<const core::TableWriteNode>(plan)) {
     return toSql(tableWriteNode);
@@ -498,6 +503,51 @@ std::optional<std::string> PrestoQueryRunner::toSql(
   return sql.str();
 }
 
+std::optional<std::string> PrestoQueryRunner::toSql(
+    const std::shared_ptr<const core::TopNRowNumberNode>& rowNumberNode) {
+  if (!isSupportedDwrfType(rowNumberNode->sources()[0]->outputType())) {
+    return std::nullopt;
+  }
+
+  std::stringstream sql;
+  sql << "SELECT * FROM (SELECT ";
+
+  const auto& inputType = rowNumberNode->sources()[0]->outputType();
+  for (auto i = 0; i < inputType->size(); ++i) {
+    appendComma(i, sql);
+    sql << inputType->nameOf(i);
+  }
+
+  sql << ", row_number() OVER (";
+
+  const auto& partitionKeys = rowNumberNode->partitionKeys();
+  if (!partitionKeys.empty()) {
+    sql << "partition by ";
+    for (auto i = 0; i < partitionKeys.size(); ++i) {
+      appendComma(i, sql);
+      sql << partitionKeys[i]->name();
+    }
+  }
+
+  const auto& sortingKeys = rowNumberNode->sortingKeys();
+  const auto& sortingOrders = rowNumberNode->sortingOrders();
+
+  if (!sortingKeys.empty()) {
+    sql << " ORDER BY ";
+    for (auto j = 0; j < sortingKeys.size(); ++j) {
+      appendComma(j, sql);
+      sql << sortingKeys[j]->name() << " " << sortingOrders[j].toString();
+    }
+  }
+
+  sql << ") as row_number FROM tmp)";
+
+  sql << " where row_number <= " << rowNumberNode->limit();
+
+  LOG(INFO) << "PrestoSQL generated " << sql.str();
+  return sql.str();
+}
+
 std::optional<std::string> PrestoQueryRunner::toSql(
     const std::shared_ptr<const core::TableWriteNode>& tableWriteNode) {
   auto insertTableHandle =

diff --git a/velox/exec/fuzzer/PrestoQueryRunner.h b/velox/exec/fuzzer/PrestoQueryRunner.h
@@ -120,6 +120,9 @@ class PrestoQueryRunner : public velox::exec::test::ReferenceQueryRunner {
   std::optional<std::string> toSql(
       const std::shared_ptr<const velox::core::RowNumberNode>& rowNumberNode);
 
+  std::optional<std::string> toSql(
+      const std::shared_ptr<const core::TopNRowNumberNode>& rowNumberNode);
+
   std::optional<std::string> toSql(
       const std::shared_ptr<const core::TableWriteNode>& tableWriteNode);