feat: Add overriding file create config for different operators (#16318)

tanjialiang · meta-codesync[bot] · commit b88ce66a7b49 · 2026-02-10T18:31:12.000-08:00
Summary: Pull Request resolved: #16318 for different operator types, we might need different file create configs. this change allows this behavior. Reviewed By: xiaoxmeng Differential Revision: D92762382 fbshipit-source-id: 2ee9a1d1038b71023029997d66b243bef90370fc
diff --git a/velox/core/QueryConfig.h b/velox/core/QueryConfig.h
@@ -420,7 +420,19 @@ class QueryConfig {
   static constexpr const char* kSpillFileCreateConfig =
       "spill_file_create_config";
 
-  /// Default offset spill start partition bit. It is used with
+  /// Config used to create aggregation spill files. This config is provided to
+  /// underlying file system and the config is free form. The form should be
+  /// defined by the underlying file system.
+  static constexpr const char* kAggregationSpillFileCreateConfig =
+      "aggregation_spill_file_create_config";
+
+  /// Config used to create hash join spill files. This config is provided to
+  /// underlying file system and the config is free form. The form should be
+  /// defined by the underlying file system.
+  static constexpr const char* kHashJoinSpillFileCreateConfig =
+      "hash_join_spill_file_create_config";
+
+  /// Default offset spill start partition bit.
   /// 'kSpillNumPartitionBits' together to
   /// calculate the spilling partition number for join spill or aggregation
   /// spill.
@@ -1193,6 +1205,14 @@ class QueryConfig {
     return get<std::string>(kSpillFileCreateConfig, "");
   }
 
+  std::string aggregationSpillFileCreateConfig() const {
+    return get<std::string>(kAggregationSpillFileCreateConfig, "");
+  }
+
+  std::string hashJoinSpillFileCreateConfig() const {
+    return get<std::string>(kHashJoinSpillFileCreateConfig, "");
+  }
+
   int32_t minSpillableReservationPct() const {
     constexpr int32_t kDefaultPct = 5;
     return get<int32_t>(kMinSpillableReservationPct, kDefaultPct);
diff --git a/velox/core/tests/QueryConfigTest.cpp b/velox/core/tests/QueryConfigTest.cpp
@@ -280,4 +280,54 @@ TEST_F(QueryConfigTest, singleSourceExchangeOptimizationConfig) {
   }
 }
 
+TEST_F(QueryConfigTest, operatorSpillFileCreateConfig) {
+  // Test default values (empty strings)
+  {
+    auto queryCtx = QueryCtx::create(nullptr, QueryConfig{{}});
+    const QueryConfig& config = queryCtx->queryConfig();
+    EXPECT_EQ(config.aggregationSpillFileCreateConfig(), "");
+    EXPECT_EQ(config.hashJoinSpillFileCreateConfig(), "");
+  }
+
+  // Test with aggregation spill file create config set
+  {
+    std::unordered_map<std::string, std::string> configData(
+        {{QueryConfig::kAggregationSpillFileCreateConfig,
+          "aggregation_config_value"}});
+    auto queryCtx =
+        QueryCtx::create(nullptr, QueryConfig{std::move(configData)});
+    const QueryConfig& config = queryCtx->queryConfig();
+    EXPECT_EQ(
+        config.aggregationSpillFileCreateConfig(), "aggregation_config_value");
+    EXPECT_EQ(config.hashJoinSpillFileCreateConfig(), "");
+  }
+
+  // Test with hash join spill file create config set
+  {
+    std::unordered_map<std::string, std::string> configData(
+        {{QueryConfig::kHashJoinSpillFileCreateConfig,
+          "hashjoin_config_value"}});
+    auto queryCtx =
+        QueryCtx::create(nullptr, QueryConfig{std::move(configData)});
+    const QueryConfig& config = queryCtx->queryConfig();
+    EXPECT_EQ(config.aggregationSpillFileCreateConfig(), "");
+    EXPECT_EQ(config.hashJoinSpillFileCreateConfig(), "hashjoin_config_value");
+  }
+
+  // Test with both configs set
+  {
+    std::unordered_map<std::string, std::string> configData(
+        {{QueryConfig::kAggregationSpillFileCreateConfig,
+          "aggregation_config_value"},
+         {QueryConfig::kHashJoinSpillFileCreateConfig,
+          "hashjoin_config_value"}});
+    auto queryCtx =
+        QueryCtx::create(nullptr, QueryConfig{std::move(configData)});
+    const QueryConfig& config = queryCtx->queryConfig();
+    EXPECT_EQ(
+        config.aggregationSpillFileCreateConfig(), "aggregation_config_value");
+    EXPECT_EQ(config.hashJoinSpillFileCreateConfig(), "hashjoin_config_value");
+  }
+}
+
 } // namespace facebook::velox::core::test
diff --git a/velox/exec/Driver.cpp b/velox/exec/Driver.cpp
@@ -107,8 +107,19 @@ velox::memory::MemoryPool* DriverCtx::addOperatorPool(
       planNodeId, splitGroupId, pipelineId, driverId, operatorType);
 }
 
