[sapphire-velox]fix: Fix to handle multiple TaskSource have the same source node

presto-native-execution/presto_cpp/main/TaskManager.cpp

@@ -542,7 +542,7 @@ std::unique_ptr<TaskInfo> TaskManager::createOrUpdateTaskImpl(
   {
     std::lock_guard<std::mutex> l(prestoTask->mutex);
     prestoTask->updateCoordinatorHeartbeatLocked();
-    if (not prestoTask->task && planFragment.planNode) {
+    if ((prestoTask->task == nullptr) && (planFragment.planNode != nullptr)) {
       // If the task is aborted, no need to do anything else.
       // This takes care of DELETE task message coming before CREATE task.
       if (prestoTask->info.taskStatus.state == protocol::TaskState::ABORTED) {
@@ -611,7 +611,33 @@ std::unique_ptr<TaskInfo> TaskManager::createOrUpdateTaskImpl(
       VLOG(1) << "Failed to update output buffers for task: " << taskId;
     }
 
+    folly::F14FastMap<protocol::PlanNodeId, protocol::TaskSource> sourcesMap;
     for (const auto& source : sources) {
+      auto it = sourcesMap.find(source.planNodeId);
+      if (it == sourcesMap.end()) {
+        // No existing source with same planNodeId, add as new
+        sourcesMap.emplace(source.planNodeId, source);
+        continue;
+      }
+
+      // Merge with existing source that has the same planNodeId
+      auto& merged = it->second;
+
+      // Merge splits
+      merged.splits.insert(
+          merged.splits.end(), source.splits.begin(), source.splits.end());
+
+      // Merge noMoreSplitsForLifespan
+      merged.noMoreSplitsForLifespan.insert(
+          merged.noMoreSplitsForLifespan.end(),
+          source.noMoreSplitsForLifespan.begin(),
+          source.noMoreSplitsForLifespan.end());
+
+      // Use OR logic for noMoreSplits flag
+      merged.noMoreSplits = merged.noMoreSplits || source.noMoreSplits;
+    }
+
+    for (const auto& [_, source] : sourcesMap) {
       // Add all splits from the source to the task.
       VLOG(1) << "Adding " << source.splits.size() << " splits to " << taskId
               << " for node " << source.planNodeId;

presto-native-execution/presto_cpp/main/tests/TaskManagerTest.cpp

@@ -715,6 +715,39 @@ TEST_P(TaskManagerTest, tableScanAllSplitsAtOnce) {
   assertResults(taskId, rowType_, "SELECT * FROM tmp WHERE c0 % 5 = 0");
 }
 
+TEST_P(TaskManagerTest, addSplitsWithSameSourceNode) {
+  const auto tableDir = exec::test::TempDirectoryPath::create();
+  auto filePaths = makeFilePaths(tableDir, 5);
+  auto vectors = makeVectors(filePaths.size(), 1'000);
+  for (int i = 0; i < filePaths.size(); i++) {
+    writeToFile(filePaths[i], vectors[i]);
+  }
+  duckDbQueryRunner_.createTable("tmp", vectors);
+
+  const auto planFragment = exec::test::PlanBuilder()
+                          .tableScan(rowType_)
+                          .filter("c0 % 5 = 0")
+                          .partitionedOutput({}, 1, {"c0", "c1"}, GetParam())
+                          .planFragment();
+
+  protocol::TaskUpdateRequest updateRequest;
+  // Create multiple task sources with the same source node id.
+  std::vector<protocol::TaskSource> taskSources;
+  taskSources.reserve(filePaths.size());
+  long splitSequenceId{0};
+  for (const auto& filePath : filePaths) {
+    taskSources.push_back(makeSource("0", {filePath}, /*noMoreSplits=*/true, splitSequenceId));
+  }
+  taskSources.reserve(filePaths.size());
+  updateRequest.sources = std::move(taskSources);
+
+  protocol::TaskId taskId = "scan.0.0.1.0";
+  auto taskInfo = createOrUpdateTask(taskId, updateRequest, planFragment);
+
+  ASSERT_GE(taskInfo->stats.queuedTimeInNanos, 0);
+  assertResults(taskId, rowType_, "SELECT * FROM tmp WHERE c0 % 5 = 0");
+}
+
 TEST_P(TaskManagerTest, fecthFromFinishedTask) {
   const auto tableDir = exec::test::TempDirectoryPath::create();
   auto filePaths = makeFilePaths(tableDir, 5);

presto-native-execution/src/test/java/com/facebook/presto/spark/TestPrestoSparkNativeGeneralQueries.java

@@ -121,9 +121,4 @@ public void testCatalogWithCacheEnabled() {}
     @Override
     @Ignore
     public void testKeyBasedSamplingInlined() {}
-
-    // VeloxRuntimeError: !noMoreSplits_
-    @Override
-    @Ignore
-    public void testUnionAllInsert() {}
 }