This is an automated email from the ASF dual-hosted git repository.
Jackie-Jiang 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 b8e2ef27671 Build METADATA push metadata from the local segment
instead of re-downloading the staged tar (#19490)
b8e2ef27671 is described below
commit b8e2ef27671ab9a62b3e2b955af664015ddb42e0
Author: Kartik Khare <[email protected]>
AuthorDate: Tue Oct 6 00:12:22 2026 +0530
Build METADATA push metadata from the local segment instead of
re-downloading the staged tar (#19490)
---
.../BaseMultipleSegmentsConversionExecutor.java | 67 ++---
.../tasks/BaseSingleSegmentConversionExecutor.java | 20 +-
.../plugin/minion/tasks/BaseTaskExecutor.java | 19 --
...BaseMultipleSegmentsConversionExecutorTest.java | 177 +++++++++++-
.../BaseSingleSegmentConversionExecutorTest.java | 42 ++-
.../segment/local/utils/SegmentPushUtils.java | 307 ++++++++++++---------
.../segment/local/utils/SegmentPushUtilsTest.java | 112 ++++++++
7 files changed, 542 insertions(+), 202 deletions(-)
diff --git
a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseMultipleSegmentsConversionExecutor.java
b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseMultipleSegmentsConversionExecutor.java
index 052dd2df721..b5df8a4306e 100644
---
a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseMultipleSegmentsConversionExecutor.java
+++
b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseMultipleSegmentsConversionExecutor.java
@@ -36,6 +36,7 @@ import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.stream.Collectors;
+import javax.annotation.Nullable;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.io.FileUtils;
import org.apache.commons.lang3.StringUtils;
@@ -57,7 +58,6 @@ import org.apache.pinot.segment.spi.SegmentMetadata;
import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl;
import org.apache.pinot.spi.auth.AuthProvider;
import org.apache.pinot.spi.config.table.TableType;
-import org.apache.pinot.spi.filesystem.PinotFS;
import org.apache.pinot.spi.ingestion.batch.BatchConfigProperties;
import org.apache.pinot.spi.ingestion.batch.spec.PushJobSpec;
import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
@@ -214,6 +214,10 @@ public abstract class
BaseMultipleSegmentsConversionExecutor extends BaseTaskExe
int numOutputSegments = segmentConversionResults.size();
List<File> tarredSegmentFiles = new ArrayList<>(numOutputSegments);
+ // METADATA pushes send only metadata.properties and creation.meta,
built here from the local converted segment
+ // so the staged tar is never downloaded back. Null entries for TAR push.
+ List<File> segmentMetadataTarFiles = new ArrayList<>(numOutputSegments);
+ BatchConfigProperties.SegmentPushType pushType =
getSegmentPushType(taskConfigs);
int count = 1;
for (SegmentConversionResult segmentConversionResult :
segmentConversionResults) {
File convertedSegmentDir = segmentConversionResult.getFile();
@@ -226,6 +230,9 @@ public abstract class
BaseMultipleSegmentsConversionExecutor extends BaseTaskExe
segmentConversionResult.getSegmentName() +
TarCompressionUtils.TAR_GZ_FILE_EXTENSION);
TarCompressionUtils.createCompressedTarFile(convertedSegmentDir,
convertedSegmentTarFile);
tarredSegmentFiles.add(convertedSegmentTarFile);
+ segmentMetadataTarFiles.add(pushType ==
BatchConfigProperties.SegmentPushType.TAR ? null
+ :
SegmentPushUtils.generateSegmentMetadataFile(convertedSegmentDir,
convertedTarredSegmentDir,
+ segmentConversionResult.getSegmentName()));
if (!FileUtils.deleteQuietly(convertedSegmentDir)) {
LOGGER.warn("Failed to delete converted segment: {}",
convertedSegmentDir.getAbsolutePath());
}
@@ -249,13 +256,14 @@ public abstract class
BaseMultipleSegmentsConversionExecutor extends BaseTaskExe
SegmentUploadContext segmentUploadContext = new
SegmentUploadContext(pinotTaskConfig, segmentConversionResults);
preUploadSegments(segmentUploadContext);
- Map<String, String> segmentUriToTarPathMap = new HashMap<>();
+ Map<String, File> segmentUriToMetadataFileMap = new HashMap<>();
PushJobSpec pushJobSpec = getPushJobSpec(taskConfigs);
boolean batchSegmentUpload = pushJobSpec.isBatchSegmentUpload();
// Upload the tarred segments
for (int i = 0; i < numOutputSegments; i++) {
File convertedTarredSegmentFile = tarredSegmentFiles.get(i);
+ File segmentMetadataTarFile = segmentMetadataTarFiles.get(i);
SegmentConversionResult segmentConversionResult =
segmentConversionResults.get(i);
String resultSegmentName = segmentConversionResult.getSegmentName();
_eventObserver.notifyProgress(_pinotTaskConfig, "Uploading segment: "
+ resultSegmentName + " (" + (i + 1)
@@ -294,20 +302,24 @@ public abstract class
BaseMultipleSegmentsConversionExecutor extends BaseTaskExe
}
if (batchSegmentUpload) {
- updateSegmentUriToTarPathMap(taskConfigs, outputSegmentTarURI,
segmentConversionResult,
- segmentUriToTarPathMap, pushJobSpec);
+ for (String segmentUri : getSegmentUris(taskConfigs,
outputSegmentTarURI, segmentConversionResult,
+ pushJobSpec)) {
+ segmentUriToMetadataFileMap.put(segmentUri,
segmentMetadataTarFile);
+ }
} else {
String rawTableName =
TableNameBuilder.extractRawTableName(tableNameWithType);
- pushSegment(rawTableName, taskConfigs, outputSegmentTarURI,
httpHeaders, parameters, segmentConversionResult);
+ pushSegment(rawTableName, taskConfigs, outputSegmentTarURI,
httpHeaders, parameters, segmentConversionResult,
+ segmentMetadataTarFile);
if (!FileUtils.deleteQuietly(convertedTarredSegmentFile)) {
LOGGER.warn("Failed to delete tarred converted segment: {}",
convertedTarredSegmentFile.getAbsolutePath());
}
+ FileUtils.deleteQuietly(segmentMetadataTarFile);
}
}
if (batchSegmentUpload) {
try {
- pushSegments(tableNameWithType, taskConfigs, pinotTaskConfig,
segmentUriToTarPathMap, pushJobSpec,
+ pushSegments(tableNameWithType, taskConfigs, pinotTaskConfig,
segmentUriToMetadataFileMap, pushJobSpec,
authProvider, segmentConversionResults);
} finally {
for (File convertedTarredSegmentFile : tarredSegmentFiles) {
@@ -316,6 +328,7 @@ public abstract class
BaseMultipleSegmentsConversionExecutor extends BaseTaskExe
convertedTarredSegmentFile.getAbsolutePath());
}
}
+ segmentMetadataTarFiles.forEach(FileUtils::deleteQuietly);
}
}
@@ -410,21 +423,17 @@ public abstract class
BaseMultipleSegmentsConversionExecutor extends BaseTaskExe
}
}
+ /// Download URIs the controller should register for a staged segment tar
(prefix/suffix rules from the push spec).
@VisibleForTesting
- void updateSegmentUriToTarPathMap(Map<String, String> taskConfigs, URI
outputSegmentTarURI,
- SegmentConversionResult segmentConversionResult, Map<String, String>
segmentUriToTarPathMap,
- PushJobSpec pushJobSpec) {
+ Set<String> getSegmentUris(Map<String, String> taskConfigs, URI
outputSegmentTarURI,
+ SegmentConversionResult segmentConversionResult, PushJobSpec
pushJobSpec) {
String segmentName = segmentConversionResult.getSegmentName();
if
(!taskConfigs.containsKey(BatchConfigProperties.OUTPUT_SEGMENT_DIR_URI)) {
throw new RuntimeException("Output dir URI missing for metadata push
while processing segment: " + segmentName);
}
URI outputSegmentDirURI =
URI.create(taskConfigs.get(BatchConfigProperties.OUTPUT_SEGMENT_DIR_URI));
- Map<String, String> localSegmentUriToTarPathMap =
- SegmentPushUtils.getSegmentUriToTarPathMap(outputSegmentDirURI,
pushJobSpec,
- new String[]{outputSegmentTarURI.toString()});
- if (!localSegmentUriToTarPathMap.isEmpty()) {
- segmentUriToTarPathMap.putAll(localSegmentUriToTarPathMap);
- }
+ return SegmentPushUtils.getSegmentUriToTarPathMap(outputSegmentDirURI,
pushJobSpec,
+ new String[]{outputSegmentTarURI.toString()}).keySet();
}
@VisibleForTesting
@@ -436,23 +445,20 @@ public abstract class
BaseMultipleSegmentsConversionExecutor extends BaseTaskExe
}
private void pushSegments(String tableNameWithType, Map<String, String>
taskConfigs, PinotTaskConfig pinotTaskConfig,
- Map<String, String> segmentUriToTarPathMap, PushJobSpec pushJobSpec,
- AuthProvider authProvider, List<SegmentConversionResult>
segmentConversionResults)
+ Map<String, File> segmentUriToMetadataFileMap, PushJobSpec pushJobSpec,
AuthProvider authProvider,
+ List<SegmentConversionResult> segmentConversionResults)
throws Exception {
String tableName = TableNameBuilder.extractRawTableName(tableNameWithType);
SegmentGenerationJobSpec spec =
generateSegmentGenerationJobSpec(tableName, taskConfigs, pushJobSpec);
List<Header> headers = getSegmentPushCommonHeaders(pinotTaskConfig,
authProvider, segmentConversionResults);
List<NameValuePair> parameters =
getSegmentPushCommonParams(tableNameWithType);
-
- URI outputSegmentDirURI =
URI.create(taskConfigs.get(BatchConfigProperties.OUTPUT_SEGMENT_DIR_URI));
- try (PinotFS outputFileFS = MinionTaskUtils.getOutputPinotFS(taskConfigs,
outputSegmentDirURI)) {
- SegmentPushUtils.sendSegmentsUriAndMetadata(spec, outputFileFS,
segmentUriToTarPathMap, headers, parameters);
- }
+ SegmentPushUtils.sendSegmentsUriAndMetadata(spec,
segmentUriToMetadataFileMap, headers, parameters);
}
private void pushSegment(String tableName, Map<String, String> taskConfigs,
URI outputSegmentTarURI,
- List<Header> headers, List<NameValuePair> parameters,
SegmentConversionResult segmentConversionResult)
+ List<Header> headers, List<NameValuePair> parameters,
SegmentConversionResult segmentConversionResult,
+ @Nullable File segmentMetadataTarFile)
throws Exception {
BatchConfigProperties.SegmentPushType pushType =
getSegmentPushType(taskConfigs);
LOGGER.info("Trying to push Pinot segment with push mode {} from {}",
pushType, outputSegmentTarURI);
@@ -470,17 +476,12 @@ public abstract class
BaseMultipleSegmentsConversionExecutor extends BaseTaskExe
uploadURL, tarFile);
break;
case METADATA:
- if
(taskConfigs.containsKey(BatchConfigProperties.OUTPUT_SEGMENT_DIR_URI)) {
- URI outputSegmentDirURI =
URI.create(taskConfigs.get(BatchConfigProperties.OUTPUT_SEGMENT_DIR_URI));
- try (PinotFS outputFileFS =
MinionTaskUtils.getOutputPinotFS(taskConfigs, outputSegmentDirURI)) {
- Map<String, String> segmentUriToTarPathMap =
-
SegmentPushUtils.getSegmentUriToTarPathMap(outputSegmentDirURI, pushJobSpec,
- new String[]{outputSegmentTarURI.toString()});
- SegmentPushUtils.sendSegmentUriAndMetadata(spec, outputFileFS,
segmentUriToTarPathMap, headers, parameters);
- }
- } else {
- throw new RuntimeException("Output dir URI missing for metadata
push");
+ Map<String, File> segmentUriToMetadataFileMap = new HashMap<>();
+ for (String segmentUri : getSegmentUris(taskConfigs,
outputSegmentTarURI, segmentConversionResult,
+ pushJobSpec)) {
+ segmentUriToMetadataFileMap.put(segmentUri, segmentMetadataTarFile);
}
+ SegmentPushUtils.sendSegmentUriAndMetadata(spec,
segmentUriToMetadataFileMap, headers, parameters);
break;
default:
throw new UnsupportedOperationException("Unrecognized push mode - " +
pushType);
diff --git
a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseSingleSegmentConversionExecutor.java
b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseSingleSegmentConversionExecutor.java
index 72311626f07..eb348cef805 100644
---
a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseSingleSegmentConversionExecutor.java
+++
b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseSingleSegmentConversionExecutor.java
@@ -23,6 +23,7 @@ import com.google.common.base.Preconditions;
import java.io.File;
import java.net.URI;
import java.util.ArrayList;
+import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
@@ -147,12 +148,16 @@ public abstract class BaseSingleSegmentConversionExecutor
extends BaseTaskExecut
BatchConfigProperties.SegmentPushType pushType =
getSegmentPushType(configs);
boolean copyToDeepStore =
pushType == BatchConfigProperties.SegmentPushType.METADATA &&
isCopyToDeepStoreForMetadataPush();
- // Controller-copy pushes send only metadata.properties and
creation.meta, built from the local converted
- // segment. An unchanged segment (same CRC) skips the tar and the
staging and only re-registers its metadata.
+ // METADATA pushes send only metadata.properties and creation.meta,
built here from the local converted segment
+ // so the staged tar is never downloaded back. A controller-copy push of
an unchanged segment (same CRC) skips
+ // the tar and the staging and only re-registers its metadata.
File segmentMetadataTarFile = null;
boolean reuseExistingSegment = false;
+ if (pushType == BatchConfigProperties.SegmentPushType.METADATA) {
+ segmentMetadataTarFile =
+ SegmentPushUtils.generateSegmentMetadataFile(convertedSegmentDir,
tempDataDir, segmentName);
+ }
if (copyToDeepStore) {
- segmentMetadataTarFile =
createSegmentMetadataTarFile(convertedSegmentDir, tempDataDir, segmentName);
long convertedSegmentCrc = new
SegmentMetadataImpl(convertedSegmentDir).getCrc();
reuseExistingSegment =
convertedSegmentCrc == Long.parseLong(originalSegmentCrc) &&
StringUtils.isNotEmpty(downloadURL);
@@ -224,7 +229,7 @@ public abstract class BaseSingleSegmentConversionExecutor
extends BaseTaskExecut
uploadURL, downloadURL, convertedTarredSegmentFile,
segmentMetadataTarFile);
} else {
uploadSegmentWithMetadata(configs, pinotTaskConfig,
segmentConversionResult, authProvider, parameters,
- tableNameWithType, convertedTarredSegmentFile);
+ tableNameWithType, convertedTarredSegmentFile,
segmentMetadataTarFile);
}
break;
default:
@@ -262,7 +267,7 @@ public abstract class BaseSingleSegmentConversionExecutor
extends BaseTaskExecut
/// [BatchConfigProperties#PUSH_CONTROLLER_URI] in configs.
private void uploadSegmentWithMetadata(Map<String, String> configs,
PinotTaskConfig pinotTaskConfig,
SegmentConversionResult segmentConversionResult, AuthProvider
authProvider, List<NameValuePair> parameters,
- String tableNameWithType, File convertedTarredSegmentFile)
+ String tableNameWithType, File convertedTarredSegmentFile, File
segmentMetadataTarFile)
throws Exception {
if (!configs.containsKey(BatchConfigProperties.OUTPUT_SEGMENT_DIR_URI)) {
throw new RuntimeException("Output dir URI missing for metadata push.
Set "
@@ -282,9 +287,10 @@ public abstract class BaseSingleSegmentConversionExecutor
extends BaseTaskExecut
try (PinotFS outputFileFS = MinionTaskUtils.getOutputPinotFS(configs,
outputSegmentDirURI)) {
Map<String, String> segmentUriToTarPathMap =
SegmentPushUtils.getSegmentUriToTarPathMap(outputSegmentDirURI,
pushJobSpec, new String[]{outputSegmentTarURI.toString()});
+ Map<String, File> segmentUriToMetadataFileMap = new HashMap<>();
+ segmentUriToTarPathMap.keySet().forEach(uri ->
segmentUriToMetadataFileMap.put(uri, segmentMetadataTarFile));
try {
- SegmentPushUtils.sendSegmentUriAndMetadata(spec, outputFileFS,
segmentUriToTarPathMap, metadataHeaders,
- parameters);
+ SegmentPushUtils.sendSegmentUriAndMetadata(spec,
segmentUriToMetadataFileMap, metadataHeaders, parameters);
} catch (Exception e) {
// The tar was already staged to the output PinotFS before this
failure. If the task is retried, the next
// moveSegmentToOutputPinotFS() would fail with "Output file already
exists" (overwriteOutput defaults to
diff --git
a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseTaskExecutor.java
b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseTaskExecutor.java
index 61fe7b74601..dce1317d5f1 100644
---
a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseTaskExecutor.java
+++
b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/BaseTaskExecutor.java
@@ -20,7 +20,6 @@ package org.apache.pinot.plugin.minion.tasks;
import com.google.common.base.Preconditions;
import java.io.File;
-import java.io.IOException;
import java.net.URI;
import java.util.ArrayList;
import java.util.Collections;
@@ -45,14 +44,12 @@ import org.apache.pinot.core.minion.PinotTaskConfig;
import org.apache.pinot.core.util.PeerServerSegmentFinder;
import org.apache.pinot.minion.MinionContext;
import org.apache.pinot.minion.executor.PinotTaskExecutor;
-import org.apache.pinot.segment.spi.store.SegmentDirectoryPaths;
import org.apache.pinot.spi.auth.AuthProvider;
import org.apache.pinot.spi.config.table.TableConfig;
import org.apache.pinot.spi.config.table.TableType;
import org.apache.pinot.spi.data.Schema;
import org.apache.pinot.spi.filesystem.PinotFS;
import org.apache.pinot.spi.ingestion.batch.BatchConfigProperties;
-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;
import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
@@ -249,22 +246,6 @@ public abstract class BaseTaskExecutor implements
PinotTaskExecutor {
}
}
- /// Tars only metadata.properties and creation.meta from the local segment,
which is all a METADATA push sends.
- protected File createSegmentMetadataTarFile(File segmentDir, File outputDir,
String segmentName)
- throws IOException {
- File metadataDir = new File(outputDir, segmentName + "-metadata");
- File metadataTarFile = new File(outputDir, segmentName +
Constants.METADATA_TAR_GZ_FILE_EXT);
- try {
- FileUtils.forceMkdir(metadataDir);
-
FileUtils.copyFileToDirectory(SegmentDirectoryPaths.findMetadataFile(segmentDir),
metadataDir);
-
FileUtils.copyFileToDirectory(SegmentDirectoryPaths.findCreationMetaFile(segmentDir),
metadataDir);
- TarCompressionUtils.createCompressedTarFile(metadataDir,
metadataTarFile);
- return metadataTarFile;
- } finally {
- FileUtils.deleteQuietly(metadataDir);
- }
- }
-
/// Returns HTTP parameters common to segment upload and metadata push
(parallel push protection, table name, type).
protected List<NameValuePair> getSegmentPushCommonParams(String
tableNameWithType) {
List<NameValuePair> params = new ArrayList<>();
diff --git
a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/test/java/org/apache/pinot/plugin/minion/tasks/BaseMultipleSegmentsConversionExecutorTest.java
b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/test/java/org/apache/pinot/plugin/minion/tasks/BaseMultipleSegmentsConversionExecutorTest.java
index 6793e1ad175..7f1a2c15529 100644
---
a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/test/java/org/apache/pinot/plugin/minion/tasks/BaseMultipleSegmentsConversionExecutorTest.java
+++
b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/test/java/org/apache/pinot/plugin/minion/tasks/BaseMultipleSegmentsConversionExecutorTest.java
@@ -22,9 +22,12 @@ import java.io.File;
import java.io.IOException;
import java.net.URI;
import java.util.ArrayList;
+import java.util.Arrays;
import java.util.HashMap;
+import java.util.HashSet;
import java.util.List;
import java.util.Map;
+import java.util.Set;
import java.util.UUID;
import org.apache.commons.io.FileUtils;
import org.apache.hc.core5.http.Header;
@@ -32,17 +35,39 @@ import org.apache.hc.core5.http.NameValuePair;
import org.apache.hc.core5.http.message.BasicNameValuePair;
import org.apache.pinot.common.auth.NullAuthProvider;
import
org.apache.pinot.common.metadata.segment.SegmentZKMetadataCustomMapModifier;
+import org.apache.pinot.common.metrics.MinionMetrics;
import org.apache.pinot.common.utils.FileUploadDownloadClient;
+import org.apache.pinot.common.utils.TarCompressionUtils;
import org.apache.pinot.core.common.MinionConstants;
import org.apache.pinot.core.minion.PinotTaskConfig;
+import org.apache.pinot.minion.MinionConf;
import org.apache.pinot.minion.MinionContext;
+import org.apache.pinot.minion.event.MinionEventObservers;
+import
org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl;
+import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader;
+import org.apache.pinot.segment.local.utils.SegmentPushUtils;
+import org.apache.pinot.segment.spi.V1Constants;
+import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig;
+import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl;
import org.apache.pinot.spi.auth.AuthProvider;
+import org.apache.pinot.spi.config.instance.InstanceType;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.data.readers.GenericRow;
+import org.apache.pinot.spi.filesystem.PinotFS;
import org.apache.pinot.spi.ingestion.batch.BatchConfigProperties;
+import org.apache.pinot.spi.ingestion.batch.spec.Constants;
import org.apache.pinot.spi.ingestion.batch.spec.PushJobSpec;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.mockito.MockedStatic;
import org.mockito.Mockito;
+import org.mockito.stubbing.Answer;
import org.testng.Assert;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
@@ -237,7 +262,7 @@ public class BaseMultipleSegmentsConversionExecutorTest {
}
@Test
- public void testUpdateSegmentURIToTarPathMap()
+ public void testGetSegmentUris()
throws IOException {
// setup
File segmentDir = new File(_tempDir, "segments");
@@ -254,14 +279,154 @@ public class BaseMultipleSegmentsConversionExecutorTest {
PushJobSpec pushJobSpec = new PushJobSpec();
SegmentConversionResult conversionResult =
new
SegmentConversionResult.Builder().setSegmentName("mySegment").build();
- Map<String, String> segmentUriToTarPathMap = new HashMap<>();
// test
- _executor.updateSegmentUriToTarPathMap(taskConfigs, outputSegmentTarURI,
conversionResult, segmentUriToTarPathMap,
- pushJobSpec);
+ Set<String> segmentUris = _executor.getSegmentUris(taskConfigs,
outputSegmentTarURI, conversionResult, pushJobSpec);
// validate
- Assert.assertEquals(segmentUriToTarPathMap.size(), 1);
- Assert.assertTrue(segmentUriToTarPathMap.containsKey(_tempDir.toURI() +
"segments/segment.tar.gz"));
+ Assert.assertEquals(segmentUris, Set.of(_tempDir.toURI() +
"segments/segment.tar.gz"));
+ }
+
+ @DataProvider(name = "batchSegmentUpload")
+ public Object[][] batchSegmentUpload() {
+ return new Object[][]{{false}, {true}};
+ }
+
+ /// On METADATA push the metadata tar of every output segment is built from
the local converted segment and pushed
+ /// with the local-metadata overload, so the staged tar is never downloaded
back. Covers per-segment and batch pushes.
+ @Test(dataProvider = "batchSegmentUpload")
+ public void testMetadataPushUsesLocalMetadataTars(boolean batchSegmentUpload)
+ throws Exception {
+ File outputDir = new File(_tempDir, "output");
+ FileUtils.forceMkdir(outputDir);
+ File dataDir = new File(_tempDir, "minionData");
+ FileUtils.forceMkdir(dataDir);
+ File capturedDir = new File(_tempDir, "captured");
+ FileUtils.forceMkdir(capturedDir);
+ File inputSegmentDir = buildSegment("input_0");
+ List<String> outputSegmentNames = List.of("merged_0", "merged_1");
+ Map<String, File> outputSegmentDirs = new HashMap<>();
+ for (String name : outputSegmentNames) {
+ outputSegmentDirs.put(name, buildSegment(name));
+ }
+ String taskId = "Task_TestMultiSegmentTask_" + UUID.randomUUID();
+ MinionContext.getInstance().setDataDir(dataDir);
+ MinionMetrics.register(Mockito.mock(MinionMetrics.class));
+ MinionEventObservers.getInstance().addMinionEventObserver(taskId,
MinionTaskTestUtils.getMinionProgressObserver());
+
+ Map<String, String> configs = new HashMap<>();
+ configs.put(MinionConstants.TABLE_NAME_KEY, "myTable_OFFLINE");
+ configs.put(MinionConstants.SEGMENT_NAME_KEY, "input_0");
+ configs.put(MinionConstants.DOWNLOAD_URL_KEY, "http://unused/download";);
+ configs.put(MinionConstants.UPLOAD_URL_KEY, "http://unused/upload";);
+ configs.put(BatchConfigProperties.PUSH_MODE,
BatchConfigProperties.SegmentPushType.METADATA.name());
+ configs.put(BatchConfigProperties.OUTPUT_SEGMENT_DIR_URI,
outputDir.toURI().toString());
+ configs.put(BatchConfigProperties.BATCH_SEGMENT_UPLOAD,
String.valueOf(batchSegmentUpload));
+ configs.put("TASK_ID", taskId);
+
+ Map<String, File> capturedUriToMetadataFile = new HashMap<>();
+ try (MockedStatic<SegmentPushUtils> mocked =
Mockito.mockStatic(SegmentPushUtils.class,
+ Mockito.CALLS_REAL_METHODS)) {
+ Answer<Void> capture = invocation -> {
+ Map<String, File> uriToMetadataFile = invocation.getArgument(1);
+ for (Map.Entry<String, File> entry : uriToMetadataFile.entrySet()) {
+ File copy = new File(capturedDir, entry.getValue().getName());
+ FileUtils.copyFile(entry.getValue(), copy);
+ capturedUriToMetadataFile.put(entry.getKey(), copy);
+ }
+ return null;
+ };
+ mocked.when(() ->
SegmentPushUtils.sendSegmentUriAndMetadata(Mockito.any(), Mockito.anyMap(),
Mockito.anyList(),
+ Mockito.anyList())).thenAnswer(capture);
+ mocked.when(() ->
SegmentPushUtils.sendSegmentsUriAndMetadata(Mockito.any(), Mockito.anyMap(),
Mockito.anyList(),
+ Mockito.anyList())).thenAnswer(capture);
+
+ new TestMultipleSegmentsConversionExecutor(inputSegmentDir,
outputSegmentDirs).executeTask(
+ new PinotTaskConfig("TestMultiSegmentTask", configs));
+
+ // Never the PinotFS variants, which download the staged tar back to
extract the metadata
+ mocked.verify(() ->
SegmentPushUtils.sendSegmentUriAndMetadata(Mockito.any(),
Mockito.any(PinotFS.class),
+ Mockito.anyMap(), Mockito.anyList(), Mockito.anyList()),
Mockito.never());
+ mocked.verify(() ->
SegmentPushUtils.sendSegmentsUriAndMetadata(Mockito.any(),
Mockito.any(PinotFS.class),
+ Mockito.anyMap(), Mockito.anyList(), Mockito.anyList()),
Mockito.never());
+ } finally {
+ MinionEventObservers.getInstance().removeMinionEventObserver(taskId);
+ MinionContext.getInstance().setDataDir(null);
+ }
+
+ Assert.assertEquals(capturedUriToMetadataFile.size(), 2);
+ for (String name : outputSegmentNames) {
+ File stagedTar = new File(outputDir, name +
TarCompressionUtils.TAR_GZ_FILE_EXTENSION);
+ Assert.assertTrue(stagedTar.isFile(), "staged tar is the download URL
and must stay");
+ File metadataTar =
capturedUriToMetadataFile.get(stagedTar.toURI().toString());
+ Assert.assertNotNull(metadataTar, "metadata pushed for " +
stagedTar.toURI());
+ Assert.assertEquals(metadataTar.getName(), name +
Constants.METADATA_TAR_GZ_FILE_EXT);
+ File untarred = TarCompressionUtils.untar(metadataTar, new
File(_tempDir, "untar-" + name)).get(0);
+ Assert.assertEquals(new HashSet<>(Arrays.asList(untarred.list())),
+ Set.of(V1Constants.MetadataKeys.METADATA_FILE_NAME,
V1Constants.SEGMENT_CREATION_META));
+ Assert.assertEquals(new SegmentMetadataImpl(untarred).getName(), name);
+ }
+ }
+
+ private File buildSegment(String segmentName)
+ throws Exception {
+ TableConfig tableConfig = new
TableConfigBuilder(TableType.OFFLINE).setTableName("myTable").build();
+ Schema schema = new Schema.SchemaBuilder().addSingleValueDimension("d1",
FieldSpec.DataType.INT).build();
+ GenericRow row = new GenericRow();
+ row.putValue("d1", 1);
+ SegmentGeneratorConfig config = new SegmentGeneratorConfig(tableConfig,
schema);
+ config.setInstanceType(InstanceType.MINION);
+ config.setOutDir(new File(_tempDir, "built").getPath());
+ config.setSegmentName(segmentName);
+ SegmentIndexCreationDriverImpl driver = new
SegmentIndexCreationDriverImpl();
+ driver.init(config, new GenericRowRecordReader(List.of(row)));
+ driver.build();
+ return new File(new File(_tempDir, "built"), segmentName);
+ }
+
+ /// Stubs download, the segment-existence pre-check and conversion so
executeTask runs to the push step offline.
+ private static class TestMultipleSegmentsConversionExecutor extends
BaseMultipleSegmentsConversionExecutor {
+ private final File _inputSegmentDir;
+ private final Map<String, File> _outputSegmentDirs;
+
+ TestMultipleSegmentsConversionExecutor(File inputSegmentDir, Map<String,
File> outputSegmentDirs) {
+ super(new MinionConf());
+ _inputSegmentDir = inputSegmentDir;
+ _outputSegmentDirs = outputSegmentDirs;
+ }
+
+ @Override
+ protected void preProcess(PinotTaskConfig pinotTaskConfig) {
+ }
+
+ @Override
+ protected File downloadSegmentToLocalAndUntar(String tableNameWithType,
String segmentName, String deepstoreURL,
+ String taskType, File tempDataDir, String suffix)
+ throws Exception {
+ File indexDir = new File(tempDataDir, "input" + suffix);
+ FileUtils.copyDirectory(_inputSegmentDir, indexDir);
+ return indexDir;
+ }
+
+ @Override
+ protected List<SegmentConversionResult> convert(PinotTaskConfig
pinotTaskConfig, List<File> segmentDirs,
+ File workingDir)
+ throws Exception {
+ List<SegmentConversionResult> results = new ArrayList<>();
+ for (Map.Entry<String, File> entry : _outputSegmentDirs.entrySet()) {
+ File convertedDir = new File(workingDir, entry.getKey());
+ FileUtils.copyDirectory(entry.getValue(), convertedDir);
+ results.add(new SegmentConversionResult.Builder().setFile(convertedDir)
+
.setTableNameWithType(pinotTaskConfig.getConfigs().get(MinionConstants.TABLE_NAME_KEY))
+ .setSegmentName(entry.getKey()).build());
+ }
+ return results;
+ }
+
+ @Override
+ protected SegmentZKMetadataCustomMapModifier
getSegmentZKMetadataCustomMapModifier(PinotTaskConfig pinotTaskConfig,
+ SegmentConversionResult segmentConversionResult) {
+ return new
SegmentZKMetadataCustomMapModifier(SegmentZKMetadataCustomMapModifier.ModifyMode.UPDATE,
Map.of());
+ }
}
}
diff --git
a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/test/java/org/apache/pinot/plugin/minion/tasks/BaseSingleSegmentConversionExecutorTest.java
b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/test/java/org/apache/pinot/plugin/minion/tasks/BaseSingleSegmentConversionExecutorTest.java
index e677ebfcb10..7bb7a542cf9 100644
---
a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/test/java/org/apache/pinot/plugin/minion/tasks/BaseSingleSegmentConversionExecutorTest.java
+++
b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/test/java/org/apache/pinot/plugin/minion/tasks/BaseSingleSegmentConversionExecutorTest.java
@@ -165,8 +165,8 @@ public class BaseSingleSegmentConversionExecutorTest {
Mockito.mockStatic(SegmentPushUtils.class,
Mockito.CALLS_REAL_METHODS)) {
minionTaskUtils.when(() ->
MinionTaskUtils.getOutputPinotFS(Mockito.any(), Mockito.any()))
.thenReturn(mockOutputFS);
- segmentPushUtils.when(() ->
SegmentPushUtils.sendSegmentUriAndMetadata(Mockito.any(), Mockito.any(),
- Mockito.any(), Mockito.anyList(), Mockito.anyList()))
+ segmentPushUtils.when(() ->
SegmentPushUtils.sendSegmentUriAndMetadata(Mockito.any(), Mockito.anyMap(),
+ Mockito.anyList(), Mockito.anyList()))
.thenThrow(new RuntimeException("simulated metadata push failure"));
TestSingleSegmentConversionExecutor executor = new
TestSingleSegmentConversionExecutor(STALE_SEGMENT_CRC);
@@ -182,30 +182,40 @@ public class BaseSingleSegmentConversionExecutorTest {
}
- /// The default METADATA push still registers the staged tar's URI and never
takes the controller-copy path.
+ /// The default METADATA push still registers the staged tar's URI, with a
metadata tar built from the local
+ /// segment instead of the staged tar being downloaded back, and never takes
the controller-copy path.
@Test
public void testDefaultMetadataPushRegistersStagedUri()
throws Exception {
File outputDir = new File(TEMP_DIR, "output-default");
FileUtils.forceMkdir(outputDir);
+ File capturedMetadataTar = new File(TEMP_DIR,
"captured-default-metadata.tar.gz");
+ List<String> capturedUris = new ArrayList<>();
try (MockedStatic<SegmentPushUtils> segmentPushUtils =
Mockito.mockStatic(SegmentPushUtils.class,
Mockito.CALLS_REAL_METHODS);
MockedStatic<SegmentConversionUtils> conversionUtils =
Mockito.mockStatic(SegmentConversionUtils.class)) {
- segmentPushUtils.when(() ->
SegmentPushUtils.sendSegmentUriAndMetadata(Mockito.any(), Mockito.any(),
- Mockito.any(), Mockito.anyList(),
Mockito.anyList())).thenAnswer(invocation -> null);
+ segmentPushUtils.when(() ->
SegmentPushUtils.sendSegmentUriAndMetadata(Mockito.any(), Mockito.anyMap(),
+ Mockito.anyList(), Mockito.anyList())).thenAnswer(invocation -> {
+ Map<String, File> uriToMetadataFile = invocation.getArgument(1);
+ capturedUris.addAll(uriToMetadataFile.keySet());
+ FileUtils.copyFile(uriToMetadataFile.values().iterator().next(),
capturedMetadataTar);
+ return null;
+ });
new TestSingleSegmentConversionExecutor(STALE_SEGMENT_CRC).executeTask(
createMetadataPushTaskConfig(STALE_SEGMENT_CRC, outputDir));
- String stagedTarName = SEGMENT_NAME +
TarCompressionUtils.TAR_GZ_FILE_EXTENSION;
- segmentPushUtils.verify(() ->
SegmentPushUtils.sendSegmentUriAndMetadata(Mockito.any(), Mockito.any(),
- Mockito.argThat((Map<String, String> uriToTar) -> uriToTar.size() ==
1
- && uriToTar.values().iterator().next().endsWith(stagedTarName)),
- Mockito.anyList(), Mockito.anyList()));
+ // No PinotFS-based push, which is the variant that downloads the staged
tar back.
+ segmentPushUtils.verify(() ->
SegmentPushUtils.sendSegmentUriAndMetadata(Mockito.any(),
+ Mockito.any(PinotFS.class), Mockito.anyMap(), Mockito.anyList(),
Mockito.anyList()), Mockito.never());
conversionUtils.verifyNoInteractions();
}
+ File stagedTar = new File(outputDir, SEGMENT_NAME +
TarCompressionUtils.TAR_GZ_FILE_EXTENSION);
+ Assert.assertEquals(capturedUris.size(), 1);
+ Assert.assertEquals(new File(URI.create(capturedUris.get(0))), stagedTar);
+ assertMetadataTarDescribesSegment(capturedMetadataTar, new File(TEMP_DIR,
"untar-default-metadata"));
// The staged tar is the segment's download URL in this mode, so it stays.
- Assert.assertTrue(new File(outputDir, SEGMENT_NAME +
TarCompressionUtils.TAR_GZ_FILE_EXTENSION).isFile());
+ Assert.assertTrue(stagedTar.isFile());
}
/// Controller-copy push of a changed segment: task-unique staging name, TAR
guards plus copy flag, a metadata-only
@@ -250,9 +260,13 @@ public class BaseSingleSegmentConversionExecutorTest {
String downloadUri = headerValue(capturedHeaders,
FileUploadDownloadClient.CustomHeaders.DOWNLOAD_URI);
Assert.assertEquals(new File(URI.create(downloadUri)), stagedTar);
- // The metadata tar carries only the two files the controller reads, and
they describe the converted segment.
- File untarredMetadataDir = TarCompressionUtils.untar(capturedMetadataTar,
new File(TEMP_DIR, "untar-metadata"))
- .get(0);
+ assertMetadataTarDescribesSegment(capturedMetadataTar, new File(TEMP_DIR,
"untar-metadata"));
+ }
+
+ /// The metadata tar carries only the two files the controller reads, and
they describe the converted segment.
+ private void assertMetadataTarDescribesSegment(File metadataTar, File
untarDir)
+ throws Exception {
+ File untarredMetadataDir = TarCompressionUtils.untar(metadataTar,
untarDir).get(0);
Set<String> fileNames =
java.util.Arrays.stream(untarredMetadataDir.listFiles()).map(File::getName).collect(Collectors.toSet());
Assert.assertEquals(fileNames,
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 cd8bcc80a06..e1785d0919f 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
@@ -59,6 +59,7 @@ import org.apache.pinot.common.utils.tls.TlsUtils;
import org.apache.pinot.segment.local.constants.SegmentUploadConstants;
import org.apache.pinot.segment.spi.V1Constants;
import org.apache.pinot.segment.spi.creator.name.SegmentNameUtils;
+import org.apache.pinot.segment.spi.store.SegmentDirectoryPaths;
import org.apache.pinot.spi.auth.AuthProvider;
import org.apache.pinot.spi.config.table.TableType;
import org.apache.pinot.spi.env.PinotConfiguration;
@@ -342,87 +343,17 @@ public class SegmentPushUtils implements Serializable {
public static void sendSegmentUriAndMetadata(SegmentGenerationJobSpec spec,
PinotFS fileSystem,
Map<String, String> segmentUriToTarPathMap, List<Header> headers,
List<NameValuePair> parameters)
throws Exception {
- String tableName = spec.getTableSpec().getTableName();
LOGGER.info("Start pushing segment metadata: {} to locations: {} for table
{}", segmentUriToTarPathMap,
- Arrays.toString(spec.getPinotClusterSpecs()), tableName);
+ Arrays.toString(spec.getPinotClusterSpecs()),
spec.getTableSpec().getTableName());
FileUploadDownloadClient fileUploadDownloadClient =
getOrCreateFileUploadDownloadClient(spec);
- int socketTimeoutMs = getSocketTimeoutMs(spec);
try {
- for (String segmentUriPath : segmentUriToTarPathMap.keySet()) {
- String tarFilePath = segmentUriToTarPathMap.get(segmentUriPath);
- String fileName = new File(tarFilePath).getName();
- // segments stored in Pinot deep store do not have .tar.gz extension
- String segmentName = fileName.endsWith(Constants.TAR_GZ_FILE_EXT)
- ? fileName.substring(0, fileName.length() -
Constants.TAR_GZ_FILE_EXT.length()) : fileName;
- SegmentNameUtils.validatePartialOrFullSegmentName(segmentName);
- File segmentMetadataFile;
- // Check if there is a segment metadata tar gz file named
`segmentName.metadata.tar.gz`, already in the remote
- // directory. This is to avoid generating a new segment metadata tar
gz file every time we push a segment,
- // which requires downloading the entire segment tar gz file.
-
- URI metadataTarGzFilePath = generateSegmentMetadataURI(tarFilePath,
segmentName);
- LOGGER.info("Checking if metadata tar gz file {} exists",
metadataTarGzFilePath);
- if (spec.getPushJobSpec().isPreferMetadataTarGz() &&
fileSystem.exists(metadataTarGzFilePath)) {
- segmentMetadataFile = new File(FileUtils.getTempDirectory(),
- "segmentMetadata-" + UUID.randomUUID() +
TarCompressionUtils.TAR_GZ_FILE_EXTENSION);
- if (segmentMetadataFile.exists()) {
- FileUtils.forceDelete(segmentMetadataFile);
- }
- fileSystem.copyToLocalFile(metadataTarGzFilePath,
segmentMetadataFile);
- } else {
- segmentMetadataFile = generateSegmentMetadataFile(fileSystem,
URI.create(tarFilePath));
- }
+ for (Map.Entry<String, String> entry :
segmentUriToTarPathMap.entrySet()) {
+ String tarFilePath = entry.getValue();
+ String segmentName = getSegmentName(tarFilePath);
+ File segmentMetadataFile = getSegmentMetadataFile(spec, fileSystem,
tarFilePath, segmentName);
try {
- for (PinotClusterSpec pinotClusterSpec :
spec.getPinotClusterSpecs()) {
- URI controllerURI;
- try {
- controllerURI = new URI(pinotClusterSpec.getControllerURI());
- } catch (URISyntaxException e) {
- throw new RuntimeException("Got invalid controller uri - '" +
pinotClusterSpec.getControllerURI() + "'");
- }
- LOGGER.info("Pushing segment: {} to location: {} for table {}",
segmentName, controllerURI, tableName);
- int attempts = 1;
- if (spec.getPushJobSpec() != null &&
spec.getPushJobSpec().getPushAttempts() > 0) {
- attempts = spec.getPushJobSpec().getPushAttempts();
- }
- long retryWaitMs = 1000L;
- if (spec.getPushJobSpec() != null &&
spec.getPushJobSpec().getPushRetryIntervalMillis() > 0) {
- retryWaitMs = spec.getPushJobSpec().getPushRetryIntervalMillis();
- }
- RetryPolicies.exponentialBackoffRetryPolicy(attempts, retryWaitMs,
5).attempt(() -> {
- List<Header> reqHttpHeaders = new ArrayList<>(headers);
- try {
- reqHttpHeaders.add(
- new
BasicHeader(FileUploadDownloadClient.CustomHeaders.DOWNLOAD_URI,
segmentUriPath));
- reqHttpHeaders.add(new
BasicHeader(FileUploadDownloadClient.CustomHeaders.UPLOAD_TYPE,
-
FileUploadDownloadClient.FileUploadType.METADATA.toString()));
- if (spec.getPushJobSpec() != null) {
- reqHttpHeaders.add(new
BasicHeader(FileUploadDownloadClient.CustomHeaders.COPY_SEGMENT_TO_DEEP_STORE,
-
String.valueOf(spec.getPushJobSpec().getCopyToDeepStoreForMetadataPush())));
- }
-
- SimpleHttpResponse response =
fileUploadDownloadClient.uploadSegmentMetadata(
-
FileUploadDownloadClient.getUploadSegmentURI(controllerURI), segmentName,
- segmentMetadataFile, reqHttpHeaders, parameters,
socketTimeoutMs);
- LOGGER.info("Response for pushing table {} segment {} to
location {} - {}: {}", tableName, segmentName,
- controllerURI, response.getStatusCode(),
response.getResponse());
- return true;
- } catch (HttpErrorStatusException e) {
- int statusCode = e.getStatusCode();
- if (statusCode >= 500) {
- // Temporary exception
- LOGGER.warn("Caught temporary exception while pushing table:
{} segment: {} to {}, will retry",
- tableName, segmentName, controllerURI, e);
- return false;
- } else {
- // Permanent exception
- LOGGER.error("Caught permanent exception while pushing
table: {} segment: {} to {}, won't retry",
- tableName, segmentName, controllerURI, e);
- throw e;
- }
- }
- });
- }
+ pushSegmentMetadata(spec, fileUploadDownloadClient, entry.getKey(),
segmentName, segmentMetadataFile,
+ headers, parameters);
} finally {
FileUtils.deleteQuietly(segmentMetadataFile);
}
@@ -432,30 +363,171 @@ public class SegmentPushUtils implements Serializable {
}
}
- public static void sendSegmentsUriAndMetadata(SegmentGenerationJobSpec spec,
PinotFS fileSystem,
- Map<String, String> segmentUriToTarPathMap, List<Header> headers,
List<NameValuePair> parameters)
+ /// Same as above for segments whose metadata tars the caller already holds
locally (see
+ /// [#generateSegmentMetadataFile(File, File, String)]), so nothing is
downloaded back from the output filesystem.
+ /// Each file must be named `<segmentName>.metadata.tar.gz`, and the caller
owns the files.
+ public static void sendSegmentUriAndMetadata(SegmentGenerationJobSpec spec,
+ Map<String, File> segmentUriToMetadataFileMap, List<Header> headers,
List<NameValuePair> parameters)
throws Exception {
- String tableName = spec.getTableSpec().getTableName();
- ConcurrentHashMap<String, File> segmentMetadataFileMap = new
ConcurrentHashMap<>();
- ConcurrentLinkedQueue<String> segmentURIs = new ConcurrentLinkedQueue<>();
- Map<String, File> allSegmentsMetadataMap = new HashMap<>();
- File allSegmentsMetadataTarFile = null;
- int nThreads =
spec.getPushJobSpec().getSegmentMetadataGenerationParallelism();
+ LOGGER.info("Start pushing local segment metadata for: {} to locations: {}
for table {}",
+ segmentUriToMetadataFileMap.keySet(),
Arrays.toString(spec.getPinotClusterSpecs()),
+ spec.getTableSpec().getTableName());
FileUploadDownloadClient fileUploadDownloadClient =
getOrCreateFileUploadDownloadClient(spec);
+ try {
+ for (Map.Entry<String, File> entry :
segmentUriToMetadataFileMap.entrySet()) {
+ File segmentMetadataFile = entry.getValue();
+ pushSegmentMetadata(spec, fileUploadDownloadClient, entry.getKey(),
+ getSegmentNameFromMetadataFile(segmentMetadataFile),
segmentMetadataFile, headers, parameters);
+ }
+ } finally {
+ closeFileUploadDownloadClient(spec, fileUploadDownloadClient);
+ }
+ }
+
+ private static void pushSegmentMetadata(SegmentGenerationJobSpec spec,
+ FileUploadDownloadClient fileUploadDownloadClient, String
segmentUriPath, String segmentName,
+ File segmentMetadataFile, List<Header> headers, List<NameValuePair>
parameters)
+ throws Exception {
+ String tableName = spec.getTableSpec().getTableName();
int socketTimeoutMs = getSocketTimeoutMs(spec);
- ExecutorService executor = Executors.newFixedThreadPool(nThreads);
+ for (PinotClusterSpec pinotClusterSpec : spec.getPinotClusterSpecs()) {
+ URI controllerURI;
+ try {
+ controllerURI = new URI(pinotClusterSpec.getControllerURI());
+ } catch (URISyntaxException e) {
+ throw new RuntimeException("Got invalid controller uri - '" +
pinotClusterSpec.getControllerURI() + "'");
+ }
+ LOGGER.info("Pushing segment: {} to location: {} for table {}",
segmentName, controllerURI, tableName);
+ int attempts = 1;
+ if (spec.getPushJobSpec() != null &&
spec.getPushJobSpec().getPushAttempts() > 0) {
+ attempts = spec.getPushJobSpec().getPushAttempts();
+ }
+ long retryWaitMs = 1000L;
+ if (spec.getPushJobSpec() != null &&
spec.getPushJobSpec().getPushRetryIntervalMillis() > 0) {
+ retryWaitMs = spec.getPushJobSpec().getPushRetryIntervalMillis();
+ }
+ RetryPolicies.exponentialBackoffRetryPolicy(attempts, retryWaitMs,
5).attempt(() -> {
+ List<Header> reqHttpHeaders = new ArrayList<>(headers);
+ try {
+ reqHttpHeaders.add(
+ new
BasicHeader(FileUploadDownloadClient.CustomHeaders.DOWNLOAD_URI,
segmentUriPath));
+ reqHttpHeaders.add(new
BasicHeader(FileUploadDownloadClient.CustomHeaders.UPLOAD_TYPE,
+ FileUploadDownloadClient.FileUploadType.METADATA.toString()));
+ if (spec.getPushJobSpec() != null) {
+ reqHttpHeaders.add(new
BasicHeader(FileUploadDownloadClient.CustomHeaders.COPY_SEGMENT_TO_DEEP_STORE,
+
String.valueOf(spec.getPushJobSpec().getCopyToDeepStoreForMetadataPush())));
+ }
+
+ SimpleHttpResponse response =
fileUploadDownloadClient.uploadSegmentMetadata(
+ FileUploadDownloadClient.getUploadSegmentURI(controllerURI),
segmentName,
+ segmentMetadataFile, reqHttpHeaders, parameters,
socketTimeoutMs);
+ LOGGER.info("Response for pushing table {} segment {} to location {}
- {}: {}", tableName, segmentName,
+ controllerURI, response.getStatusCode(), response.getResponse());
+ return true;
+ } catch (HttpErrorStatusException e) {
+ int statusCode = e.getStatusCode();
+ if (statusCode >= 500) {
+ // Temporary exception
+ LOGGER.warn("Caught temporary exception while pushing table: {}
segment: {} to {}, will retry",
+ tableName, segmentName, controllerURI, e);
+ return false;
+ } else {
+ // Permanent exception
+ LOGGER.error("Caught permanent exception while pushing table: {}
segment: {} to {}, won't retry",
+ tableName, segmentName, controllerURI, e);
+ throw e;
+ }
+ }
+ });
+ }
+ }
+
+ /// Metadata tar for a staged segment tar: the
`<segmentName>.metadata.tar.gz` sidecar when the spec prefers it and
+ /// it exists, otherwise extracted from the segment tar (downloaded first
unless the filesystem is local).
+ private static File getSegmentMetadataFile(SegmentGenerationJobSpec spec,
PinotFS fileSystem, String tarFilePath,
+ String segmentName)
+ throws Exception {
+ URI metadataTarGzFilePath = generateSegmentMetadataURI(tarFilePath,
segmentName);
+ LOGGER.info("Checking if metadata tar gz file {} exists",
metadataTarGzFilePath);
+ if (spec.getPushJobSpec().isPreferMetadataTarGz() &&
fileSystem.exists(metadataTarGzFilePath)) {
+ File segmentMetadataFile = new File(FileUtils.getTempDirectory(),
+ SegmentUploadConstants.SEGMENT_METADATA_DIR_PREFIX +
UUID.randomUUID()
+ + TarCompressionUtils.TAR_GZ_FILE_EXTENSION);
+ if (segmentMetadataFile.exists()) {
+ FileUtils.forceDelete(segmentMetadataFile);
+ }
+ fileSystem.copyToLocalFile(metadataTarGzFilePath, segmentMetadataFile);
+ return segmentMetadataFile;
+ }
+ 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) {
+ 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;
+ SegmentNameUtils.validatePartialOrFullSegmentName(segmentName);
+ return segmentName;
+ }
+
+ private static String getSegmentNameFromMetadataFile(File
segmentMetadataFile) {
+ String fileName = segmentMetadataFile.getName();
+
Preconditions.checkArgument(fileName.endsWith(Constants.METADATA_TAR_GZ_FILE_EXT),
+ "Segment metadata file: %s must be named <segmentName>%s", fileName,
Constants.METADATA_TAR_GZ_FILE_EXT);
+ String segmentName = fileName.substring(0, fileName.length() -
Constants.METADATA_TAR_GZ_FILE_EXT.length());
+ SegmentNameUtils.validatePartialOrFullSegmentName(segmentName);
+ return segmentName;
+ }
+
+ public static void sendSegmentsUriAndMetadata(SegmentGenerationJobSpec spec,
PinotFS fileSystem,
+ Map<String, String> segmentUriToTarPathMap, List<Header> headers,
List<NameValuePair> parameters)
+ throws Exception {
LOGGER.info("Start pushing segment metadata: {} to locations: {} for
table: {} with parallelism: {}",
- segmentUriToTarPathMap, Arrays.toString(spec.getPinotClusterSpecs()),
tableName,
+ segmentUriToTarPathMap, Arrays.toString(spec.getPinotClusterSpecs()),
spec.getTableSpec().getTableName(),
spec.getPushJobSpec().getPushParallelism());
-
+ ConcurrentHashMap<String, File> segmentMetadataFileMap = new
ConcurrentHashMap<>();
+ ConcurrentLinkedQueue<String> segmentURIs = new ConcurrentLinkedQueue<>();
+ ExecutorService executor =
+
Executors.newFixedThreadPool(spec.getPushJobSpec().getSegmentMetadataGenerationParallelism());
try {
generateSegmentMetadataFiles(spec, fileSystem, segmentUriToTarPathMap,
segmentMetadataFileMap, segmentURIs,
executor);
- allSegmentsMetadataTarFile = createSegmentsMetadataTarFile(segmentURIs,
segmentMetadataFileMap);
- // the key is unused in batch upload mode and hence 'noopKey'
- allSegmentsMetadataMap.put("noopKey", allSegmentsMetadataTarFile);
+ pushSegmentsMetadata(spec, segmentURIs, segmentMetadataFileMap, headers,
parameters);
+ } finally {
+ for (File segmentMetadataFile : segmentMetadataFileMap.values()) {
+ FileUtils.deleteQuietly(segmentMetadataFile);
+ }
+ executor.shutdown();
+ }
+ }
- // perform metadata push in batch mode for every cluster
+ /// Batch variant of [#sendSegmentUriAndMetadata(SegmentGenerationJobSpec,
Map, List, List)] for locally held
+ /// metadata tars named `<segmentName>.metadata.tar.gz`. The caller owns the
files.
+ public static void sendSegmentsUriAndMetadata(SegmentGenerationJobSpec spec,
+ Map<String, File> segmentUriToMetadataFileMap, List<Header> headers,
List<NameValuePair> parameters)
+ throws Exception {
+ Map<String, File> segmentMetadataFileMap = new HashMap<>();
+ List<String> segmentURIs = new ArrayList<>();
+ for (Map.Entry<String, File> entry :
segmentUriToMetadataFileMap.entrySet()) {
+ String segmentName = getSegmentNameFromMetadataFile(entry.getValue());
+ segmentMetadataFileMap.put(segmentName, entry.getValue());
+ segmentURIs.add(segmentName);
+ segmentURIs.add(entry.getKey());
+ }
+ pushSegmentsMetadata(spec, segmentURIs, segmentMetadataFileMap, headers,
parameters);
+ }
+
+ private static void pushSegmentsMetadata(SegmentGenerationJobSpec spec,
Collection<String> segmentURIs,
+ Map<String, File> segmentMetadataFileMap, List<Header> headers,
List<NameValuePair> parameters)
+ throws Exception {
+ String tableName = spec.getTableSpec().getTableName();
+ FileUploadDownloadClient fileUploadDownloadClient =
getOrCreateFileUploadDownloadClient(spec);
+ int socketTimeoutMs = getSocketTimeoutMs(spec);
+ File allSegmentsMetadataTarFile =
createSegmentsMetadataTarFile(segmentURIs, segmentMetadataFileMap);
+ // the key is unused in batch upload mode and hence 'noopKey'
+ Map<String, File> allSegmentsMetadataMap = Map.of("noopKey",
allSegmentsMetadataTarFile);
+ try {
for (PinotClusterSpec pinotClusterSpec : spec.getPinotClusterSpecs()) {
URI controllerURI;
try {
@@ -500,14 +572,8 @@ public class SegmentPushUtils implements Serializable {
});
}
} finally {
- for (Map.Entry<String, File> metadataFileEntry :
segmentMetadataFileMap.entrySet()) {
- FileUtils.deleteQuietly(metadataFileEntry.getValue());
- }
- if (allSegmentsMetadataTarFile != null) {
- FileUtils.deleteQuietly(allSegmentsMetadataTarFile);
- }
+ FileUtils.deleteQuietly(allSegmentsMetadataTarFile);
closeFileUploadDownloadClient(spec, fileUploadDownloadClient);
- executor.shutdown();
}
}
@@ -522,30 +588,8 @@ public class SegmentPushUtils implements Serializable {
futures.add(
executor.submit(() -> {
String tarFilePath = segmentUriToTarPathMap.get(segmentUriPath);
- String fileName = new File(tarFilePath).getName();
- // segments stored in Pinot deep store do not have .tar.gz
extension
- String segmentName = fileName.endsWith(Constants.TAR_GZ_FILE_EXT)
- ? fileName.substring(0, fileName.length() -
Constants.TAR_GZ_FILE_EXT.length()) : fileName;
- SegmentNameUtils.validatePartialOrFullSegmentName(segmentName);
- File segmentMetadataFile;
- // Check if there is a segment metadata tar gz file named
`segmentName.metadata.tar.gz`, already in the
- // remote directory. This is to avoid generating a new segment
metadata tar gz file every time we push a
- // segment, which requires downloading the entire segment tar gz
file.
-
- URI metadataTarGzFilePath =
generateSegmentMetadataURI(tarFilePath, segmentName);
- LOGGER.info("Checking if metadata tar gz file {} exists",
metadataTarGzFilePath);
- if (spec.getPushJobSpec().isPreferMetadataTarGz() &&
fileSystem.exists(metadataTarGzFilePath)) {
- segmentMetadataFile = new File(FileUtils.getTempDirectory(),
- SegmentUploadConstants.SEGMENT_METADATA_DIR_PREFIX +
UUID.randomUUID()
- + TarCompressionUtils.TAR_GZ_FILE_EXTENSION);
- if (segmentMetadataFile.exists()) {
- FileUtils.forceDelete(segmentMetadataFile);
- }
- fileSystem.copyToLocalFile(metadataTarGzFilePath,
segmentMetadataFile);
- } else {
- segmentMetadataFile = generateSegmentMetadataFile(fileSystem,
URI.create(tarFilePath));
- }
- segmentMetadataFileMap.put(segmentName, segmentMetadataFile);
+ String segmentName = getSegmentName(tarFilePath);
+ segmentMetadataFileMap.put(segmentName,
getSegmentMetadataFile(spec, fileSystem, tarFilePath, segmentName));
segmentURIs.add(segmentName);
segmentURIs.add(segmentUriPath);
return null;
@@ -707,6 +751,23 @@ public class SegmentPushUtils implements Serializable {
}
}
+ /// Builds the metadata-only tar a METADATA push sends
(`metadata.properties` and `creation.meta`) from a local
+ /// segment directory, named `<segmentName>.metadata.tar.gz`, so the segment
tar never has to be downloaded back.
+ public static File generateSegmentMetadataFile(File segmentDir, File
outputDir, String segmentName)
+ throws IOException {
+ File metadataDir = new File(outputDir, segmentName + "-metadata");
+ File metadataTarFile = new File(outputDir, segmentName +
Constants.METADATA_TAR_GZ_FILE_EXT);
+ try {
+ FileUtils.forceMkdir(metadataDir);
+
FileUtils.copyFileToDirectory(SegmentDirectoryPaths.findMetadataFile(segmentDir),
metadataDir);
+
FileUtils.copyFileToDirectory(SegmentDirectoryPaths.findCreationMetaFile(segmentDir),
metadataDir);
+ TarCompressionUtils.createCompressedTarFile(metadataDir,
metadataTarFile);
+ return metadataTarFile;
+ } finally {
+ FileUtils.deleteQuietly(metadataDir);
+ }
+ }
+
public static URI generateSegmentMetadataURI(String segmentTarPath, String
segmentName)
throws URISyntaxException {
URI segmentTarURI = URI.create(segmentTarPath);
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 0bf6005b8c3..a1763c47252 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
@@ -34,13 +34,17 @@ import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.security.SecureRandom;
+import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
+import java.util.HashSet;
import java.util.List;
import java.util.Map;
+import java.util.Set;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
@@ -58,6 +62,7 @@ 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.PinotFS;
+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;
import org.apache.pinot.spi.ingestion.batch.spec.SegmentGenerationJobSpec;
@@ -375,6 +380,113 @@ public class SegmentPushUtilsTest {
}
}
+ @Test
+ public void testGenerateSegmentMetadataFileFromLocalSegment()
+ throws Exception {
+ // A v3 segment layout with a data file that must not be included
+ File segmentDir = new File(_tempDir, TEST_SEGMENT_NAME);
+ File v3Dir = new File(segmentDir, "v3");
+ FileUtils.forceMkdir(v3Dir);
+ FileUtils.writeStringToFile(new File(v3Dir, "metadata.properties"),
"segment.name = " + TEST_SEGMENT_NAME,
+ StandardCharsets.UTF_8);
+ FileUtils.writeStringToFile(new File(v3Dir, "creation.meta"), "crc",
StandardCharsets.UTF_8);
+ FileUtils.writeStringToFile(new File(v3Dir, "columns.psf"), "data",
StandardCharsets.UTF_8);
+
+ File metadataTar =
SegmentPushUtils.generateSegmentMetadataFile(segmentDir, _tempDir,
TEST_SEGMENT_NAME);
+
+ assertEquals(metadataTar.getName(), TEST_SEGMENT_NAME +
Constants.METADATA_TAR_GZ_FILE_EXT);
+ 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);
+ }
+
+ @Test
+ public void testSendSegmentUriAndMetadataWithLocalMetadataFiles()
+ throws Exception {
+ List<Headers> requests = new CopyOnWriteArrayList<>();
+ HttpsServer httpsServer = createHttpsServer("/v2/segments", new
RecordingHandler(requests));
+ try {
+ URI controllerUri = new URI("https://localhost:"; +
httpsServer.getAddress().getPort());
+ SegmentGenerationJobSpec spec =
createSegmentGenerationJobSpec(controllerUri, createTlsSpec());
+ spec.getPushJobSpec().setCopyToDeepStoreForMetadataPush(true);
+ File metadataTar = createMetadataTar(TEST_SEGMENT_NAME);
+
+ SegmentPushUtils.sendSegmentUriAndMetadata(spec,
Map.of(TEST_SEGMENT_URI, metadataTar), new ArrayList<>(),
+ new ArrayList<>());
+
+ assertEquals(requests.size(), 1);
+ Headers headers = requests.get(0);
+
assertEquals(headers.getFirst(FileUploadDownloadClient.CustomHeaders.UPLOAD_TYPE),
+ FileUploadDownloadClient.FileUploadType.METADATA.toString());
+
assertEquals(headers.getFirst(FileUploadDownloadClient.CustomHeaders.DOWNLOAD_URI),
TEST_SEGMENT_URI);
+
assertEquals(headers.getFirst(FileUploadDownloadClient.CustomHeaders.COPY_SEGMENT_TO_DEEP_STORE),
"true");
+ // The caller owns the metadata file
+ assertTrue(metadataTar.isFile());
+ } finally {
+ httpsServer.stop(0);
+ }
+ }
+
+ @Test
+ public void testSendSegmentsUriAndMetadataWithLocalMetadataFiles()
+ throws Exception {
+ List<Headers> requests = new CopyOnWriteArrayList<>();
+ HttpsServer httpsServer = createHttpsServer("/segments/batchUpload", new
RecordingHandler(requests));
+ try {
+ URI controllerUri = new URI("https://localhost:"; +
httpsServer.getAddress().getPort());
+ SegmentGenerationJobSpec spec =
createSegmentGenerationJobSpec(controllerUri, createTlsSpec());
+ Map<String, File> segmentUriToMetadataFileMap = new HashMap<>();
+ segmentUriToMetadataFileMap.put("file:///tmp/segment1.tar.gz",
createMetadataTar("segment1"));
+ segmentUriToMetadataFileMap.put("file:///tmp/segment2.tar.gz",
createMetadataTar("segment2"));
+
+ SegmentPushUtils.sendSegmentsUriAndMetadata(spec,
segmentUriToMetadataFileMap, new ArrayList<>(),
+ new ArrayList<>());
+
+ // One batch request for both segments, and the caller still owns the
metadata files
+ assertEquals(requests.size(), 1);
+
assertEquals(requests.get(0).getFirst(FileUploadDownloadClient.CustomHeaders.UPLOAD_TYPE),
+ FileUploadDownloadClient.FileUploadType.METADATA.toString());
+ segmentUriToMetadataFileMap.values().forEach(file ->
assertTrue(file.isFile()));
+ } finally {
+ httpsServer.stop(0);
+ }
+ }
+
+ @Test
+ public void testSendSegmentUriAndMetadataRejectsUnnamedMetadataFile() {
+ SegmentGenerationJobSpec spec =
createSegmentGenerationJobSpec(URI.create("https://localhost:1";), null);
+ File notAMetadataTar = new File(_tempDir, "segment.tar.gz");
+ assertThrows(IllegalArgumentException.class, () ->
SegmentPushUtils.sendSegmentUriAndMetadata(spec,
+ Map.of(TEST_SEGMENT_URI, notAMetadataTar), new ArrayList<>(), new
ArrayList<>()));
+ }
+
+ private File createMetadataTar(String segmentName)
+ throws IOException {
+ File segmentDir = new File(_tempDir, segmentName);
+ FileUtils.forceMkdir(segmentDir);
+ FileUtils.touch(new File(segmentDir, "metadata.properties"));
+ FileUtils.touch(new File(segmentDir, "creation.meta"));
+ return SegmentPushUtils.generateSegmentMetadataFile(segmentDir, _tempDir,
segmentName);
+ }
+
+ /// Records request headers and answers OK, for both single and batch
metadata pushes.
+ private static class RecordingHandler implements HttpHandler {
+ private final List<Headers> _requests;
+
+ private RecordingHandler(List<Headers> requests) {
+ _requests = requests;
+ }
+
+ @Override
+ public void handle(HttpExchange httpExchange)
+ throws IOException {
+ _requests.add(httpExchange.getRequestHeaders());
+ httpExchange.getRequestBody().readAllBytes();
+ writeResponse(httpExchange, OK_RESPONSE);
+ }
+ }
+
@Test
public void testGetOrCreateFileUploadDownloadClientUsesSharedDefaultClient()
throws Exception {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
