feat(fuzzer): Add TopNRowNumberFuzzer (#12103)

Summary:
topNRowNumber node is an optimized planNode for SQL with ranking window functions but which limits them to only the topN results. Add a TopNRowNumberFuzzer for plans with this planNode.

This fuzzer is closely modeled after the RowNumberFuzzer. So the common code is abstracted to a RowNumberFuzzerBase class which is used as the parent class for both RowNumberFuzzer and TopnRowNumberFuzzer.

The fuzzer generates plans only for row_number function right now. It will be enhanced to support rank and dense_rank functions after #11554

Pull Request resolved: #12103

Reviewed By: xiaoxmeng

Differential Revision: D69936162

Pulled By: kagamiori

fbshipit-source-id: 81214748c874f219cf7bc57b5eeeb8039325b06c

velox/exec/fuzzer/CMakeLists.txt

@@ -33,7 +33,6 @@ target_link_libraries(
   velox_type_parser
   Folly::folly
   velox_hive_connector
-  velox_dwio_dwrf_reader
   velox_dwio_dwrf_writer
   velox_dwio_catalog_fbhive
   velox_dwio_faulty_file_sink)
@@ -79,22 +78,39 @@ target_link_libraries(
   velox_aggregation_fuzzer_base
   velox_temp_path)
 
+add_library(velox_row_number_fuzzer_base_lib RowNumberFuzzerBase.cpp)
+
+target_link_libraries(
+  velox_row_number_fuzzer_base_lib
+  velox_dwio_dwrf_reader
+  velox_fuzzer_util
+  velox_vector_fuzzer
+  velox_exec_test_lib)
+
 add_library(velox_row_number_fuzzer_lib RowNumberFuzzer.cpp)
 
+target_link_libraries(
+  velox_row_number_fuzzer_lib velox_row_number_fuzzer_base_lib velox_type
+  velox_expression_test_utility)
+
 # RowNumber Fuzzer.
 add_executable(velox_row_number_fuzzer RowNumberFuzzerRunner.cpp)
 
 target_link_libraries(
   velox_row_number_fuzzer velox_row_number_fuzzer_lib)
 
+add_library(velox_topn_row_number_fuzzer_lib TopNRowNumberFuzzer.cpp)
+
 target_link_libraries(
-  velox_row_number_fuzzer_lib
-  velox_fuzzer_util
-  velox_type
-  velox_vector_fuzzer
-  velox_exec_test_lib
+  velox_topn_row_number_fuzzer_lib velox_row_number_fuzzer_base_lib velox_type
   velox_expression_test_utility)
 
+# TopNRowNumber Fuzzer.
+add_executable(velox_topn_row_number_fuzzer TopNRowNumberFuzzerRunner.cpp)
+
+target_link_libraries(
+  velox_topn_row_number_fuzzer velox_topn_row_number_fuzzer_lib)
+
 add_library(velox_join_fuzzer JoinFuzzer.cpp)
 
 target_link_libraries(

velox/exec/fuzzer/DuckQueryRunner.cpp

@@ -166,6 +166,11 @@ std::optional<std::string> DuckQueryRunner::toSql(
     return toSql(rowNumberNode);
   }
 
+  if (const auto topNRowNumberNode =
+          std::dynamic_pointer_cast<const core::TopNRowNumberNode>(plan)) {
+    return toSql(topNRowNumberNode);
+  }
+
   if (const auto joinNode =
           std::dynamic_pointer_cast<const core::HashJoinNode>(plan)) {
     return toSql(joinNode);
@@ -377,4 +382,55 @@ std::optional<std::string> DuckQueryRunner::toSql(
 
   return sql.str();
 }
+
+std::optional<std::string> DuckQueryRunner::toSql(
+    const std::shared_ptr<const core::TopNRowNumberNode>& topNRowNumberNode) {
+  std::stringstream sql;
+  sql << "SELECT * FROM (SELECT ";
+
+  const auto& inputType = topNRowNumberNode->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 = topNRowNumberNode->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 = topNRowNumberNode->sortingKeys();
+  const auto& sortingOrders = topNRowNumberNode->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();
+    }
+  }
+
+  std::string rowNumberColumnName = topNRowNumberNode->generateRowNumber()
+      ? topNRowNumberNode->outputType()->nameOf(
+            topNRowNumberNode->outputType()->children().size() - 1)
+      : "row_number";
+
+  // TopNRowNumberNode should have a single source.
+  std::optional<std::string> source = toSql(topNRowNumberNode->sources()[0]);
+  if (!source) {
+    return std::nullopt;
+  }
+  sql << ") as " << rowNumberColumnName << " FROM " << *source << ") ";
+  sql << " where " << rowNumberColumnName
+      << " <= " << topNRowNumberNode->limit();
+
+  return sql.str();
+}
+
 } // namespace facebook::velox::exec::test

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>& topNRowNumberNode);
+
   std::unordered_set<std::string> aggregateFunctionNames_;
 };
 

velox/exec/fuzzer/PrestoQueryRunner.cpp

@@ -183,6 +183,11 @@ std::optional<std::string> PrestoQueryRunner::toSql(
     return toSql(rowNumberNode);
   }
 
+  if (const auto topNRowNumberNode =
+          std::dynamic_pointer_cast<const core::TopNRowNumberNode>(plan)) {
+    return toSql(topNRowNumberNode);
+  }
+
   if (auto tableWriteNode =
           std::dynamic_pointer_cast<const core::TableWriteNode>(plan)) {
     return toSql(tableWriteNode);
@@ -497,6 +502,60 @@ std::optional<std::string> PrestoQueryRunner::toSql(
   return sql.str();
 }
 
+std::optional<std::string> PrestoQueryRunner::toSql(
+    const std::shared_ptr<const core::TopNRowNumberNode>& topNRowNumberNode) {
+  if (!isSupportedDwrfType(topNRowNumberNode->sources()[0]->outputType())) {
+    return std::nullopt;
+  }
+
+  std::stringstream sql;
+  sql << "SELECT * FROM (SELECT ";
+
+  const auto& inputType = topNRowNumberNode->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 = topNRowNumberNode->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 = topNRowNumberNode->sortingKeys();
+  const auto& sortingOrders = topNRowNumberNode->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();
+    }
+  }
+
+  std::string rowNumberColumnName = topNRowNumberNode->generateRowNumber()
+      ? topNRowNumberNode->outputType()->nameOf(
+            topNRowNumberNode->outputType()->children().size() - 1)
+      : "row_number";
+
+  // TopNRowNumberNode should have a single source.
+  std::optional<std::string> source = toSql(topNRowNumberNode->sources()[0]);
+  if (!source) {
+    return std::nullopt;
+  }
+  sql << ") as " << rowNumberColumnName << " FROM " << *source << ") ";
+  sql << " where " << rowNumberColumnName
+      << " <= " << topNRowNumberNode->limit();
+
+  return sql.str();
+}
+
 std::optional<std::string> PrestoQueryRunner::toSql(
     const std::shared_ptr<const core::TableWriteNode>& tableWriteNode) {
   auto insertTableHandle =

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);