From e56ac7905969fc57c9bdf0aeba6acc244e71ccdc Mon Sep 17 00:00:00 2001
From: Krishna Pai <kpai@fb.com>
Date: Tue, 27 Feb 2024 15:55:40 -0800
Subject: [PATCH] Add support of velox vectors in ReferenceQueryRunner.

---
 velox/exec/fuzzer/PrestoQueryRunner.cpp  | 24 ++++++++++++++++++++----
 velox/exec/fuzzer/PrestoQueryRunner.h    | 12 ++++++++++++
 velox/exec/fuzzer/ReferenceQueryRunner.h | 17 +++++++++++++++++
 3 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/velox/exec/fuzzer/PrestoQueryRunner.cpp b/velox/exec/fuzzer/PrestoQueryRunner.cpp
index e092ad5a84aaf..9d6a8853795ea 100644
--- a/velox/exec/fuzzer/PrestoQueryRunner.cpp
+++ b/velox/exec/fuzzer/PrestoQueryRunner.cpp
@@ -491,6 +491,13 @@ std::multiset<std::vector<variant>> PrestoQueryRunner::execute(
     const std::string& sql,
     const std::vector<RowVectorPtr>& input,
     const RowTypePtr& resultType) {
+  return exec::test::materialize(executeImpl(sql, input, resultType));
+}
+
+std::vector<velox::RowVectorPtr> PrestoQueryRunner::executeImpl(
+    const std::string& sql,
+    const std::vector<velox::RowVectorPtr>& input,
+    const velox::RowTypePtr& resultType) {
   auto inputType = asRowType(input[0]->type());
   if (inputType->size() == 0) {
     // The query doesn't need to read any columns, but it needs to see a
@@ -509,7 +516,7 @@ std::multiset<std::vector<variant>> PrestoQueryRunner::execute(
         nullptr,
         numInput,
         std::vector<VectorPtr>{column});
-    return execute(sql, {rowVector}, resultType);
+    return executeImpl(sql, {rowVector}, resultType);
   }
 
   // Create tmp table in Presto using DWRF file format and add a single
@@ -548,9 +555,7 @@ std::multiset<std::vector<variant>> PrestoQueryRunner::execute(
   writeToFile(newFilePath, input, writerPool.get());
 
   // Run the query.
-  results = execute(sql);
-
-  return exec::test::materialize(results);
+  return execute(sql);
 }
 
 std::vector<RowVectorPtr> PrestoQueryRunner::execute(const std::string& sql) {
@@ -611,4 +616,15 @@ std::string PrestoQueryRunner::fetchNext(const std::string& nextUri) {
   return response.text;
 }
 
+bool PrestoQueryRunner::supportsVeloxVectors() {
+  return true;
+}
+
+std::vector<RowVectorPtr> PrestoQueryRunner::executeVector(
+    const std::string& sql,
+    const std::vector<RowVectorPtr>& input,
+    const RowTypePtr& resultType) {
+  return executeImpl(sql, input, resultType);
+}
+
 } // namespace facebook::velox::exec::test
diff --git a/velox/exec/fuzzer/PrestoQueryRunner.h b/velox/exec/fuzzer/PrestoQueryRunner.h
index a52f6359b9783..6e85698788343 100644
--- a/velox/exec/fuzzer/PrestoQueryRunner.h
+++ b/velox/exec/fuzzer/PrestoQueryRunner.h
@@ -60,6 +60,18 @@ class PrestoQueryRunner : public velox::exec::test::ReferenceQueryRunner {
   /// the query must already exist.
   std::vector<velox::RowVectorPtr> execute(const std::string& sql);
 
+  std::vector<velox::RowVectorPtr> executeImpl(
+      const std::string& sql,
+      const std::vector<velox::RowVectorPtr>& input,
+      const velox::RowTypePtr& resultType);
+
+  bool supportsVeloxVectors();
+
+  std::vector<RowVectorPtr> executeVector(
+      const std::string& sql,
+      const std::vector<RowVectorPtr>& input,
+      const RowTypePtr& resultType);
+
  private:
   velox::memory::MemoryPool* rootPool() {
     return rootPool_.get();
diff --git a/velox/exec/fuzzer/ReferenceQueryRunner.h b/velox/exec/fuzzer/ReferenceQueryRunner.h
index b9e5b3dd048c7..dd3e8854fbae4 100644
--- a/velox/exec/fuzzer/ReferenceQueryRunner.h
+++ b/velox/exec/fuzzer/ReferenceQueryRunner.h
@@ -35,6 +35,23 @@ class ReferenceQueryRunner {
       const std::string& sql,
       const std::vector<RowVectorPtr>& input,
       const RowTypePtr& resultType) = 0;
+
+  /// Returns true if the underlying Query Runner supports
+  /// returning the results of a query in Velox Vector format. If supported
+  /// then 'executeVector' can be called to get results as Velox Vector.
+  virtual bool supportsVeloxVectors() const {
+    return false;
+  }
+
+  /// Executes SQL query returned by the 'toSql' method using 'input' data
+  /// and returns the results in RowVector format. Check if this function is
+  /// supported by calling the 'supportsVeloxVectors' function before calling.
+  virtual std::vector<RowVectorPtr> executeVector(
+      const std::string& sql,
+      const std::vector<RowVectorPtr>& input,
+      const RowTypePtr& resultType) {
+    VELOX_NYI("Not supported in this Reference Query Runner.");
+  }
 };
 
 } // namespace facebook::velox::exec::test