+namespace {
+bool isHashJoinSpillOperator(const std::string& operatorType) {
+  return operatorType == "HashBuild" || operatorType == "HashProbe";
+}
+
+bool isAggregationSpillOperator(const std::string& operatorType) {
+  return operatorType == "Aggregation" || operatorType == "PartialAggregation";
+}
+} // namespace
+
 std::optional<common::SpillConfig> DriverCtx::makeSpillConfig(
-    int32_t operatorId) const {
+    int32_t operatorId,
+    const std::string& operatorType) const {
   const auto& queryConfig = task->queryCtx()->queryConfig();
   if (!queryConfig.spillEnabled()) {
     return std::nullopt;
@@ -126,6 +137,21 @@ std::optional<common::SpillConfig> DriverCtx::makeSpillConfig(
       [this](uint64_t bytes) {
         task->queryCtx()->updateSpilledBytesAndCheckLimit(bytes);
       };
+
+  std::string fileCreateConfig = queryConfig.spillFileCreateConfig();
+  if (isHashJoinSpillOperator(operatorType)) {
+    const auto& hashJoinConfig = queryConfig.hashJoinSpillFileCreateConfig();
+    if (!hashJoinConfig.empty()) {
+      fileCreateConfig = hashJoinConfig;
+    }
+  } else if (isAggregationSpillOperator(operatorType)) {
+    const auto& aggregationConfig =
+        queryConfig.aggregationSpillFileCreateConfig();
+    if (!aggregationConfig.empty()) {
+      fileCreateConfig = aggregationConfig;
+    }
+  }
+
   return common::SpillConfig(
       std::move(getSpillDirPathCb),
       std::move(updateAndCheckSpillLimitCb),
@@ -146,7 +172,7 @@ std::optional<common::SpillConfig> DriverCtx::makeSpillConfig(
       queryConfig.spillPrefixSortEnabled()
           ? std::optional<common::PrefixSortConfig>(prefixSortConfig())
           : std::nullopt,
-      queryConfig.spillFileCreateConfig(),
+      fileCreateConfig,
       queryConfig.windowSpillMinReadBatchRows());
 }
 
diff --git a/velox/exec/Driver.h b/velox/exec/Driver.h
@@ -257,8 +257,11 @@ struct DriverCtx {
       const core::PlanNodeId& planNodeId,
       const std::string& operatorType);
 
-  /// Builds the spill config for the operator with specified 'operatorId'.
-  std::optional<common::SpillConfig> makeSpillConfig(int32_t operatorId) const;
+  /// Builds the spill config for the operator with specified 'operatorId' and
+  /// 'operatorType'.
+  std::optional<common::SpillConfig> makeSpillConfig(
+      int32_t operatorId,
+      const std::string& operatorType) const;
 
   common::PrefixSortConfig prefixSortConfig() const {
     return common::PrefixSortConfig{
diff --git a/velox/exec/HashAggregation.cpp b/velox/exec/HashAggregation.cpp
@@ -35,7 +35,12 @@ HashAggregation::HashAggregation(
               ? "PartialAggregation"
               : "Aggregation",
           aggregationNode->canSpill(driverCtx->queryConfig())
-              ? driverCtx->makeSpillConfig(operatorId)
+              ? driverCtx->makeSpillConfig(
+                    operatorId,
+                    aggregationNode->step() ==
+                            core::AggregationNode::Step::kPartial
+                        ? "PartialAggregation"
+                        : "Aggregation")
               : std::nullopt),
       aggregationNode_(aggregationNode),
       isPartialOutput_(isPartialOutput(aggregationNode->step())),
diff --git a/velox/exec/HashBuild.cpp b/velox/exec/HashBuild.cpp
@@ -58,7 +58,7 @@ HashBuild::HashBuild(
           joinNode->id(),
           "HashBuild",
           joinNode->canSpill(driverCtx->queryConfig())
-              ? driverCtx->makeSpillConfig(operatorId)
+              ? driverCtx->makeSpillConfig(operatorId, "HashBuild")
               : std::nullopt),
       joinNode_(std::move(joinNode)),
       joinType_{joinNode_->joinType()},
diff --git a/velox/exec/HashProbe.cpp b/velox/exec/HashProbe.cpp
@@ -120,7 +120,7 @@ HashProbe::HashProbe(
           joinNode->id(),
           "HashProbe",
           joinNode->canSpill(driverCtx->queryConfig())
-              ? driverCtx->makeSpillConfig(operatorId)
+              ? driverCtx->makeSpillConfig(operatorId, "HashProbe")
               : std::nullopt),
       outputBatchSize_{outputBatchRows()},
       joinNode_(std::move(joinNode)),
diff --git a/velox/exec/Merge.cpp b/velox/exec/Merge.cpp
@@ -726,7 +726,7 @@ LocalMerge::LocalMerge(
           localMergeNode->id(),
           "LocalMerge",
           localMergeNode->canSpill(driverCtx->queryConfig())
-              ? driverCtx->makeSpillConfig(operatorId)
+              ? driverCtx->makeSpillConfig(operatorId, "LocalMerge")
               : std::nullopt) {
   VELOX_CHECK_EQ(
       operatorCtx_->driverCtx()->driverId,
diff --git a/velox/exec/Operator.h b/velox/exec/Operator.h
@@ -522,6 +522,12 @@ class Operator : public BaseRuntimeStatWriter {
     return input_ != nullptr;
   }
 
+  /// Returns the spill config for this operator. This method is only used for
+  /// test.
+  const common::SpillConfig* testingSpillConfig() const {
+    return spillConfig();
+  }
+
  protected:
   static std::vector<std::unique_ptr<PlanNodeTranslator>>& translators();
   friend class NonReclaimableSection;
diff --git a/velox/exec/OrderBy.cpp b/velox/exec/OrderBy.cpp
@@ -41,7 +41,7 @@ OrderBy::OrderBy(
           orderByNode->id(),
           "OrderBy",
           orderByNode->canSpill(driverCtx->queryConfig())
-              ? driverCtx->makeSpillConfig(operatorId)
+              ? driverCtx->makeSpillConfig(operatorId, "OrderBy")
               : std::nullopt) {
   maxOutputRows_ = outputBatchRows(std::nullopt);
   VELOX_CHECK(pool()->trackUsage());
diff --git a/velox/exec/RowNumber.cpp b/velox/exec/RowNumber.cpp
@@ -30,7 +30,7 @@ RowNumber::RowNumber(
           rowNumberNode->id(),
           "RowNumber",
           rowNumberNode->canSpill(driverCtx->queryConfig())
-              ? driverCtx->makeSpillConfig(operatorId)
+              ? driverCtx->makeSpillConfig(operatorId, "RowNumber")
               : std::nullopt),
       limit_{rowNumberNode->limit()},
       generateRowNumber_{rowNumberNode->generateRowNumber()} {
diff --git a/velox/exec/TableWriter.cpp b/velox/exec/TableWriter.cpp
@@ -30,7 +30,7 @@ TableWriter::TableWriter(
           tableWriteNode->id(),
           "TableWrite",
           tableWriteNode->canSpill(driverCtx->queryConfig())
-              ? driverCtx->makeSpillConfig(operatorId)
+              ? driverCtx->makeSpillConfig(operatorId, "TableWrite")
               : std::nullopt),
       driverCtx_(driverCtx),
       connectorPool_(driverCtx_->task->addConnectorPoolLocked(
diff --git a/velox/exec/TopNRowNumber.cpp b/velox/exec/TopNRowNumber.cpp
@@ -133,7 +133,7 @@ TopNRowNumber::TopNRowNumber(
           node->id(),
           "TopNRowNumber",
           node->canSpill(driverCtx->queryConfig())
-              ? driverCtx->makeSpillConfig(operatorId)
+              ? driverCtx->makeSpillConfig(operatorId, "TopNRowNumber")
               : std::nullopt),
       rankFunction_(node->rankFunction()),
       limit_{node->limit()},
diff --git a/velox/exec/Window.cpp b/velox/exec/Window.cpp
@@ -45,7 +45,7 @@ Window::Window(
           windowNode->id(),
           "Window",
           windowNode->canSpill(driverCtx->queryConfig())
-              ? driverCtx->makeSpillConfig(operatorId)
+              ? driverCtx->makeSpillConfig(operatorId, "Window")
               : std::nullopt),
       numInputColumns_(windowNode->inputType()->size()),
       windowNode_(windowNode),
diff --git a/velox/exec/tests/AggregationTest.cpp b/velox/exec/tests/AggregationTest.cpp
@@ -4245,4 +4245,64 @@ TEST_F(AggregationTest, keysProvideCustomComparison) {
       {makeRowVector({c0, c1}), c1},
       {makeRowVector({e0, e1}), e1});
 }
+
+// Test that aggregation spill uses the aggregation_spill_file_create_config
+// when set, and other spillable operators use the default
+// spill_file_create_config.
+DEBUG_ONLY_TEST_F(AggregationTest, aggregationSpillFileCreateConfig) {
+  auto vectors = makeVectors(rowType_, 32, 100);
+  createDuckDbTable(vectors);
+
+  auto tempDirectory = exec::test::TempDirectoryPath::create();
+
+  std::atomic_bool aggregationConfigVerified{false};
+  std::atomic_bool defaultConfigVerified{false};
+  SCOPED_TESTVALUE_SET(
+      "facebook::velox::exec::Driver::runInternal::isBlocked",
+      std::function<void(exec::Operator*)>([&](exec::Operator* op) {
+        const auto* spillConfig = op->testingSpillConfig();
+        if (spillConfig == nullptr) {
+          return;
+        }
+        const auto& opType = op->operatorType();
+        if (opType == "Aggregation" || opType == "PartialAggregation") {
+          // Aggregation operators should use
+          // aggregation_spill_file_create_config.
+          ASSERT_EQ(spillConfig->fileCreateConfig, "test_aggregation_config")
+              << "Operator: " << opType;
+          aggregationConfigVerified = true;
+        } else {
+          // Other spillable operators (e.g., OrderBy) should use the default
+          // spill_file_create_config.
+          ASSERT_EQ(spillConfig->fileCreateConfig, "test_default_config")
+              << "Operator: " << opType;
+          defaultConfigVerified = true;
+        }
+      }));
+
+  // Build a plan with aggregation and orderBy. Aggregation operators should use
+  // aggregation_spill_file_create_config and orderBy should use the default
+  // spill_file_create_config.
+  TestScopedSpillInjection scopedSpillInjection(100);
+  AssertQueryBuilder(duckDbQueryRunner_)
+      .spillDirectory(tempDirectory->getPath())
+      .config(QueryConfig::kSpillEnabled, true)
+      .config(QueryConfig::kAggregationSpillEnabled, true)
+      .config(QueryConfig::kOrderBySpillEnabled, true)
+      .config(QueryConfig::kSpillFileCreateConfig, "test_default_config")
+      .config(
+          QueryConfig::kAggregationSpillFileCreateConfig,
+          "test_aggregation_config")
+      .plan(
+          PlanBuilder()
+              .values(vectors)
+              .singleAggregation({"c0", "c1"}, {"sum(c2)"})
+              .orderBy({"c0 ASC NULLS LAST"}, false)
+              .planNode())
+      .assertResults(
+          "SELECT c0, c1, sum(c2) FROM tmp GROUP BY c0, c1 ORDER BY c0 ASC NULLS LAST");
+
+  ASSERT_TRUE(aggregationConfigVerified.load());
+  ASSERT_TRUE(defaultConfigVerified.load());
+}
 } // namespace facebook::velox::exec::test
diff --git a/velox/exec/tests/HashJoinTest.cpp b/velox/exec/tests/HashJoinTest.cpp
@@ -8724,5 +8724,74 @@ VELOX_INSTANTIATE_TEST_SUITE_P(
       return TestParamToName(info.param);
     });
 
+// Test that hash join spill uses the hash_join_spill_file_create_config when
+// set, and other spillable operators use the default spill_file_create_config.
+DEBUG_ONLY_TEST_P(HashJoinTest, hashJoinSpillFileCreateConfig) {
+  const auto rowType =
+      ROW({"c0", "c1", "c2"}, {INTEGER(), INTEGER(), VARCHAR()});
+  const auto probeVectors = createVectors(rowType, 128, 10);
+  const auto buildVectors = createVectors(rowType, 128, 10);
+
+  auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
+  // Build a plan with hash join and orderBy. Hash join operators should use
+  // hash_join_spill_file_create_config and orderBy should use the default
+  // spill_file_create_config.
+  auto plan = PlanBuilder(planNodeIdGenerator)
+                  .values(probeVectors, true)
+                  .hashJoin(
+                      {"c0"},
+                      {"u0"},
+                      PlanBuilder(planNodeIdGenerator)
+                          .values(buildVectors, true)
+                          .project({"c0 AS u0", "c1 AS u1", "c2 AS u2"})
+                          .planNode(),
+                      "",
+                      {"c0", "c1", "c2"},
+                      core::JoinType::kInner)
+                  .orderBy({"c0 ASC NULLS LAST"}, false)
+                  .planNode();
+
+  auto spillDirectory = exec::test::TempDirectoryPath::create();
+
+  std::atomic_bool hashJoinConfigVerified{false};
+  std::atomic_bool defaultConfigVerified{false};
+  SCOPED_TESTVALUE_SET(
+      "facebook::velox::exec::Driver::runInternal::isBlocked",
+      std::function<void(exec::Operator*)>([&](exec::Operator* op) {
+        const auto* spillConfig = op->testingSpillConfig();
+        if (spillConfig == nullptr) {
+          return;
+        }
+        const auto& opType = op->operatorType();
+        if (opType == "HashBuild" || opType == "HashProbe") {
+          // Hash join operators should use hash_join_spill_file_create_config.
+          ASSERT_EQ(spillConfig->fileCreateConfig, "test_hashjoin_config")
+              << "Operator: " << opType;
+          hashJoinConfigVerified = true;
+        } else {
+          // Other spillable operators (e.g., OrderBy) should use the default
+          // spill_file_create_config.
+          ASSERT_EQ(spillConfig->fileCreateConfig, "test_default_config")
+              << "Operator: " << opType;
+          defaultConfigVerified = true;
+        }
+      }));
+
+  TestScopedSpillInjection scopedSpillInjection(100);
+  AssertQueryBuilder(plan)
+      .spillDirectory(spillDirectory->getPath())
+      .config(core::QueryConfig::kSpillEnabled, true)
+      .config(core::QueryConfig::kJoinSpillEnabled, true)
+      .config(core::QueryConfig::kOrderBySpillEnabled, true)
+      .config(core::QueryConfig::kSpillFileCreateConfig, "test_default_config")
+      .config(
+          core::QueryConfig::kHashJoinSpillFileCreateConfig,
+          "test_hashjoin_config")
+      .copyResults(pool_.get());
+
+  ASSERT_TRUE(hashJoinConfigVerified.load());
+  ASSERT_TRUE(defaultConfigVerified.load());
+}
+
 } // namespace
 } // namespace facebook::velox::exec
diff --git a/velox/experimental/cudf/exec/CudfHashAggregation.cpp b/velox/experimental/cudf/exec/CudfHashAggregation.cpp
