This is an automated email from the ASF dual-hosted git repository.
xiangfu0 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new adb87d861c8 Avoid redundant local input copies and support Zstandard
batch archives (#19898)
adb87d861c8 is described below
commit adb87d861c87a09bdd902cbf4e56bf7d6832483c
Author: Xiang Fu <[email protected]>
AuthorDate: Fri Oct 23 15:42:45 2026 -0700
Avoid redundant local input copies and support Zstandard batch archives
(#19898)
* Avoid redundant local input copies and support Zstandard batch archives
**Usage:** set `segmentTarCompressionCodec: zstd` in batch YAML for
standalone, Spark, or Hadoop; the [README](README.md) includes an example. Gzip
remains the default, and metadata sidecars remain gzip.
**Validation:** 44 focused tests, Hadoop/Spark compilation, and
style/license checks passed, including source ownership and archive discovery.
* Document batch archive compatibility and upgrade behavior
Keep gzip as default and enable zstd only when archive consumers support
it; distinguish generation from LZ4 discovery.
Document duplicate-name rejection and validate the SPI and README
style/license checks.
* Filter batch metadata archives before duplicate validation
Reproduce with same-named gzip and zstd archives and a pushFileNamePattern
that selects only one.
Filter on the driver before checking duplicates, and skip empty selections.
* Centralize batch archive codecs and trim redundant tests
Use named codec and archive extension constants, including the gzip default.
Replace repeated mock-based runners and filesystem adapters with focused
checks.
* Address archive ingestion review feedback
Split out LZ4 and duplicate-name changes, move the usage guide to
pinot-docs, and fix codec-aware backfill discovery. Keep archive selection
consistent and isolate the input ownership test filesystem.
* Preserve metadata push URIs and stabilize LZ4 assertions
Handle metadata pushes with a URI prefix and no output directory. Verify
default LZ4 output by content because reading the input can change its tar
access-time header.
* Address empty archive selection and local URI review feedback
Skip consistent replacement when the push pattern selects no archives, and
warn with selection context. Cover literal percent-encoded filenames in the
existing URI tests.
---
.../common/utils/TarCompressionUtilsTest.java | 4 +-
.../batch/common/BaseSegmentPushJobRunner.java | 11 +-
.../batch/hadoop/HadoopSegmentCreationMapper.java | 3 +-
.../hadoop/HadoopSegmentMetadataPushJobRunner.java | 18 +--
.../hadoop/HadoopSegmentTarPushJobRunner.java | 14 +-
.../hadoop/HadoopSegmentUriPushJobRunner.java | 15 +-
.../spark3/SparkSegmentGenerationJobRunner.java | 3 +-
.../spark3/SparkSegmentMetadataPushJobRunner.java | 16 +--
.../BaseSparkSegmentMetadataPushJobRunner.java | 10 +-
.../common/BaseSparkSegmentTarPushJobRunner.java | 10 +-
.../common/BaseSparkSegmentUriPushJobRunner.java | 14 +-
.../BaseSparkSegmentMetadataPushJobRunnerTest.java | 125 +++++++++++++++++
.../standalone/SegmentGenerationJobRunner.java | 23 +++-
.../standalone/SegmentGenerationJobRunnerTest.java | 151 +++++++++++++++++++++
.../segmentuploader/SegmentUploaderDefault.java | 2 +-
.../segment/local/utils/SegmentPushUtils.java | 39 ++++--
.../segment/local/utils/SegmentPushUtilsTest.java | 140 +++++++++++++++++++
.../pinot/spi/ingestion/batch/spec/Constants.java | 20 +++
.../batch/spec/SegmentGenerationJobSpec.java | 20 +++
.../command/LaunchBackfillIngestionJobCommand.java | 35 +++--
.../LaunchBackfillIngestionJobCommandTest.java | 59 ++++++++
21 files changed, 626 insertions(+), 106 deletions(-)
diff --git
a/pinot-common/src/test/java/org/apache/pinot/common/utils/TarCompressionUtilsTest.java
b/pinot-common/src/test/java/org/apache/pinot/common/utils/TarCompressionUtilsTest.java
index 8941168644f..afe4d84fe70 100644
---
a/pinot-common/src/test/java/org/apache/pinot/common/utils/TarCompressionUtilsTest.java
+++
b/pinot-common/src/test/java/org/apache/pinot/common/utils/TarCompressionUtilsTest.java
@@ -126,9 +126,11 @@ public class TarCompressionUtilsTest {
File archive = new File(TAR_DIR, "data" +
TarCompressionUtils.TAR_LZ4_FILE_EXTENSION);
TarCompressionUtils.createCompressedTarFile(inputFile, archive,
compressorName);
if (compressorName.equals(LZ4_FAST)) {
+ assertEquals(TarCompressionUtils.COMPRESSOR_NAME_BY_FILE_EXTENSIONS.get(
+ TarCompressionUtils.TAR_LZ4_FILE_EXTENSION), LZ4_FAST);
File defaultArchive = new File(TAR_DIR, "default" +
TarCompressionUtils.TAR_LZ4_FILE_EXTENSION);
TarCompressionUtils.createCompressedTarFile(inputFile, defaultArchive);
- assertEquals(Files.readAllBytes(defaultArchive.toPath()),
Files.readAllBytes(archive.toPath()));
+
assertEquals(Files.readAllBytes(TarCompressionUtils.untar(defaultArchive,
UNTAR_DIR).get(0).toPath()), content);
}
List<File> extractedFiles = TarCompressionUtils.untar(archive, UNTAR_DIR);
diff --git
a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-common/src/main/java/org/apache/pinot/plugin/ingestion/batch/common/BaseSegmentPushJobRunner.java
b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-common/src/main/java/org/apache/pinot/plugin/ingestion/batch/common/BaseSegmentPushJobRunner.java
index 305d05520b8..efe28e59162 100644
---
a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-common/src/main/java/org/apache/pinot/plugin/ingestion/batch/common/BaseSegmentPushJobRunner.java
+++
b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-common/src/main/java/org/apache/pinot/plugin/ingestion/batch/common/BaseSegmentPushJobRunner.java
@@ -18,7 +18,6 @@
*/
package org.apache.pinot.plugin.ingestion.batch.common;
-import com.google.common.base.Preconditions;
import java.io.File;
import java.io.IOException;
import java.net.URI;
@@ -35,7 +34,6 @@ import org.apache.pinot.spi.env.PinotConfiguration;
import org.apache.pinot.spi.filesystem.PinotFS;
import org.apache.pinot.spi.filesystem.PinotFSFactory;
import org.apache.pinot.spi.ingestion.batch.runner.IngestionJobRunner;
-import org.apache.pinot.spi.ingestion.batch.spec.Constants;
import org.apache.pinot.spi.ingestion.batch.spec.PinotFSSpec;
import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
@@ -105,11 +103,7 @@ public abstract class BaseSegmentPushJobRunner implements
IngestionJobRunner {
Collection<String> tarFilePaths = segmentsUriToTarPathMap.values();
List<String> segmentNames = new ArrayList<>(tarFilePaths.size());
for (String tarFilePath : tarFilePaths) {
- File tarFile = new File(tarFilePath);
- String fileName = tarFile.getName();
-
Preconditions.checkArgument(fileName.endsWith(Constants.TAR_GZ_FILE_EXT));
- String segmentName = fileName.substring(0, fileName.length() -
Constants.TAR_GZ_FILE_EXT.length());
- segmentNames.add(segmentName);
+
segmentNames.add(SegmentPushUtils.getSegmentNameFromTarFile(tarFilePath));
}
return segmentNames;
}
@@ -131,6 +125,9 @@ public abstract class BaseSegmentPushJobRunner implements
IngestionJobRunner {
try {
Map<String, String> segmentsUriToTarPathMap =
SegmentPushUtils.getSegmentUriToTarPathMap(_outputDirURI,
_spec.getPushJobSpec(), _files);
+ if (segmentsUriToTarPathMap.isEmpty()) {
+ return;
+ }
if (_consistentPushEnabled) {
List<String> segmentsToReplace =
getSegmentsToReplace(segmentsUriToTarPathMap);
uriToLineageEntryIdMap = ConsistentDataPushUtils.preUpload(_spec,
segmentsToReplace);
diff --git
a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-hadoop/src/main/java/org/apache/pinot/plugin/ingestion/batch/hadoop/HadoopSegmentCreationMapper.java
b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-hadoop/src/main/java/org/apache/pinot/plugin/ingestion/batch/hadoop/HadoopSegmentCreationMapper.java
index 8400bc69072..d8a0c2897ea 100644
---
a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-hadoop/src/main/java/org/apache/pinot/plugin/ingestion/batch/hadoop/HadoopSegmentCreationMapper.java
+++
b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-hadoop/src/main/java/org/apache/pinot/plugin/ingestion/batch/hadoop/HadoopSegmentCreationMapper.java
@@ -176,7 +176,8 @@ public class HadoopSegmentCreationMapper extends
Mapper<LongWritable, Text, Long
// Tar segment directory to compress file
File localSegmentDir = new File(localOutputTempDir, segmentName);
- String segmentTarFileName = URIUtils.encode(segmentName +
Constants.TAR_GZ_FILE_EXT);
+ String segmentTarFileName = URIUtils.encode(
+ segmentName +
Constants.SEGMENT_TAR_FILE_EXTENSIONS.get(_spec.getSegmentTarCompressionCodec()));
File localSegmentTarFile = new File(localOutputTempDir,
segmentTarFileName);
LOGGER.info("Tarring segment from: {} to: {}", localSegmentDir,
localSegmentTarFile);
TarCompressionUtils.createCompressedTarFile(localSegmentDir,
localSegmentTarFile);
diff --git
a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-hadoop/src/main/java/org/apache/pinot/plugin/ingestion/batch/hadoop/HadoopSegmentMetadataPushJobRunner.java
b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-hadoop/src/main/java/org/apache/pinot/plugin/ingestion/batch/hadoop/HadoopSegmentMetadataPushJobRunner.java
index fed62025a31..216bcf6fae3 100644
---
a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-hadoop/src/main/java/org/apache/pinot/plugin/ingestion/batch/hadoop/HadoopSegmentMetadataPushJobRunner.java
+++
b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-hadoop/src/main/java/org/apache/pinot/plugin/ingestion/batch/hadoop/HadoopSegmentMetadataPushJobRunner.java
@@ -23,7 +23,6 @@ import java.io.IOException;
import java.io.Serializable;
import java.net.URI;
import java.net.URISyntaxException;
-import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import org.apache.pinot.segment.local.utils.SegmentPushUtils;
@@ -31,7 +30,6 @@ import org.apache.pinot.spi.env.PinotConfiguration;
import org.apache.pinot.spi.filesystem.PinotFS;
import org.apache.pinot.spi.filesystem.PinotFSFactory;
import org.apache.pinot.spi.ingestion.batch.runner.IngestionJobRunner;
-import org.apache.pinot.spi.ingestion.batch.spec.Constants;
import org.apache.pinot.spi.ingestion.batch.spec.PinotFSSpec;
import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
@@ -79,22 +77,14 @@ public class HadoopSegmentMetadataPushJobRunner implements
IngestionJobRunner, S
throw new RuntimeException("Unable to list all files under outputDirURI
- '" + outputDirURI + "'");
}
- List<String> segmentsToPush = new ArrayList<>();
- for (String file : files) {
- if (file.endsWith(Constants.TAR_GZ_FILE_EXT)) {
- segmentsToPush.add(file);
- }
- }
-
- int pushParallelism = _spec.getPushJobSpec().getPushParallelism();
- if (pushParallelism < 1) {
- pushParallelism = segmentsToPush.size();
+ Map<String, String> segmentUriToTarPathMap =
+ SegmentPushUtils.getSegmentUriToTarPathMap(outputDirURI,
_spec.getPushJobSpec(), files);
+ if (segmentUriToTarPathMap.isEmpty()) {
+ return;
}
// Push from driver
try {
- Map<String, String> segmentUriToTarPathMap = SegmentPushUtils
- .getSegmentUriToTarPathMap(outputDirURI, _spec.getPushJobSpec(),
segmentsToPush.toArray(new String[0]));
SegmentPushUtils
.sendSegmentUriAndMetadata(_spec,
PinotFSFactory.create(outputDirURI.getScheme()), segmentUriToTarPathMap);
} catch (Exception e) {
diff --git
a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-hadoop/src/main/java/org/apache/pinot/plugin/ingestion/batch/hadoop/HadoopSegmentTarPushJobRunner.java
b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-hadoop/src/main/java/org/apache/pinot/plugin/ingestion/batch/hadoop/HadoopSegmentTarPushJobRunner.java
index 2dd6d39dd23..6771db1a3e8 100644
---
a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-hadoop/src/main/java/org/apache/pinot/plugin/ingestion/batch/hadoop/HadoopSegmentTarPushJobRunner.java
+++
b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-hadoop/src/main/java/org/apache/pinot/plugin/ingestion/batch/hadoop/HadoopSegmentTarPushJobRunner.java
@@ -30,7 +30,6 @@ import org.apache.pinot.spi.env.PinotConfiguration;
import org.apache.pinot.spi.filesystem.PinotFS;
import org.apache.pinot.spi.filesystem.PinotFSFactory;
import org.apache.pinot.spi.ingestion.batch.runner.IngestionJobRunner;
-import org.apache.pinot.spi.ingestion.batch.spec.Constants;
import org.apache.pinot.spi.ingestion.batch.spec.PinotFSSpec;
import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
import org.apache.pinot.spi.utils.retry.AttemptsExceededException;
@@ -80,17 +79,12 @@ public class HadoopSegmentTarPushJobRunner implements
IngestionJobRunner, Serial
throw new RuntimeException("Unable to list all files under outputDirURI
- '" + outputDirURI + "'");
}
- List<String> segmentsToPush = new ArrayList<>();
- for (String file : files) {
- if (file.endsWith(Constants.TAR_GZ_FILE_EXT)) {
- segmentsToPush.add(file);
- }
+ List<String> segmentsToPush = new ArrayList<>(
+ SegmentPushUtils.getSegmentUriToTarPathMap(outputDirURI,
_spec.getPushJobSpec(), files).values());
+ if (segmentsToPush.isEmpty()) {
+ return;
}
- int pushParallelism = _spec.getPushJobSpec().getPushParallelism();
- if (pushParallelism < 1) {
- pushParallelism = segmentsToPush.size();
- }
// Push from driver
try {
SegmentPushUtils.pushSegments(_spec, outputDirFS, segmentsToPush);
diff --git
a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-hadoop/src/main/java/org/apache/pinot/plugin/ingestion/batch/hadoop/HadoopSegmentUriPushJobRunner.java
b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-hadoop/src/main/java/org/apache/pinot/plugin/ingestion/batch/hadoop/HadoopSegmentUriPushJobRunner.java
index 0f253386f94..0352859145e 100644
---
a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-hadoop/src/main/java/org/apache/pinot/plugin/ingestion/batch/hadoop/HadoopSegmentUriPushJobRunner.java
+++
b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-hadoop/src/main/java/org/apache/pinot/plugin/ingestion/batch/hadoop/HadoopSegmentUriPushJobRunner.java
@@ -30,7 +30,6 @@ import org.apache.pinot.spi.env.PinotConfiguration;
import org.apache.pinot.spi.filesystem.PinotFS;
import org.apache.pinot.spi.filesystem.PinotFSFactory;
import org.apache.pinot.spi.ingestion.batch.runner.IngestionJobRunner;
-import org.apache.pinot.spi.ingestion.batch.spec.Constants;
import org.apache.pinot.spi.ingestion.batch.spec.PinotFSSpec;
import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
import org.apache.pinot.spi.utils.retry.AttemptsExceededException;
@@ -82,16 +81,12 @@ public class HadoopSegmentUriPushJobRunner implements
IngestionJobRunner, Serial
} catch (IOException e) {
throw new RuntimeException("Unable to list all files under outputDirURI
- '" + outputDirURI + "'");
}
- List<String> segmentUris = new ArrayList<>();
- for (String file : files) {
- URI uri = URI.create(file);
- if (uri.getPath().endsWith(Constants.TAR_GZ_FILE_EXT)) {
- URI updatedURI = SegmentPushUtils
- .generateSegmentTarURI(outputDirURI, uri,
_spec.getPushJobSpec().getSegmentUriPrefix(),
- _spec.getPushJobSpec().getSegmentUriSuffix());
- segmentUris.add(updatedURI.toString());
- }
+ List<String> segmentUris = new ArrayList<>(
+ SegmentPushUtils.getSegmentUriToTarPathMap(outputDirURI,
_spec.getPushJobSpec(), files).keySet());
+ if (segmentUris.isEmpty()) {
+ return;
}
+
// Push from driver
try {
SegmentPushUtils.sendSegmentUris(_spec, segmentUris);
diff --git
a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-3/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark3/SparkSegmentGenerationJobRunner.java
b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-3/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark3/SparkSegmentGenerationJobRunner.java
index 2d16805798d..8415302fd2d 100644
---
a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-3/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark3/SparkSegmentGenerationJobRunner.java
+++
b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-3/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark3/SparkSegmentGenerationJobRunner.java
@@ -293,7 +293,8 @@ public class SparkSegmentGenerationJobRunner implements
IngestionJobRunner, Seri
// Tar segment directory to compress file
File localSegmentDir = new File(localOutputTempDir, segmentName);
- String segmentTarFileName = URIUtils.encode(segmentName +
Constants.TAR_GZ_FILE_EXT);
+ String segmentTarFileName = URIUtils.encode(
+ segmentName +
Constants.SEGMENT_TAR_FILE_EXTENSIONS.get(_spec.getSegmentTarCompressionCodec()));
File localSegmentTarFile = new File(localOutputTempDir,
segmentTarFileName);
LOGGER.info("Tarring segment from: {} to: {}", localSegmentDir,
localSegmentTarFile);
TarCompressionUtils.createCompressedTarFile(localSegmentDir,
localSegmentTarFile);
diff --git
a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-3/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark3/SparkSegmentMetadataPushJobRunner.java
b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-3/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark3/SparkSegmentMetadataPushJobRunner.java
index 644b8f810e0..a34a1e66cd0 100644
---
a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-3/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark3/SparkSegmentMetadataPushJobRunner.java
+++
b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-3/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark3/SparkSegmentMetadataPushJobRunner.java
@@ -18,19 +18,16 @@
*/
package org.apache.pinot.plugin.ingestion.batch.spark3;
-import com.google.common.base.Preconditions;
import java.io.File;
import java.io.IOException;
import java.io.Serializable;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.ArrayList;
-import java.util.Arrays;
import java.util.Collection;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
-import java.util.stream.Collectors;
import org.apache.pinot.common.segment.generation.SegmentGenerationUtils;
import org.apache.pinot.segment.local.utils.ConsistentDataPushUtils;
import org.apache.pinot.segment.local.utils.SegmentPushUtils;
@@ -39,7 +36,6 @@ import org.apache.pinot.spi.env.PinotConfiguration;
import org.apache.pinot.spi.filesystem.PinotFS;
import org.apache.pinot.spi.filesystem.PinotFSFactory;
import org.apache.pinot.spi.ingestion.batch.runner.IngestionJobRunner;
-import org.apache.pinot.spi.ingestion.batch.spec.Constants;
import org.apache.pinot.spi.ingestion.batch.spec.PinotClusterSpec;
import org.apache.pinot.spi.ingestion.batch.spec.PinotFSSpec;
import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
@@ -107,6 +103,9 @@ public class SparkSegmentMetadataPushJobRunner implements
IngestionJobRunner, Se
// Collect files to process
List<String> segmentsToPush = getSegmentsToPush(outputDirFS, outputDirURI);
+ if (segmentsToPush.isEmpty()) {
+ return;
+ }
// Ensure tableConfigURI is set if missing
setupTableConfigURI();
@@ -239,7 +238,8 @@ public class SparkSegmentMetadataPushJobRunner implements
IngestionJobRunner, Se
private List<String> getSegmentsToPush(PinotFS outputDirFS, URI
outputDirURI) {
String[] files = listFiles(outputDirFS, outputDirURI);
- return Arrays.stream(files).filter(file ->
file.endsWith(Constants.TAR_GZ_FILE_EXT)).collect(Collectors.toList());
+ return new ArrayList<>(
+ SegmentPushUtils.getSegmentUriToTarPathMap(outputDirURI,
_spec.getPushJobSpec(), files).values());
}
private void setupFileSystems() {
@@ -308,11 +308,7 @@ public class SparkSegmentMetadataPushJobRunner implements
IngestionJobRunner, Se
Collection<String> tarFilePaths = segmentsUriToTarPathMap.values();
List<String> segmentNames = new ArrayList<>(tarFilePaths.size());
for (String tarFilePath : tarFilePaths) {
- File tarFile = new File(tarFilePath);
- String fileName = tarFile.getName();
-
Preconditions.checkArgument(fileName.endsWith(Constants.TAR_GZ_FILE_EXT));
- String segmentName = fileName.substring(0, fileName.length() -
Constants.TAR_GZ_FILE_EXT.length());
- segmentNames.add(segmentName);
+
segmentNames.add(SegmentPushUtils.getSegmentNameFromTarFile(tarFilePath));
}
return segmentNames;
}
diff --git
a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-base/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark/common/BaseSparkSegmentMetadataPushJobRunner.java
b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-base/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark/common/BaseSparkSegmentMetadataPushJobRunner.java
index 049a0eb3b3b..cf98c4447f8 100644
---
a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-base/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark/common/BaseSparkSegmentMetadataPushJobRunner.java
+++
b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-base/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark/common/BaseSparkSegmentMetadataPushJobRunner.java
@@ -30,7 +30,6 @@ import org.apache.pinot.spi.env.PinotConfiguration;
import org.apache.pinot.spi.filesystem.PinotFS;
import org.apache.pinot.spi.filesystem.PinotFSFactory;
import org.apache.pinot.spi.ingestion.batch.runner.IngestionJobRunner;
-import org.apache.pinot.spi.ingestion.batch.spec.Constants;
import org.apache.pinot.spi.ingestion.batch.spec.PinotFSSpec;
import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
import org.apache.pinot.spi.utils.retry.AttemptsExceededException;
@@ -79,11 +78,10 @@ public abstract class BaseSparkSegmentMetadataPushJobRunner
implements Ingestion
throw new RuntimeException("Unable to list all files under outputDirURI
- '" + outputDirURI + "'");
}
- List<String> segmentsToPush = new ArrayList<>();
- for (String file : files) {
- if (file.endsWith(Constants.TAR_GZ_FILE_EXT)) {
- segmentsToPush.add(file);
- }
+ List<String> segmentsToPush = new ArrayList<>(
+ SegmentPushUtils.getSegmentUriToTarPathMap(outputDirURI,
_spec.getPushJobSpec(), files).values());
+ if (segmentsToPush.isEmpty()) {
+ return;
}
int pushParallelism = _spec.getPushJobSpec().getPushParallelism();
diff --git
a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-base/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark/common/BaseSparkSegmentTarPushJobRunner.java
b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-base/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark/common/BaseSparkSegmentTarPushJobRunner.java
index f04c4d6fd17..b9f1e4480ef 100644
---
a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-base/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark/common/BaseSparkSegmentTarPushJobRunner.java
+++
b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-base/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark/common/BaseSparkSegmentTarPushJobRunner.java
@@ -30,7 +30,6 @@ import org.apache.pinot.spi.env.PinotConfiguration;
import org.apache.pinot.spi.filesystem.PinotFS;
import org.apache.pinot.spi.filesystem.PinotFSFactory;
import org.apache.pinot.spi.ingestion.batch.runner.IngestionJobRunner;
-import org.apache.pinot.spi.ingestion.batch.spec.Constants;
import org.apache.pinot.spi.ingestion.batch.spec.PinotFSSpec;
import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
import org.apache.pinot.spi.utils.retry.AttemptsExceededException;
@@ -79,11 +78,10 @@ public abstract class BaseSparkSegmentTarPushJobRunner
implements IngestionJobRu
throw new RuntimeException("Unable to list all files under outputDirURI
- '" + outputDirURI + "'");
}
- List<String> segmentsToPush = new ArrayList<>();
- for (String file : files) {
- if (file.endsWith(Constants.TAR_GZ_FILE_EXT)) {
- segmentsToPush.add(file);
- }
+ List<String> segmentsToPush = new ArrayList<>(
+ SegmentPushUtils.getSegmentUriToTarPathMap(outputDirURI,
_spec.getPushJobSpec(), files).values());
+ if (segmentsToPush.isEmpty()) {
+ return;
}
int pushParallelism = _spec.getPushJobSpec().getPushParallelism();
diff --git
a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-base/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark/common/BaseSparkSegmentUriPushJobRunner.java
b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-base/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark/common/BaseSparkSegmentUriPushJobRunner.java
index 883eb23f56f..11c88975bbc 100644
---
a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-base/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark/common/BaseSparkSegmentUriPushJobRunner.java
+++
b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-base/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark/common/BaseSparkSegmentUriPushJobRunner.java
@@ -30,7 +30,6 @@ import org.apache.pinot.spi.env.PinotConfiguration;
import org.apache.pinot.spi.filesystem.PinotFS;
import org.apache.pinot.spi.filesystem.PinotFSFactory;
import org.apache.pinot.spi.ingestion.batch.runner.IngestionJobRunner;
-import org.apache.pinot.spi.ingestion.batch.spec.Constants;
import org.apache.pinot.spi.ingestion.batch.spec.PinotFSSpec;
import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
import org.apache.pinot.spi.utils.retry.AttemptsExceededException;
@@ -83,15 +82,10 @@ public abstract class BaseSparkSegmentUriPushJobRunner
implements IngestionJobRu
} catch (IOException e) {
throw new RuntimeException("Unable to list all files under outputDirURI
- '" + outputDirURI + "'");
}
- List<String> segmentUris = new ArrayList<>();
- for (String file : files) {
- URI uri = URI.create(file);
- if (uri.getPath().endsWith(Constants.TAR_GZ_FILE_EXT)) {
- URI updatedURI =
- SegmentPushUtils.generateSegmentTarURI(outputDirURI, uri,
_spec.getPushJobSpec().getSegmentUriPrefix(),
- _spec.getPushJobSpec().getSegmentUriSuffix());
- segmentUris.add(updatedURI.toString());
- }
+ List<String> segmentUris = new ArrayList<>(
+ SegmentPushUtils.getSegmentUriToTarPathMap(outputDirURI,
_spec.getPushJobSpec(), files).keySet());
+ if (segmentUris.isEmpty()) {
+ return;
}
int pushParallelism = _spec.getPushJobSpec().getPushParallelism();
diff --git
a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-base/src/test/java/org/apache/pinot/plugin/ingestion/batch/spark/common/BaseSparkSegmentMetadataPushJobRunnerTest.java
b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-base/src/test/java/org/apache/pinot/plugin/ingestion/batch/spark/common/BaseSparkSegmentMetadataPushJobRunnerTest.java
new file mode 100644
index 00000000000..39b6ec26962
--- /dev/null
+++
b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-base/src/test/java/org/apache/pinot/plugin/ingestion/batch/spark/common/BaseSparkSegmentMetadataPushJobRunnerTest.java
@@ -0,0 +1,125 @@
+/**
+ * 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 java.util.Map;
+import org.apache.commons.io.FileUtils;
+import org.apache.pinot.plugin.ingestion.batch.common.BaseSegmentPushJobRunner;
+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.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.fail;
+
+
+/// Exercises shared archive selection before dispatching Spark push jobs.
+public class BaseSparkSegmentMetadataPushJobRunnerTest {
+ @Test
+ public void testEmptyStandaloneSelectionSkipsConsistentPush()
+ throws Exception {
+ BaseSegmentPushJobRunner runner = new BaseSegmentPushJobRunner() {
+ @Override
+ public void initFileSys() {
+ _spec = new SegmentGenerationJobSpec();
+ _spec.setPushJobSpec(new PushJobSpec());
+ _spec.getPushJobSpec().setPushFileNamePattern("glob:**/*.tar.zst");
+ _outputDirURI = URI.create("file:/segments/");
+ _files = new String[]{"/segments/segment.tar.gz"};
+ _consistentPushEnabled = true;
+ }
+
+ @Override
+ public List<String> getSegmentsToReplace(Map<String, String> segments) {
+ throw new AssertionError("Empty selection must not start consistent
push");
+ }
+
+ @Override
+ public void uploadSegments(Map<String, String> segments) {
+ fail("Empty selection must not upload segments");
+ }
+ };
+ runner.run();
+ }
+
+ @Test
+ public void testPushPatternSelectsSegmentArchives()
+ throws Exception {
+ Path outputDir = Files.createTempDirectory("archive push-");
+ try {
+ Path current = Files.createDirectory(outputDir.resolve("current"));
+ Path selected = Files.createFile(current.resolve("segment.tar.zst"));
+ Files.createFile(current.resolve("segment.tar.gz"));
+
Files.createFile(Files.createDirectory(outputDir.resolve("stale")).resolve("segment.tar.zst"));
+ 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.zst");
+ spec.setPushJobSpec(pushSpec);
+ List<String> pushed = new ArrayList<>();
+ BaseSparkSegmentMetadataPushJobRunner runner = new
BaseSparkSegmentMetadataPushJobRunner(spec) {
+ @Override
+ public void parallelizeMetadataPushJob(List<String> segments,
List<PinotFSSpec> fileSystems,
+ int parallelism, URI output) {
+ assertFalse(segments.isEmpty());
+ pushed.addAll(segments);
+ }
+ };
+ BaseSparkSegmentTarPushJobRunner tarRunner = new
BaseSparkSegmentTarPushJobRunner(spec) {
+ @Override
+ public void parallelizeTarPushJob(List<PinotFSSpec> fileSystems,
List<String> segments,
+ int parallelism, URI output) {
+ assertFalse(segments.isEmpty());
+ pushed.addAll(segments);
+ }
+ };
+ BaseSparkSegmentUriPushJobRunner uriRunner = new
BaseSparkSegmentUriPushJobRunner(spec) {
+ @Override
+ public void parallelizeUriPushJob(List<PinotFSSpec> fileSystems,
List<String> segments, int parallelism) {
+ assertFalse(segments.isEmpty());
+
segments.stream().map(URI::create).map(URI::getPath).forEach(pushed::add);
+ }
+ };
+ for (IngestionJobRunner pushRunner : List.of(runner, tarRunner,
uriRunner)) {
+ pushed.clear();
+ pushSpec.setPushParallelism(2);
+ pushSpec.setPushFileNamePattern("glob:**/current/*.tar.zst");
+ pushRunner.run();
+ assertEquals(pushed, List.of(selected.toString()));
+ pushed.clear();
+ pushSpec.setPushParallelism(0);
+ pushSpec.setPushFileNamePattern("glob:**/missing/*");
+ pushRunner.run();
+ assertEquals(pushed, List.of());
+ }
+ } finally {
+ FileUtils.deleteDirectory(outputDir.toFile());
+ }
+ }
+}
diff --git
a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-standalone/src/main/java/org/apache/pinot/plugin/ingestion/batch/standalone/SegmentGenerationJobRunner.java
b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-standalone/src/main/java/org/apache/pinot/plugin/ingestion/batch/standalone/SegmentGenerationJobRunner.java
index dc6cc6cb730..0163a7da90a 100644
---
a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-standalone/src/main/java/org/apache/pinot/plugin/ingestion/batch/standalone/SegmentGenerationJobRunner.java
+++
b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-standalone/src/main/java/org/apache/pinot/plugin/ingestion/batch/standalone/SegmentGenerationJobRunner.java
@@ -42,7 +42,9 @@ import
org.apache.pinot.plugin.ingestion.batch.common.SegmentGenerationTaskRunne
import org.apache.pinot.segment.local.utils.ConsistentDataPushUtils;
import org.apache.pinot.spi.config.table.TableConfig;
import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.data.readers.RecordReaderUtils;
import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.filesystem.LocalPinotFS;
import org.apache.pinot.spi.filesystem.PinotFS;
import org.apache.pinot.spi.filesystem.PinotFSFactory;
import org.apache.pinot.spi.ingestion.batch.BatchConfigProperties;
@@ -230,9 +232,17 @@ public class SegmentGenerationJobRunner implements
IngestionJobRunner {
File localOutputTempDir = new File(localTempDir, "output");
FileUtils.forceMkdir(localOutputTempDir);
- //copy input path to local
- File localInputDataFile = createLocalInputDateFile(inputFileURI,
localInputTempDir);
- _inputDirFS.copyToLocalFile(inputFileURI, localInputDataFile);
+ // Gzip readers can unpack beside the input, so keep those writes inside
an owned staging directory.
+ File inputFile = new File(URIUtils.decode(inputFileURI.getRawPath()));
+ boolean borrowedInputFile =
PinotFSFactory.isFileSystemInstanceOf(_inputDirFS, LocalPinotFS.class)
+ && !RecordReaderUtils.isGZippedFile(inputFile);
+ File localInputDataFile;
+ if (borrowedInputFile) {
+ localInputDataFile = inputFile;
+ } else {
+ localInputDataFile = createLocalInputDateFile(inputFileURI,
localInputTempDir);
+ _inputDirFS.copyToLocalFile(inputFileURI, localInputDataFile);
+ }
//create task spec
SegmentGenerationTaskSpec taskSpec = new SegmentGenerationTaskSpec();
@@ -265,7 +275,8 @@ public class SegmentGenerationJobRunner implements
IngestionJobRunner {
String segmentName = taskRunner.run();
// Tar segment directory to compress file
localSegmentDir = new File(localOutputTempDir, segmentName);
- String segmentTarFileName = URIUtils.encode(segmentName +
Constants.TAR_GZ_FILE_EXT);
+ String segmentTarFileName = URIUtils.encode(
+ segmentName +
Constants.SEGMENT_TAR_FILE_EXTENSIONS.get(_spec.getSegmentTarCompressionCodec()));
localSegmentTarFile = new File(localOutputTempDir, segmentTarFileName);
LOGGER.info("Tarring segment from: {} to: {}", localSegmentDir,
localSegmentTarFile);
TarCompressionUtils.createCompressedTarFile(localSegmentDir,
localSegmentTarFile);
@@ -295,7 +306,9 @@ public class SegmentGenerationJobRunner implements
IngestionJobRunner {
_segmentCreationTaskCountDownLatch.countDown();
FileUtils.deleteQuietly(localSegmentDir);
FileUtils.deleteQuietly(localSegmentTarFile);
- FileUtils.deleteQuietly(localInputDataFile);
+ if (!borrowedInputFile) {
+ FileUtils.deleteQuietly(localInputDataFile);
+ }
}
});
}
diff --git
a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-standalone/src/test/java/org/apache/pinot/plugin/ingestion/batch/standalone/SegmentGenerationJobRunnerTest.java
b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-standalone/src/test/java/org/apache/pinot/plugin/ingestion/batch/standalone/SegmentGenerationJobRunnerTest.java
index 6bc9218e30f..873caa76fa8 100644
---
a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-standalone/src/test/java/org/apache/pinot/plugin/ingestion/batch/standalone/SegmentGenerationJobRunnerTest.java
+++
b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-standalone/src/test/java/org/apache/pinot/plugin/ingestion/batch/standalone/SegmentGenerationJobRunnerTest.java
@@ -22,13 +22,21 @@ import com.google.common.collect.Lists;
import java.io.File;
import java.io.FilenameFilter;
import java.io.IOException;
+import java.lang.reflect.Field;
+import java.net.URI;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
+import java.util.ArrayList;
import java.util.List;
+import java.util.Set;
+import java.util.zip.GZIPOutputStream;
+import javax.annotation.Nullable;
import org.apache.commons.io.FileUtils;
+import org.apache.pinot.common.utils.TarCompressionUtils;
import
org.apache.pinot.plugin.ingestion.batch.common.SegmentGenerationTaskRunner;
import org.apache.pinot.plugin.inputformat.csv.CSVRecordReader;
import org.apache.pinot.plugin.inputformat.csv.CSVRecordReaderConfig;
+import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl;
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;
@@ -36,6 +44,8 @@ import
org.apache.pinot.spi.config.table.ingestion.IngestionConfig;
import org.apache.pinot.spi.data.FieldSpec.DataType;
import org.apache.pinot.spi.data.Schema;
import org.apache.pinot.spi.data.Schema.SchemaBuilder;
+import org.apache.pinot.spi.data.readers.RecordReaderConfig;
+import org.apache.pinot.spi.data.readers.RecordReaderUtils;
import org.apache.pinot.spi.filesystem.LocalPinotFS;
import org.apache.pinot.spi.ingestion.batch.BatchConfigProperties;
import org.apache.pinot.spi.ingestion.batch.spec.ExecutionFrameworkSpec;
@@ -46,14 +56,134 @@ import
org.apache.pinot.spi.ingestion.batch.spec.SegmentNameGeneratorSpec;
import org.apache.pinot.spi.ingestion.batch.spec.TableSpec;
import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
import org.testng.Assert;
+import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertThrows;
import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.expectThrows;
import static org.testng.Assert.fail;
public class SegmentGenerationJobRunnerTest {
+ @DataProvider
+ public Object[][] localInputOwnershipCases() {
+ return new Object[][]{{false}, {true}};
+ }
+
+ @Test(dataProvider = "localInputOwnershipCases")
+ public void testLocalInputIsBorrowedAndPreserved(boolean failGeneration)
+ throws Exception {
+ File testDir = makeTestDir();
+ try {
+ File inputDir = new File(testDir, "input");
+ FileUtils.forceMkdir(inputDir);
+ File inputFile = new File(inputDir, "input with spaces.csv");
+ FileUtils.writeLines(inputFile, List.of("col1,col2", "value1,1",
failGeneration ? "value2,invalid" : "value2,2"));
+ byte[] originalBytes = Files.readAllBytes(inputFile.toPath());
+ File outputDir = new File(testDir, "output");
+ SegmentGenerationJobSpec spec = makeJobSpec(inputDir, outputDir,
makeSchemaFile(testDir, "mySchema"),
+ makeTableConfigFile(testDir, "mySchema"));
+ spec.setPinotFSSpecs(List.of());
+ SegmentGenerationJobRunner runner = new SegmentGenerationJobRunner(spec);
+ TrackingLocalPinotFS fileSystem = new TrackingLocalPinotFS();
+ Field inputDirFS =
SegmentGenerationJobRunner.class.getDeclaredField("_inputDirFS");
+ inputDirFS.setAccessible(true);
+ inputDirFS.set(runner, fileSystem);
+ if (failGeneration) {
+ RuntimeException failure = expectThrows(RuntimeException.class,
runner::run);
+ assertTrue(failure.getMessage().contains("Failed to generate Pinot
segment"));
+ } else {
+ runner.run();
+ File unpacked = new File(testDir, "unpacked");
+ TarCompressionUtils.untar(new File(outputDir,
"myTable_OFFLINE_0.tar.gz"), unpacked);
+ assertEquals(new SegmentMetadataImpl(new File(unpacked,
"myTable_OFFLINE_0")).getTotalDocs(), 2);
+ }
+ assertTrue(fileSystem._copiedInputs.isEmpty(), "Local inputs must be
borrowed without copying");
+ assertEquals(Files.readAllBytes(inputFile.toPath()), originalBytes);
+ } finally {
+ FileUtils.deleteDirectory(testDir);
+ }
+ }
+
+ @DataProvider
+ public Object[][] compressedInputOwnershipCases() {
+ return new Object[][]{{false, false}, {true, true}};
+ }
+
+ @Test(dataProvider = "compressedInputOwnershipCases")
+ public void testCompressedLocalInputPreservesSourceAndSibling(boolean
gzipSuffix, boolean failGeneration)
+ throws Exception {
+ File testDir = makeTestDir();
+ try {
+ File inputDir = new File(testDir, "input");
+ FileUtils.forceMkdir(inputDir);
+ String inputName = gzipSuffix ? "input.csv.gz" : "compressed-input";
+ File inputFile = new File(inputDir, inputName);
+ File sibling = new File(inputDir, gzipSuffix ? "input.csv" :
"compressed-input.csv");
+ byte[] sentinel = "preserve the existing source
sibling".getBytes(StandardCharsets.UTF_8);
+ Files.write(sibling.toPath(), sentinel);
+ String csv = "col1,col2\nvalue1,1\nvalue2," + (failGeneration ?
"invalid" : "2") + "\n";
+ try (GZIPOutputStream compressed = new
GZIPOutputStream(Files.newOutputStream(inputFile.toPath()))) {
+ compressed.write(csv.getBytes(StandardCharsets.UTF_8));
+ }
+ byte[] originalBytes = Files.readAllBytes(inputFile.toPath());
+ File outputDir = new File(testDir, "output");
+ SegmentGenerationJobSpec spec = makeJobSpec(inputDir, outputDir,
makeSchemaFile(testDir, "mySchema"),
+ makeTableConfigFile(testDir, "mySchema"));
+ spec.setIncludeFileNamePattern("glob:**/" + inputName);
+
spec.getRecordReaderSpec().setClassName(UnpackingCSVRecordReader.class.getName());
+ SegmentGenerationJobRunner runner = new SegmentGenerationJobRunner(spec);
+ if (failGeneration) {
+ expectThrows(RuntimeException.class, runner::run);
+ } else {
+ runner.run();
+ assertTrue(new File(outputDir, "myTable_OFFLINE_0.tar.gz").isFile());
+ }
+ assertEquals(Files.readAllBytes(inputFile.toPath()), originalBytes);
+ assertEquals(Files.readAllBytes(sibling.toPath()), sentinel, "Existing
source siblings must not be overwritten");
+ } finally {
+ FileUtils.deleteDirectory(testDir);
+ }
+ }
+
+ @DataProvider(name = "segmentArchiveCodecs")
+ public Object[][] segmentArchiveCodecs() {
+ return new Object[][]{{"gzip", ".tar.gz"}, {"zstd", ".tar.zst"}};
+ }
+
+ @Test(dataProvider = "segmentArchiveCodecs")
+ public void testSegmentGenerationArchiveCodecs(String codec, String
extension)
+ throws Exception {
+ File testDir = makeTestDir();
+ try {
+ File inputDir = new File(testDir, "input");
+ FileUtils.forceMkdir(inputDir);
+ FileUtils.writeLines(new File(inputDir, "input.csv"),
List.of("col1,col2", "value1,1", "value2,2"));
+ File outputDir = new File(testDir, "output");
+ SegmentGenerationJobSpec spec = makeJobSpec(inputDir, outputDir,
makeSchemaFile(testDir, "mySchema"),
+ makeTableConfigFile(testDir, "mySchema"));
+ spec.setSegmentTarCompressionCodec(codec);
+ new SegmentGenerationJobRunner(spec).run();
+
+ File archive = new File(outputDir, "myTable_OFFLINE_0" + extension);
+ assertTrue(archive.isFile());
+ File unpacked = new File(testDir, "unpacked");
+ TarCompressionUtils.untar(archive, unpacked);
+ assertEquals(new SegmentMetadataImpl(new File(unpacked,
"myTable_OFFLINE_0")).getTotalDocs(), 2);
+ } finally {
+ FileUtils.deleteDirectory(testDir);
+ }
+ }
+
+ @Test
+ public void testSegmentGenerationRejectsUnsupportedArchiveCodec() {
+ SegmentGenerationJobSpec spec = new SegmentGenerationJobSpec();
+ assertEquals(spec.getSegmentTarCompressionCodec(), "gzip");
+ assertThrows(IllegalArgumentException.class, () ->
spec.setSegmentTarCompressionCodec("bzip2"));
+ }
+
@Test
public void testSegmentGeneration() throws Exception {
// TODO use common resource definitions & code shared with Hadoop unit
test.
@@ -227,6 +357,27 @@ public class SegmentGenerationJobRunnerTest {
}
}
+ /// Exercises the same sibling-writing unpacker as the Parquet and ORC
readers without another plugin dependency.
+ public static class UnpackingCSVRecordReader extends CSVRecordReader {
+ @Override
+ public void init(File dataFile, @Nullable Set<String> fieldsToRead,
@Nullable RecordReaderConfig config)
+ throws IOException {
+ super.init(RecordReaderUtils.unpackIfRequired(dataFile, "csv"),
fieldsToRead, config);
+ }
+ }
+
+ /// Records copies to verify that plain local inputs are borrowed.
+ private static class TrackingLocalPinotFS extends LocalPinotFS {
+ private final List<File> _copiedInputs = new ArrayList<>();
+
+ @Override
+ public void copyToLocalFile(URI source, File destination)
+ throws Exception {
+ _copiedInputs.add(destination);
+ super.copyToLocalFile(source, destination);
+ }
+ }
+
private File makeTestDir() throws IOException {
File testDir =
Files.createTempDirectory("testSegmentGeneration-").toFile();
testDir.delete();
diff --git
a/pinot-plugins/pinot-segment-uploader/pinot-segment-uploader-default/src/main/java/org/apache/pinot/plugin/segmentuploader/SegmentUploaderDefault.java
b/pinot-plugins/pinot-segment-uploader/pinot-segment-uploader-default/src/main/java/org/apache/pinot/plugin/segmentuploader/SegmentUploaderDefault.java
index 4ec4c19c2a3..7d8a79f3601 100644
---
a/pinot-plugins/pinot-segment-uploader/pinot-segment-uploader-default/src/main/java/org/apache/pinot/plugin/segmentuploader/SegmentUploaderDefault.java
+++
b/pinot-plugins/pinot-segment-uploader/pinot-segment-uploader-default/src/main/java/org/apache/pinot/plugin/segmentuploader/SegmentUploaderDefault.java
@@ -99,7 +99,7 @@ public class SegmentUploaderDefault implements
SegmentUploader {
String[] filePaths = outputPinotFS.listFiles(segmentDir, true);
for (String filePath : filePaths) {
URI uri = URI.create(filePath);
- if (!outputPinotFS.isDirectory(uri) &&
filePath.endsWith(Constants.TAR_GZ_FILE_EXT)) {
+ if (!outputPinotFS.isDirectory(uri) &&
Constants.isSegmentTarFile(filePath)) {
segmentTarURIs.add(uri);
}
}
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SegmentPushUtils.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SegmentPushUtils.java
index e1785d0919f..5ac39f52693 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SegmentPushUtils.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/SegmentPushUtils.java
@@ -210,10 +210,7 @@ public class SegmentPushUtils implements Serializable {
try {
for (String tarFilePath : tarFilePaths) {
URI tarFileURI = URI.create(tarFilePath);
- File tarFile = new File(tarFilePath);
- String fileName = tarFile.getName();
-
Preconditions.checkArgument(fileName.endsWith(Constants.TAR_GZ_FILE_EXT));
- String segmentName = fileName.substring(0, fileName.length() -
Constants.TAR_GZ_FILE_EXT.length());
+ String segmentName = getSegmentNameFromTarFile(tarFilePath);
for (PinotClusterSpec pinotClusterSpec : spec.getPinotClusterSpecs()) {
URI controllerURI;
try {
@@ -462,15 +459,27 @@ public class SegmentPushUtils implements Serializable {
return generateSegmentMetadataFile(fileSystem, URI.create(tarFilePath));
}
- /// Segments stored in the deep store do not have the .tar.gz extension.
- private static String getSegmentName(String tarFilePath) {
+ /// 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();
- String segmentName = fileName.endsWith(Constants.TAR_GZ_FILE_EXT)
- ? fileName.substring(0, fileName.length() -
Constants.TAR_GZ_FILE_EXT.length()) : fileName;
+ Preconditions.checkArgument(Constants.isSegmentTarFile(fileName),
"Unsupported segment archive: %s", fileName);
+ String extension =
Constants.SEGMENT_TAR_FILE_EXTENSIONS.values().stream().filter(fileName::endsWith)
+ .findFirst().orElseThrow();
+ String segmentName = fileName.substring(0, fileName.length() -
extension.length());
SegmentNameUtils.validatePartialOrFullSegmentName(segmentName);
return segmentName;
}
+ /// Segments stored in the deep store can omit the archive extension.
+ private static String getSegmentName(String tarFilePath) {
+ String fileName = new File(tarFilePath).getName();
+ if (Constants.isSegmentTarFile(fileName)) {
+ return getSegmentNameFromTarFile(tarFilePath);
+ }
+ SegmentNameUtils.validatePartialOrFullSegmentName(fileName);
+ return fileName;
+ }
+
private static String getSegmentNameFromMetadataFile(File
segmentMetadataFile) {
String fileName = segmentMetadataFile.getName();
Preconditions.checkArgument(fileName.endsWith(Constants.METADATA_TAR_GZ_FILE_EXT),
@@ -684,17 +693,22 @@ public class SegmentPushUtils implements Serializable {
}
}
- URI uri = URI.create(file);
+ URI uri = outputDirURI != null &&
"file".equals(outputDirURI.getScheme()) && new File(file).isAbsolute()
+ ? new File(file).toURI() : URI.create(file);
if (uri.getPath().endsWith(Constants.METADATA_TAR_GZ_FILE_EXT)) {
// Skip segment metadata tar gz files
continue;
}
- if (uri.getPath().endsWith(Constants.TAR_GZ_FILE_EXT)) {
+ if (Constants.isSegmentTarFile(uri.getPath())) {
URI updatedURI = SegmentPushUtils.generateSegmentTarURI(outputDirURI,
uri, pushSpec.getSegmentUriPrefix(),
pushSpec.getSegmentUriSuffix());
segmentUriToTarPathMap.put(updatedURI.toString(), file);
}
}
+ if (segmentUriToTarPathMap.isEmpty()) {
+ LOGGER.warn("No segment archives selected under outputDirURI: {} with
pushFileNamePattern: {}", outputDirURI,
+ pushSpec.getPushFileNamePattern());
+ }
return segmentUriToTarPathMap;
}
@@ -711,8 +725,9 @@ public class SegmentPushUtils implements Serializable {
File tarFile =
new File(FileUtils.getTempDirectory(), "segmentTar-" + uuid +
TarCompressionUtils.TAR_GZ_FILE_EXTENSION);
File segmentMetadataDir = new File(FileUtils.getTempDirectory(),
"segmentMetadataDir-" + uuid);
+ boolean localFileSystem =
PinotFSFactory.isFileSystemInstanceOf(fileSystem, LocalPinotFS.class);
try {
- if (fileSystem instanceof LocalPinotFS) {
+ if (localFileSystem) {
// For local file system, we don't need to copy the tar file.
tarFile = new File(URIUtils.decode(tarFileURI.getRawPath()));
} else {
@@ -743,7 +758,7 @@ public class SegmentPushUtils implements Serializable {
TarCompressionUtils.createCompressedTarFile(segmentMetadataDir,
segmentMetadataTarFile);
return segmentMetadataTarFile;
} finally {
- if (!(fileSystem instanceof LocalPinotFS)) {
+ if (!localFileSystem) {
// For local file system, we don't need to delete the tar file.
FileUtils.deleteQuietly(tarFile);
}
diff --git
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/SegmentPushUtilsTest.java
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/SegmentPushUtilsTest.java
index a1763c47252..80d2cf929d8 100644
---
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/SegmentPushUtilsTest.java
+++
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/SegmentPushUtilsTest.java
@@ -50,6 +50,7 @@ import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLParameters;
@@ -61,7 +62,9 @@ import org.apache.pinot.common.utils.FileUploadDownloadClient;
import org.apache.pinot.common.utils.TarCompressionUtils;
import org.apache.pinot.common.utils.http.HttpClientConfig;
import org.apache.pinot.common.utils.tls.TlsUtils;
+import org.apache.pinot.spi.filesystem.LocalPinotFS;
import org.apache.pinot.spi.filesystem.PinotFS;
+import org.apache.pinot.spi.filesystem.PinotFSFactory;
import org.apache.pinot.spi.ingestion.batch.spec.Constants;
import org.apache.pinot.spi.ingestion.batch.spec.PinotClusterSpec;
import org.apache.pinot.spi.ingestion.batch.spec.PushJobSpec;
@@ -72,14 +75,23 @@ import org.mockito.Mockito;
import org.testng.Assert;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNotNull;
import static org.testng.Assert.assertThrows;
import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.expectThrows;
public class SegmentPushUtilsTest {
@@ -401,6 +413,95 @@ public class SegmentPushUtilsTest {
"segment.name = " + TEST_SEGMENT_NAME);
}
+ @Test(dataProvider = "segmentTarExtensions")
+ public void testGenerateSegmentMetadataFileFromWrappedLocalArchive(String
extension)
+ throws Exception {
+ File archive = createSegmentMetadataSourceArchive(extension, true);
+ byte[] originalBytes = Files.readAllBytes(archive.toPath());
+ PinotFS fileSystem = spy(PinotFSFactory.create("file"));
+ assertFalse(fileSystem instanceof LocalPinotFS);
+ assertTrue(PinotFSFactory.isFileSystemInstanceOf(fileSystem,
LocalPinotFS.class));
+ File metadataTar = null;
+ try {
+ metadataTar = SegmentPushUtils.generateSegmentMetadataFile(fileSystem,
archive.toURI());
+ assertMetadataArchiveContents(metadataTar);
+ verify(fileSystem, never()).copyToLocalFile(any(URI.class),
any(File.class));
+ verify(fileSystem, never()).close();
+ assertEquals(Files.readAllBytes(archive.toPath()), originalBytes);
+ } finally {
+ FileUtils.deleteQuietly(metadataTar);
+ }
+ }
+
+ @Test
+ public void
testGenerateSegmentMetadataFileFailurePreservesWrappedLocalArchive()
+ throws Exception {
+ File archive = createSegmentMetadataSourceArchive(".tar.zst", false);
+ byte[] originalBytes = Files.readAllBytes(archive.toPath());
+ PinotFS fileSystem = spy(PinotFSFactory.create("file"));
+ IOException failure = expectThrows(IOException.class,
+ () -> SegmentPushUtils.generateSegmentMetadataFile(fileSystem,
archive.toURI()));
+ assertTrue(failure.getMessage().contains("creation.meta"));
+ verify(fileSystem, never()).copyToLocalFile(any(URI.class),
any(File.class));
+ verify(fileSystem, never()).close();
+ assertEquals(Files.readAllBytes(archive.toPath()), originalBytes);
+ }
+
+ @Test
+ public void
testGenerateSegmentMetadataFileCopiesCustomFileUriAndDeletesOnlyStaging()
+ throws Exception {
+ File archive = createSegmentMetadataSourceArchive(".tar.gz", true);
+ byte[] originalBytes = Files.readAllBytes(archive.toPath());
+ URI customSourceUri = URI.create("file:///virtual/remote-segment.tar.gz");
+ PinotFS fileSystem = mock(PinotFS.class);
+ AtomicReference<File> stagedArchive = new AtomicReference<>();
+ doAnswer(invocation -> {
+ URI source = invocation.getArgument(0);
+ assertEquals(source, customSourceUri);
+ File destination = invocation.getArgument(1);
+ assertFalse(destination.equals(archive));
+ FileUtils.copyFile(archive, destination);
+ stagedArchive.set(destination);
+ return null;
+ }).when(fileSystem).copyToLocalFile(any(URI.class), any(File.class));
+ File metadataTar = null;
+ try {
+ metadataTar = SegmentPushUtils.generateSegmentMetadataFile(fileSystem,
customSourceUri);
+ assertMetadataArchiveContents(metadataTar);
+ assertNotNull(stagedArchive.get());
+ assertFalse(stagedArchive.get().exists());
+ verify(fileSystem).copyToLocalFile(customSourceUri, stagedArchive.get());
+ verify(fileSystem, never()).close();
+ assertEquals(Files.readAllBytes(archive.toPath()), originalBytes);
+ } finally {
+ FileUtils.deleteQuietly(metadataTar);
+ }
+ }
+
+ private File createSegmentMetadataSourceArchive(String extension, boolean
includeCreationMeta)
+ throws IOException {
+ File segmentDir = new File(_tempDir, TEST_SEGMENT_NAME);
+ FileUtils.forceMkdir(segmentDir);
+ FileUtils.writeStringToFile(new File(segmentDir, "metadata.properties"),
"segment.name = " + TEST_SEGMENT_NAME,
+ StandardCharsets.UTF_8);
+ if (includeCreationMeta) {
+ FileUtils.writeStringToFile(new File(segmentDir, "creation.meta"),
"crc", StandardCharsets.UTF_8);
+ }
+ FileUtils.writeStringToFile(new File(segmentDir, "columns.psf"), "data",
StandardCharsets.UTF_8);
+ File archive = new File(_tempDir, "segment with spaces and percent%" +
extension);
+ TarCompressionUtils.createCompressedTarFile(segmentDir, archive);
+ return archive;
+ }
+
+ private void assertMetadataArchiveContents(File metadataTar)
+ throws IOException {
+ File untarred = TarCompressionUtils.untar(metadataTar, new File(_tempDir,
"untar")).get(0);
+ assertEquals(new HashSet<>(Arrays.asList(untarred.list())),
Set.of("metadata.properties", "creation.meta"));
+ assertEquals(FileUtils.readFileToString(new File(untarred,
"metadata.properties"), StandardCharsets.UTF_8),
+ "segment.name = " + TEST_SEGMENT_NAME);
+ assertEquals(FileUtils.readFileToString(new File(untarred,
"creation.meta"), StandardCharsets.UTF_8), "crc");
+ }
+
@Test
public void testSendSegmentUriAndMetadataWithLocalMetadataFiles()
throws Exception {
@@ -528,6 +629,45 @@ public class SegmentPushUtilsTest {
assertEquals(result.get(segmentFiles[3]), segmentFiles[3]);
}
+ @DataProvider(name = "segmentTarExtensions")
+ public Object[][] segmentTarExtensions() {
+ return new Object[][]{{".tar.gz"}, {".tar.lz4"}, {".tar.zst"}};
+ }
+
+ @Test(dataProvider = "segmentTarExtensions")
+ public void testSegmentArchiveScanning(String extension)
+ throws IOException {
+ String archiveUri = _tempDir.toURI().resolve(TEST_SEGMENT_NAME +
extension).toString();
+ assertEquals(SegmentPushUtils.getSegmentNameFromTarFile(archiveUri),
TEST_SEGMENT_NAME);
+ String[] files = {archiveUri, new File(_tempDir,
"testSegment.metadata.tar.gz").toURI().toString(),
+ new File(_tempDir, "testSegment.tar.bz2").toURI().toString()};
+ assertEquals(SegmentPushUtils.getSegmentUriToTarPathMap(_tempDir.toURI(),
new PushJobSpec(), files),
+ Map.of(archiveUri, archiveUri));
+
+ PushJobSpec pushSpec = new PushJobSpec();
+ pushSpec.setSegmentUriPrefix("s3://segments-bucket");
+ assertEquals(SegmentPushUtils.getSegmentUriToTarPathMap(null, pushSpec,
files),
+ Map.of("s3://segments-bucket" + URI.create(archiveUri).getRawPath(),
archiveUri));
+
+ File localArchive = new File(_tempDir, "a%3Ab" + extension);
+ assertTrue(localArchive.createNewFile());
+ String localArchivePath = localArchive.getAbsolutePath();
+ String localArchiveUri = _tempDir.toURI().resolve("a%253Ab" +
extension).toString();
+ String[] localFiles = {localArchivePath};
+ assertEquals(SegmentPushUtils.getSegmentUriToTarPathMap(_tempDir.toURI(),
new PushJobSpec(), localFiles),
+ Map.of(localArchiveUri, localArchivePath));
+ assertEquals(SegmentPushUtils.getSegmentUriToTarPathMap(_tempDir.toURI(),
pushSpec, localFiles),
+ Map.of("s3://segments-bucket" +
URI.create(localArchiveUri).getRawPath(), localArchivePath));
+ }
+
+ @Test
+ public void testSegmentArchiveNameRejectsSidecarsAndUnsupportedCodecs() {
+ assertThrows(IllegalArgumentException.class,
+ () ->
SegmentPushUtils.getSegmentNameFromTarFile("testSegment.metadata.tar.gz"));
+ assertThrows(IllegalArgumentException.class,
+ () ->
SegmentPushUtils.getSegmentNameFromTarFile("testSegment.tar.bz2"));
+ }
+
@Test
public void testGenerateSegmentMetadataURI()
throws URISyntaxException {
diff --git
a/pinot-spi/src/main/java/org/apache/pinot/spi/ingestion/batch/spec/Constants.java
b/pinot-spi/src/main/java/org/apache/pinot/spi/ingestion/batch/spec/Constants.java
index 7523fa6e48a..71294dcc7d0 100644
---
a/pinot-spi/src/main/java/org/apache/pinot/spi/ingestion/batch/spec/Constants.java
+++
b/pinot-spi/src/main/java/org/apache/pinot/spi/ingestion/batch/spec/Constants.java
@@ -18,11 +18,31 @@
*/
package org.apache.pinot.spi.ingestion.batch.spec;
+import java.util.Map;
+
public class Constants {
private Constants() {
}
/// By default Pinot segments are compressed in 'tar.gz' format then pushed
to controller.
public static final String TAR_GZ_FILE_EXT = ".tar.gz";
+ public static final String TAR_LZ4_FILE_EXT = ".tar.lz4";
+ public static final String TAR_ZST_FILE_EXT = ".tar.zst";
public static final String METADATA_TAR_GZ_FILE_EXT = ".metadata.tar.gz";
+
+ public static final String SEGMENT_COMPRESSION_CODEC_GZIP = "gzip";
+ public static final String SEGMENT_COMPRESSION_CODEC_LZ4 = "lz4";
+ public static final String SEGMENT_COMPRESSION_CODEC_ZSTD = "zstd";
+ public static final String DEFAULT_SEGMENT_COMPRESSION_CODEC =
SEGMENT_COMPRESSION_CODEC_GZIP;
+
+ /// Archive extensions understood by segment upload and metadata extraction,
including existing LZ4 archives.
+ public static final Map<String, String> SEGMENT_TAR_FILE_EXTENSIONS =
+ Map.of(SEGMENT_COMPRESSION_CODEC_GZIP, TAR_GZ_FILE_EXT,
SEGMENT_COMPRESSION_CODEC_LZ4, TAR_LZ4_FILE_EXT,
+ SEGMENT_COMPRESSION_CODEC_ZSTD, TAR_ZST_FILE_EXT);
+
+ /// Returns whether a path names a supported segment archive, excluding the
gzip metadata sidecars.
+ public static boolean isSegmentTarFile(String fileName) {
+ return !fileName.endsWith(METADATA_TAR_GZ_FILE_EXT)
+ &&
SEGMENT_TAR_FILE_EXTENSIONS.values().stream().anyMatch(fileName::endsWith);
+ }
}
diff --git
a/pinot-spi/src/main/java/org/apache/pinot/spi/ingestion/batch/spec/SegmentGenerationJobSpec.java
b/pinot-spi/src/main/java/org/apache/pinot/spi/ingestion/batch/spec/SegmentGenerationJobSpec.java
index 3992c2e4202..da6b22a2cb2 100644
---
a/pinot-spi/src/main/java/org/apache/pinot/spi/ingestion/batch/spec/SegmentGenerationJobSpec.java
+++
b/pinot-spi/src/main/java/org/apache/pinot/spi/ingestion/batch/spec/SegmentGenerationJobSpec.java
@@ -19,6 +19,7 @@
package org.apache.pinot.spi.ingestion.batch.spec;
import com.fasterxml.jackson.databind.node.ObjectNode;
+import com.google.common.base.Preconditions;
import java.io.Serializable;
import java.util.List;
import org.apache.pinot.spi.utils.JsonUtils;
@@ -58,6 +59,9 @@ public class SegmentGenerationJobSpec implements Serializable
{
/// Segment creation job parallelism.
private int _segmentCreationJobParallelism;
+ /// Codec for segment archives: `gzip` (default) or `zstd`. Metadata
sidecars always use gzip.
+ private String _segmentTarCompressionCodec =
Constants.DEFAULT_SEGMENT_COMPRESSION_CODEC;
+
/// Should overwrite output segments if existed.
private boolean _overwriteOutput;
@@ -228,6 +232,22 @@ public class SegmentGenerationJobSpec implements
Serializable {
_segmentCreationJobParallelism = segmentCreationJobParallelism;
}
+ public String getSegmentTarCompressionCodec() {
+ return _segmentTarCompressionCodec == null
+ ? Constants.DEFAULT_SEGMENT_COMPRESSION_CODEC
+ : _segmentTarCompressionCodec;
+ }
+
+ /// Selects the archive codec for standalone, Spark, and Hadoop segment
generation jobs.
+ /// Generation accepts `gzip` and `zstd`; archive discovery also consumes
existing LZ4 archives.
+ public void setSegmentTarCompressionCodec(String segmentTarCompressionCodec)
{
+
Preconditions.checkArgument(Constants.SEGMENT_COMPRESSION_CODEC_GZIP.equals(segmentTarCompressionCodec)
+ ||
Constants.SEGMENT_COMPRESSION_CODEC_ZSTD.equals(segmentTarCompressionCodec),
+ "Unsupported segmentTarCompressionCodec: %s; supported codecs: %s,
%s", segmentTarCompressionCodec,
+ Constants.SEGMENT_COMPRESSION_CODEC_GZIP,
Constants.SEGMENT_COMPRESSION_CODEC_ZSTD);
+ _segmentTarCompressionCodec = segmentTarCompressionCodec;
+ }
+
public void setCleanUpOutputDir(boolean cleanUpOutputDir) {
_cleanUpOutputDir = cleanUpOutputDir;
}
diff --git
a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/LaunchBackfillIngestionJobCommand.java
b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/LaunchBackfillIngestionJobCommand.java
index bf327153a89..6122e7ebdaf 100644
---
a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/LaunchBackfillIngestionJobCommand.java
+++
b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/LaunchBackfillIngestionJobCommand.java
@@ -23,6 +23,9 @@ import java.io.File;
import java.io.IOException;
import java.net.URI;
import java.net.URISyntaxException;
+import java.nio.file.FileSystems;
+import java.nio.file.PathMatcher;
+import java.nio.file.Paths;
import java.time.LocalDate;
import java.time.ZoneId;
import java.util.ArrayList;
@@ -41,6 +44,7 @@ import org.apache.pinot.common.utils.http.HttpClient;
import
org.apache.pinot.plugin.ingestion.batch.common.SegmentGenerationTaskRunner;
import org.apache.pinot.plugin.minion.tasks.SegmentConversionUtils;
import org.apache.pinot.segment.local.utils.ConsistentDataPushUtils;
+import org.apache.pinot.segment.local.utils.SegmentPushUtils;
import org.apache.pinot.segment.spi.partition.PartitionFunctionFactory;
import org.apache.pinot.segment.spi.partition.metadata.ColumnPartitionMetadata;
import org.apache.pinot.spi.auth.AuthProvider;
@@ -247,7 +251,7 @@ public class LaunchBackfillIngestionJobCommand extends
LaunchDataIngestionJobCom
return generatedSegments;
}
- private List<String> collectSegmentNames(SegmentGenerationJobSpec spec) {
+ List<String> collectSegmentNames(SegmentGenerationJobSpec spec) {
URI outputDirURI;
try {
outputDirURI = new URI(spec.getOutputDirURI());
@@ -259,21 +263,28 @@ public class LaunchBackfillIngestionJobCommand extends
LaunchDataIngestionJobCom
}
PinotFS outputDirFS = PinotFSFactory.create(outputDirURI.getScheme());
- // Get list of files to process
- String[] files;
+ PathMatcher pushFilePathMatcher =
+ spec.getPushJobSpec() != null &&
spec.getPushJobSpec().getPushFileNamePattern() != null
+ ?
FileSystems.getDefault().getPathMatcher(spec.getPushJobSpec().getPushFileNamePattern())
+ : null;
+ List<String> generatedSegments = new ArrayList<>();
try {
- files = outputDirFS.listFiles(outputDirURI, true);
+ for (String file : outputDirFS.listFiles(outputDirURI, true)) {
+ if (!Constants.isSegmentTarFile(file)
+ || pushFilePathMatcher != null &&
!pushFilePathMatcher.matches(Paths.get(file))) {
+ continue;
+ }
+ URI fileURI =
PinotFSFactory.LOCAL_PINOT_FS_SCHEME.equals(outputDirURI.getScheme())
+ && !file.startsWith(PinotFSFactory.LOCAL_PINOT_FS_SCHEME + ":")
+ ? new File(file).toURI()
+ : URI.create(file);
+ if (!outputDirFS.isDirectory(fileURI)) {
+
generatedSegments.add(SegmentPushUtils.getSegmentNameFromTarFile(file));
+ }
+ }
} catch (IOException e) {
throw new RuntimeException("Unable to list all files under outputDirURI
- '" + outputDirURI + "'");
}
-
- List<String> generatedSegments = new ArrayList<>();
- for (String file : files) {
- if (file.endsWith(Constants.TAR_GZ_FILE_EXT)) {
- generatedSegments.add(
- file.substring(file.lastIndexOf(File.separator) + 1,
file.indexOf(Constants.TAR_GZ_FILE_EXT)));
- }
- }
return generatedSegments;
}
diff --git
a/pinot-tools/src/test/java/org/apache/pinot/tools/admin/command/LaunchBackfillIngestionJobCommandTest.java
b/pinot-tools/src/test/java/org/apache/pinot/tools/admin/command/LaunchBackfillIngestionJobCommandTest.java
new file mode 100644
index 00000000000..db9be4152dd
--- /dev/null
+++
b/pinot-tools/src/test/java/org/apache/pinot/tools/admin/command/LaunchBackfillIngestionJobCommandTest.java
@@ -0,0 +1,59 @@
+/**
+ * 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.tools.admin.command;
+
+import java.io.File;
+import java.nio.file.Files;
+import java.util.List;
+import java.util.Set;
+import org.apache.commons.io.FileUtils;
+import org.apache.pinot.spi.ingestion.batch.spec.PushJobSpec;
+import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+
+
+/// Tests archive discovery for backfill replacement lineage.
+public class LaunchBackfillIngestionJobCommandTest {
+ @Test
+ public void testCollectSegmentNames()
+ throws Exception {
+ File outputDir = Files.createTempDirectory("backfill segments ").toFile();
+ try {
+ File nestedDir = new File(outputDir, "nested.tar.gz");
+ FileUtils.forceMkdir(nestedDir);
+ for (String file : List.of("gzip.tar.gz", "zstd.tar.zst", "lz4.tar.lz4",
"gzip.metadata.tar.gz", "input.csv")) {
+ FileUtils.touch(new File(nestedDir, file));
+ }
+ FileUtils.forceMkdir(new File(nestedDir, "directory.tar.zst"));
+ SegmentGenerationJobSpec spec = new SegmentGenerationJobSpec();
+ spec.setOutputDirURI(outputDir.toURI().toString());
+ LaunchBackfillIngestionJobCommand command = new
LaunchBackfillIngestionJobCommand();
+ assertEquals(Set.copyOf(command.collectSegmentNames(spec)),
Set.of("gzip", "zstd", "lz4"));
+
+ PushJobSpec pushSpec = new PushJobSpec();
+ pushSpec.setPushFileNamePattern("glob:**/*.tar.zst");
+ spec.setPushJobSpec(pushSpec);
+ assertEquals(command.collectSegmentNames(spec), List.of("zstd"));
+ } finally {
+ FileUtils.deleteDirectory(outputDir);
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
