Add Spark runner type

facebookincubator · Aug 19, 2024 · b5159f4 · b5159f4
1 parent 62db265
commit b5159f4
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 1 deletion.
diff --git a/velox/exec/fuzzer/ReferenceQueryRunner.h b/velox/exec/fuzzer/ReferenceQueryRunner.h
@@ -29,7 +29,7 @@ class ReferenceQueryRunner {
 
   virtual ~ReferenceQueryRunner() = default;
 
-  enum class RunnerType { kPrestoQueryRunner, kDuckQueryRunner };
+  enum class RunnerType { kPrestoQueryRunner, kDuckQueryRunner, kSparkQueryRunner };
 
   virtual RunnerType runnerType() const = 0;
 

diff --git a/velox/functions/sparksql/fuzzer/SparkQueryRunner.cpp b/velox/functions/sparksql/fuzzer/SparkQueryRunner.cpp
@@ -73,6 +73,23 @@ void writeToFile(
 
 } // namespace
 
+const std::vector<TypePtr>& SparkQueryRunner::supportedScalarTypes() const {
+  static const std::vector<TypePtr> kScalarTypes{
+      BOOLEAN(),
+      TINYINT(),
+      SMALLINT(),
+      INTEGER(),
+      BIGINT(),
+      REAL(),
+      DOUBLE(),
+      VARCHAR(),
+      VARBINARY(),
+      TIMESTAMP(),
+      DATE(),
+  };
+  return kScalarTypes;
+}
+
 std::optional<std::string> SparkQueryRunner::toSql(
     const velox::core::PlanNodePtr& plan) {
   if (const auto aggregationNode =

diff --git a/velox/functions/sparksql/fuzzer/SparkQueryRunner.h b/velox/functions/sparksql/fuzzer/SparkQueryRunner.h
@@ -47,6 +47,12 @@ class SparkQueryRunner : public velox::exec::test::ReferenceQueryRunner {
         pool_(rootPool->addLeafChild("leaf")),
         copyPool_(rootPool->addLeafChild("copy")) {};
 
+  RunnerType runnerType() const override {
+    return RunnerType::kSparkQueryRunner;
+  }
+
+  const std::vector<TypePtr>& supportedScalarTypes() const override;
+
   /// Converts Velox query plan to Spark SQL. Supports Values -> Aggregation.
   /// Values node is converted into reading from 'tmp' table.
   /// @return std::nullopt for unsupported cases.