Reject duplicate segment names before batch push

pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/CustomDataQueryClusterIntegrationTest.java

@@ -18,6 +18,7 @@
  */
 package org.apache.pinot.integration.tests.custom;
 
+import com.fasterxml.jackson.databind.JsonNode;
 import java.io.Closeable;
 import java.io.File;
 import java.io.IOException;
@@ -53,6 +54,7 @@
 import org.apache.pinot.spi.stream.StreamConfigProperties;
 import org.apache.pinot.spi.utils.builder.ControllerRequestURLBuilder;
 import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.apache.pinot.spi.utils.builder.TableNameBuilder;
 import org.apache.pinot.util.TestUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -161,6 +163,17 @@ protected void setUpTable()
             _tarDir);
         uploadSegments(getTableName(), _tarDir);
       }
+      if (!avroFiles.isEmpty()) {
+        String offlineTableName = TableNameBuilder.OFFLINE.tableNameWithType(getTableName());
+        // MSE skips optional segments even when SSE can already query them.
+        TestUtils.waitForCondition(() -> {
+          int numRoutableSegments = 0;
+          for (JsonNode segments : getDebugInfo("debug/routingTable/" + offlineTableName).path(offlineTableName)) {
+            numRoutableSegments += segments.size();
+          }
+          return numRoutableSegments == avroFiles.size();
+        }, 100L, 60_000L, "Broker did not route all offline segments for " + offlineTableName, Duration.ofSeconds(5));
+      }
     }
   }
 

pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-common/src/main/java/org/apache/pinot/plugin/ingestion/batch/common/BaseSegmentPushJobRunner.java

@@ -101,6 +101,7 @@ public void initFileSys() {
   ///                                tarPaths (values), or both may be used depending on upload mode.
   public List<String> getSegmentsToReplace(Map<String, String> segmentsUriToTarPathMap) {
     Collection<String> tarFilePaths = segmentsUriToTarPathMap.values();
+    SegmentPushUtils.validateUniqueSegmentNames(tarFilePaths);
     List<String> segmentNames = new ArrayList<>(tarFilePaths.size());
     for (String tarFilePath : tarFilePaths) {
       segmentNames.add(SegmentPushUtils.getSegmentNameFromTarFile(tarFilePath));

pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-3/pom.xml

@@ -75,6 +75,11 @@
       <artifactId>pinot-csv</artifactId>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>org.mockito</groupId>
+      <artifactId>mockito-core</artifactId>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 
   <profiles>

pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-3/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark3/SparkSegmentMetadataPushJobRunner.java

@@ -103,9 +103,6 @@ public void run()
 
     // Collect files to process
     List<String> segmentsToPush = getSegmentsToPush(outputDirFS, outputDirURI);
-    if (segmentsToPush.isEmpty()) {
-      return;
-    }
 
     // Ensure tableConfigURI is set if missing
     setupTableConfigURI();
@@ -114,6 +111,9 @@ public void run()
     TableConfig tableConfig =
         SegmentGenerationUtils.getTableConfig(_spec.getTableSpec().getTableConfigURI(), _spec.getAuthToken());
     boolean consistentPushEnabled = ConsistentDataPushUtils.consistentDataPushEnabled(tableConfig);
+    if (segmentsToPush.isEmpty() && !consistentPushEnabled) {
+      return;
+    }
 
     // Determine push parallelism
     int pushParallelism = _spec.getPushJobSpec().getPushParallelism();
@@ -136,6 +136,11 @@ private void handleConsistentPush(List<String> segmentsToPush, URI outputDirURI,
             segmentsToPush.toArray(new String[0]));
     Map<URI, String> uriToLineageEntryIdMap =
         ConsistentDataPushUtils.preUpload(_spec, getSegmentsToReplace(segmentsToPushUriToTarPathMap));
+    if (segmentsToPush.isEmpty()) {
+      // An empty REFRESH snapshot still replaces the existing segments.
+      executePostUpload(uriToLineageEntryIdMap);
+      return;
+    }
 
     if (pushParallelism == 1) {
       // Single push
@@ -306,6 +311,7 @@ private void executePostUpload(Map<URI, String> uriToLineageEntryIdMap) {
   ///                                tarPaths (values), or both may be used depending on upload mode.
   public List<String> getSegmentsToReplace(Map<String, String> segmentsUriToTarPathMap) {
     Collection<String> tarFilePaths = segmentsUriToTarPathMap.values();
+    SegmentPushUtils.validateUniqueSegmentNames(tarFilePaths);
     List<String> segmentNames = new ArrayList<>(tarFilePaths.size());
     for (String tarFilePath : tarFilePaths) {
       segmentNames.add(SegmentPushUtils.getSegmentNameFromTarFile(tarFilePath));

pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-3/src/test/java/org/apache/pinot/plugin/ingestion/batch/spark3/SparkSegmentMetadataPushJobRunnerTest.java

@@ -0,0 +1,87 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.plugin.ingestion.batch.spark3;
+
+import java.net.URI;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.List;
+import java.util.Map;
+import org.apache.commons.io.FileUtils;
+import org.apache.pinot.segment.local.utils.ConsistentDataPushUtils;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.ingestion.BatchIngestionConfig;
+import org.apache.pinot.spi.config.table.ingestion.IngestionConfig;
+import org.apache.pinot.spi.ingestion.batch.spec.PushJobSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.TableSpec;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.mockito.MockedStatic;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mockStatic;
+import static org.mockito.Mockito.times;
+
+
+/// Verifies that empty selections retain consistent REFRESH replacement without starting Spark uploads.
+public class SparkSegmentMetadataPushJobRunnerTest {
+  @DataProvider
+  public Object[][] emptyPushModes() {
+    return new Object[][]{{true, 0}, {true, 1}, {true, 2}, {false, 0}};
+  }
+
+  @Test(dataProvider = "emptyPushModes")
+  public void testEmptySelectionHonorsConsistentPush(boolean consistentPush, int parallelism)
+      throws Exception {
+    Path outputDir = Files.createTempDirectory("empty-metadata-push-");
+    try {
+      IngestionConfig ingestion = new IngestionConfig();
+      ingestion.setBatchIngestionConfig(new BatchIngestionConfig(null, "REFRESH", "DAILY", consistentPush));
+      TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName("test")
+          .setIngestionConfig(ingestion).build();
+      Path tableConfigFile = outputDir.resolve("table.json");
+      Files.writeString(tableConfigFile, tableConfig.toJsonString());
+      TableSpec tableSpec = new TableSpec();
+      tableSpec.setTableName("test");
+      tableSpec.setTableConfigURI(tableConfigFile.toUri().toString());
+      SegmentGenerationJobSpec spec = new SegmentGenerationJobSpec();
+      spec.setTableSpec(tableSpec);
+      spec.setOutputDirURI(outputDir.toUri().toString());
+      spec.setPinotFSSpecs(List.of());
+      PushJobSpec pushSpec = new PushJobSpec();
+      pushSpec.setPushParallelism(parallelism);
+      spec.setPushJobSpec(pushSpec);
+      Map<URI, String> lineage = Map.of(URI.create("http://localhost:9000"), "lineage");
+      try (MockedStatic<ConsistentDataPushUtils> consistent = mockStatic(ConsistentDataPushUtils.class)) {
+        consistent.when(() -> ConsistentDataPushUtils.consistentDataPushEnabled(any())).thenCallRealMethod();
+        consistent.when(() -> ConsistentDataPushUtils.preUpload(spec, List.of())).thenReturn(lineage);
+        SparkSegmentMetadataPushJobRunner runner = new SparkSegmentMetadataPushJobRunner();
+        runner.init(spec);
+        runner.run();
+        consistent.verify(() -> ConsistentDataPushUtils.preUpload(spec, List.of()), times(consistentPush ? 1 : 0));
+        consistent.verify(() -> ConsistentDataPushUtils.postUpload(spec, lineage), times(consistentPush ? 1 : 0));
+      }
+    } finally {
+      FileUtils.deleteDirectory(outputDir.toFile());
+    }
+  }
+}

pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-base/src/test/java/org/apache/pinot/plugin/ingestion/batch/spark/common/DuplicateSegmentMetadataPushJobRunnerTest.java

@@ -0,0 +1,99 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.plugin.ingestion.batch.spark.common;
+
+import java.net.URI;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.commons.io.FileUtils;
+import org.apache.pinot.spi.ingestion.batch.runner.IngestionJobRunner;
+import org.apache.pinot.spi.ingestion.batch.spec.PinotFSSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.PushJobSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertThrows;
+
+
+/// Exercises driver-side filtering with real local files before dispatching segment pushes.
+public class DuplicateSegmentMetadataPushJobRunnerTest {
+  @DataProvider
+  public Object[][] pushModes() {
+    return new Object[][]{{"metadata"}, {"tar"}, {"uri"}};
+  }
+
+  @Test(dataProvider = "pushModes")
+  public void testPushPatternFiltersBeforeDuplicateValidation(String mode)
+      throws Exception {
+    Path outputDir = Files.createTempDirectory("segment-push-% ");
+    try {
+      Path current = Files.createDirectory(outputDir.resolve("current"));
+      Path selected = Files.createFile(current.resolve("segment.tar.gz"));
+      Files.createFile(Files.createDirectory(outputDir.resolve("stale")).resolve("segment.tar.gz"));
+      SegmentGenerationJobSpec spec = new SegmentGenerationJobSpec();
+      spec.setOutputDirURI(outputDir.toUri().toString());
+      spec.setPinotFSSpecs(List.of());
+      PushJobSpec pushSpec = new PushJobSpec();
+      pushSpec.setPushParallelism(2);
+      pushSpec.setPushFileNamePattern("glob:**/current/*.tar.gz");
+      spec.setPushJobSpec(pushSpec);
+      List<List<String>> pushed = new ArrayList<>();
+      IngestionJobRunner runner = switch (mode) {
+        case "metadata" -> new BaseSparkSegmentMetadataPushJobRunner(spec) {
+          @Override
+          public void parallelizeMetadataPushJob(List<String> segments, List<PinotFSSpec> fileSystems,
+              int parallelism, URI output) {
+            pushed.add(List.copyOf(segments));
+          }
+        };
+        case "tar" -> new BaseSparkSegmentTarPushJobRunner(spec) {
+          @Override
+          public void parallelizeTarPushJob(List<PinotFSSpec> fileSystems, List<String> segments,
+              int parallelism, URI output) {
+            pushed.add(List.copyOf(segments));
+          }
+        };
+        case "uri" -> new BaseSparkSegmentUriPushJobRunner(spec) {
+          @Override
+          public void parallelizeUriPushJob(List<PinotFSSpec> fileSystems, List<String> segments, int parallelism) {
+            pushed.add(List.copyOf(segments));
+          }
+        };
+        default -> throw new IllegalArgumentException(mode);
+      };
+      List<List<String>> selectedPaths = List.of(
+          List.of(mode.equals("uri") ? selected.toFile().toURI().toString() : selected.toString()));
+      runner.run();
+      assertEquals(pushed, selectedPaths);
+      pushSpec.setPushParallelism(0);
+      pushSpec.setPushFileNamePattern("glob:**/missing/*");
+      runner.run();
+      assertEquals(pushed, selectedPaths);
+      pushSpec.setPushFileNamePattern(null);
+      assertThrows(IllegalArgumentException.class, runner::run);
+      assertEquals(pushed, selectedPaths);
+    } finally {
+      FileUtils.deleteDirectory(outputDir.toFile());
+    }
+  }
+}

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SegmentPushUtils.java

@@ -124,6 +124,9 @@ static int getSocketTimeoutMs(SegmentGenerationJobSpec spec) {
 
   public static URI generateSegmentTarURI(URI dirURI, URI fileURI, String prefix, String suffix) {
     if (StringUtils.isEmpty(prefix) && StringUtils.isEmpty(suffix)) {
+      if (dirURI == null) {
+        return fileURI;
+      }
       // In case the FS doesn't provide scheme or host, will fill it up from dirURI.
       String scheme = fileURI.getScheme();
       if (StringUtils.isEmpty(fileURI.getScheme())) {
@@ -199,6 +202,7 @@ public static void sendSegmentUriAndMetadata(SegmentGenerationJobSpec spec, Pino
   public static void pushSegments(SegmentGenerationJobSpec spec, PinotFS fileSystem, List<String> tarFilePaths,
       List<Header> headers, List<NameValuePair> parameters)
       throws RetriableOperationException, AttemptsExceededException {
+    validateUniqueSegmentNames(tarFilePaths);
     String tableName = spec.getTableSpec().getTableName();
     TableType tableType = tableName.endsWith("_" + TableType.REALTIME.name()) ? TableType.REALTIME : TableType.OFFLINE;
     boolean cleanUpOutputDir = spec.isCleanUpOutputDir();
@@ -340,6 +344,7 @@ public static void sendSegmentUris(SegmentGenerationJobSpec spec, List<String> s
   public static void sendSegmentUriAndMetadata(SegmentGenerationJobSpec spec, PinotFS fileSystem,
       Map<String, String> segmentUriToTarPathMap, List<Header> headers, List<NameValuePair> parameters)
       throws Exception {
+    validateUniqueSegmentNames(segmentUriToTarPathMap.values());
     LOGGER.info("Start pushing segment metadata: {} to locations: {} for table {}", segmentUriToTarPathMap,
         Arrays.toString(spec.getPinotClusterSpecs()), spec.getTableSpec().getTableName());
     FileUploadDownloadClient fileUploadDownloadClient = getOrCreateFileUploadDownloadClient(spec);
@@ -459,6 +464,20 @@ private static File getSegmentMetadataFile(SegmentGenerationJobSpec spec, PinotF
     return generateSegmentMetadataFile(fileSystem, URI.create(tarFilePath));
   }
 
+  /// Rejects ambiguous uploads before HTTP or segment replacement starts. Archives in different output directories
+  /// can otherwise overwrite the same segment in an arbitrary order.
+  /// Extensionless deep-store segment paths remain supported for metadata push.
+  public static void validateUniqueSegmentNames(Collection<String> segmentPaths) {
+    Map<String, String> pathsBySegmentName = new HashMap<>();
+    for (String path : segmentPaths) {
+      String segmentName = getSegmentName(path);
+      String previous = pathsBySegmentName.putIfAbsent(segmentName, path);
+      Preconditions.checkArgument(previous == null,
+          "Multiple archives for segment %s: %s and %s. Remove stale archives or narrow the push file pattern.",
+          segmentName, previous, path);
+    }
+  }
+
   /// Resolves and validates the segment name from a supported gzip, LZ4, or Zstandard archive path.
   public static String getSegmentNameFromTarFile(String tarFilePath) {
     String fileName = new File(tarFilePath).getName();
@@ -492,6 +511,7 @@ private static String getSegmentNameFromMetadataFile(File segmentMetadataFile) {
   public static void sendSegmentsUriAndMetadata(SegmentGenerationJobSpec spec, PinotFS fileSystem,
       Map<String, String> segmentUriToTarPathMap, List<Header> headers, List<NameValuePair> parameters)
       throws Exception {
+    validateUniqueSegmentNames(segmentUriToTarPathMap.values());
     LOGGER.info("Start pushing segment metadata: {} to locations: {} for table: {} with parallelism: {}",
         segmentUriToTarPathMap, Arrays.toString(spec.getPinotClusterSpecs()), spec.getTableSpec().getTableName(),
         spec.getPushJobSpec().getPushParallelism());
@@ -681,6 +701,7 @@ static File createSegmentsMetadataTarFile(Collection<String> segmentURIs, Map<St
   public static Map<String, String> getSegmentUriToTarPathMap(URI outputDirURI, PushJobSpec pushSpec,
       String[] files) {
     Map<String, String> segmentUriToTarPathMap = new HashMap<>();
+    List<String> segmentPaths = new ArrayList<>();
     PathMatcher pushFilePathMatcher = null;
     if (pushSpec.getPushFileNamePattern() != null) {
       pushFilePathMatcher = FileSystems.getDefault().getPathMatcher(pushSpec.getPushFileNamePattern());
@@ -700,11 +721,13 @@ public static Map<String, String> getSegmentUriToTarPathMap(URI outputDirURI, Pu
         continue;
       }
       if (Constants.isSegmentTarFile(uri.getPath())) {
+        segmentPaths.add(file);
         URI updatedURI = SegmentPushUtils.generateSegmentTarURI(outputDirURI, uri, pushSpec.getSegmentUriPrefix(),
             pushSpec.getSegmentUriSuffix());
         segmentUriToTarPathMap.put(updatedURI.toString(), file);
       }
     }
+    validateUniqueSegmentNames(segmentPaths);
     if (segmentUriToTarPathMap.isEmpty()) {
       LOGGER.warn("No segment archives selected under outputDirURI: {} with pushFileNamePattern: {}", outputDirURI,
           pushSpec.getPushFileNamePattern());