Skip to content

Commit 76a246f

Browse files
ymuppalaymuppala
authored
[vpj] enabling kif repush spark job (linkedin#2468)
1 parent 16a81ce commit 76a246f

8 files changed

Lines changed: 306 additions & 11 deletions

File tree

‎clients/venice-push-job/src/main/java/com/linkedin/venice/hadoop/VenicePushJob.java‎

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@
7070
import static com.linkedin.venice.vpj.VenicePushJobConstants.SOURCE_ETL;
7171
import static com.linkedin.venice.vpj.VenicePushJobConstants.SOURCE_GRID_FABRIC;
7272
import static com.linkedin.venice.vpj.VenicePushJobConstants.SOURCE_KAFKA;
73+
import static com.linkedin.venice.vpj.VenicePushJobConstants.SPARK_KIF_REPUSH_ENABLED;
7374
import static com.linkedin.venice.vpj.VenicePushJobConstants.SUPPRESS_END_OF_PUSH_MESSAGE;
7475
import static com.linkedin.venice.vpj.VenicePushJobConstants.SYSTEM_SCHEMA_READER_ENABLED;
7576
import static com.linkedin.venice.vpj.VenicePushJobConstants.TARGETED_REGION_PUSH_ENABLED;
@@ -542,13 +543,25 @@ private PushJobSetting getPushJobSetting(VeniceProperties props) {
542543
// Compute-engine abstraction related configs
543544
String dataWriterComputeJobClass = props.getString(DATA_WRITER_COMPUTE_JOB_CLASS, (String) null);
544545

545-
// Currently, only MR mode supports KIF. This is temporary.
546-
if (dataWriterComputeJobClass == null || pushJobSettingToReturn.isSourceKafka) {
546+
if (dataWriterComputeJobClass == null) {
547547
pushJobSettingToReturn.dataWriterComputeJobClass = DataWriterMRJob.class;
548548
} else {
549549
Class objectClass = ReflectUtils.loadClass(dataWriterComputeJobClass);
550550
Validate.isAssignableFrom(DataWriterComputeJob.class, objectClass);
551-
pushJobSettingToReturn.dataWriterComputeJobClass = objectClass;
551+
552+
// For KIF repush jobs, only use the configured Spark compute job class
553+
// if the spark.kif.repush.enabled flag is explicitly set to true. This allows gradual rollout
554+
// of Spark for KIF repush without affecting all jobs at once.
555+
if (pushJobSettingToReturn.isSourceKafka && !props.getBoolean(SPARK_KIF_REPUSH_ENABLED, false)) {
556+
LOGGER.info(
557+
"KIF repush detected but {} is not enabled. Falling back to MapReduce. "
558+
+ "Set {}=true to use Spark for KIF repush.",
559+
SPARK_KIF_REPUSH_ENABLED,
560+
SPARK_KIF_REPUSH_ENABLED);
561+
pushJobSettingToReturn.dataWriterComputeJobClass = DataWriterMRJob.class;
562+
} else {
563+
pushJobSettingToReturn.dataWriterComputeJobClass = objectClass;
564+
}
552565
}
553566
return pushJobSettingToReturn;
554567
}

‎clients/venice-push-job/src/main/java/com/linkedin/venice/hadoop/task/datawriter/AbstractPartitionWriter.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -326,7 +326,9 @@ protected VeniceWriterMessage extract(
326326

327327
VeniceRecordWithMetadata valueRecord = values.next();
328328
byte[] valueBytes = valueRecord.getValue();
329-
ByteBuffer rmd = valueRecord.getRmd() == null ? null : ByteBuffer.wrap(valueRecord.getRmd());
329+
// Handle empty RMD the same way as null - don't wrap empty byte array into ByteBuffer
330+
byte[] rmdBytes = valueRecord.getRmd();
331+
ByteBuffer rmd = (rmdBytes == null || rmdBytes.length == 0) ? null : ByteBuffer.wrap(rmdBytes);
330332

331333
if (duplicateKeyPrinter == null) {
332334
throw new VeniceException("'DuplicateKeyPrinter' is not initialized properly");

‎clients/venice-push-job/src/main/java/com/linkedin/venice/spark/datawriter/jobs/DataWriterSparkJob.java‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
package com.linkedin.venice.spark.datawriter.jobs;
22

3+
import static com.linkedin.venice.ConfigKeys.PUBSUB_BROKER_ADDRESS;
4+
import static com.linkedin.venice.ConfigKeys.PUBSUB_SECURITY_PROTOCOL;
35
import static com.linkedin.venice.spark.SparkConstants.DEFAULT_SCHEMA;
46
import static com.linkedin.venice.vpj.VenicePushJobConstants.ETL_VALUE_SCHEMA_TRANSFORMATION;
57
import static com.linkedin.venice.vpj.VenicePushJobConstants.FILE_KEY_SCHEMA;
@@ -15,6 +17,11 @@
1517
import static com.linkedin.venice.vpj.VenicePushJobConstants.RMD_SCHEMA_PROP;
1618
import static com.linkedin.venice.vpj.VenicePushJobConstants.SCHEMA_STRING_PROP;
1719
import static com.linkedin.venice.vpj.VenicePushJobConstants.SPARK_NATIVE_INPUT_FORMAT_ENABLED;
20+
import static com.linkedin.venice.vpj.VenicePushJobConstants.SSL_CONFIGURATOR_CLASS_CONFIG;
21+
import static com.linkedin.venice.vpj.VenicePushJobConstants.SSL_KEY_PASSWORD_PROPERTY_NAME;
22+
import static com.linkedin.venice.vpj.VenicePushJobConstants.SSL_KEY_STORE_PASSWORD_PROPERTY_NAME;
23+
import static com.linkedin.venice.vpj.VenicePushJobConstants.SSL_KEY_STORE_PROPERTY_NAME;
24+
import static com.linkedin.venice.vpj.VenicePushJobConstants.SSL_TRUST_STORE_PROPERTY_NAME;
1825
import static com.linkedin.venice.vpj.VenicePushJobConstants.UPDATE_SCHEMA_STRING_PROP;
1926
import static com.linkedin.venice.vpj.VenicePushJobConstants.VALUE_FIELD_PROP;
2027
import static com.linkedin.venice.vpj.VenicePushJobConstants.VSON_PUSH;
@@ -24,6 +31,7 @@
2431
import com.linkedin.venice.hadoop.input.kafka.KafkaInputUtils;
2532
import com.linkedin.venice.hadoop.input.recordreader.avro.VeniceAvroRecordReader;
2633
import com.linkedin.venice.hadoop.input.recordreader.vson.VeniceVsonRecordReader;
34+
import com.linkedin.venice.pubsub.api.PubSubSecurityProtocol;
2735
import com.linkedin.venice.spark.input.hdfs.VeniceHdfsSource;
2836
import com.linkedin.venice.spark.input.pubsub.raw.VeniceRawPubsubSource;
2937
import com.linkedin.venice.spark.utils.RowToAvroConverter;
@@ -160,6 +168,7 @@ protected Dataset<Row> getKafkaInputDataFrame() {
160168
// Configure Kafka input connection
161169
setInputConf(sparkSession, dataFrameReader, KAFKA_INPUT_TOPIC, pushJobSetting.kafkaInputTopic);
162170
setInputConf(sparkSession, dataFrameReader, KAFKA_INPUT_BROKER_URL, pushJobSetting.kafkaInputBrokerUrl);
171+
setInputConf(sparkSession, dataFrameReader, PUBSUB_BROKER_ADDRESS, pushJobSetting.kafkaInputBrokerUrl);
163172
setInputConf(
164173
sparkSession,
165174
dataFrameReader,
@@ -170,6 +179,28 @@ protected Dataset<Row> getKafkaInputDataFrame() {
170179
KafkaInputUtils.putSchemaMapIntoProperties(pushJobSetting.newKmeSchemasFromController)
171180
.forEach((key, value) -> setInputConf(sparkSession, dataFrameReader, key, value));
172181

182+
// Pass SSL configuration if enabled
183+
if (pushJobSetting.enableSSL) {
184+
passSSLConfigToDataFrameReader(sparkSession, dataFrameReader);
185+
}
186+
173187
return dataFrameReader.load();
174188
}
189+
190+
/**
191+
* Passes SSL metadata properties to the DataFrameReader so they are available to
192+
* {@link com.linkedin.venice.spark.input.pubsub.SparkPubSubInputFormat} (driver) and
193+
* {@link com.linkedin.venice.spark.input.pubsub.SparkPubSubPartitionReaderFactory} (executors).
194+
*/
195+
private void passSSLConfigToDataFrameReader(SparkSession sparkSession, DataFrameReader dataFrameReader) {
196+
VeniceProperties jobProps = getJobProperties();
197+
String[] sslMetadataKeys = { SSL_CONFIGURATOR_CLASS_CONFIG, SSL_KEY_STORE_PROPERTY_NAME,
198+
SSL_TRUST_STORE_PROPERTY_NAME, SSL_KEY_PASSWORD_PROPERTY_NAME, SSL_KEY_STORE_PASSWORD_PROPERTY_NAME };
199+
for (String key: sslMetadataKeys) {
200+
if (jobProps.containsKey(key)) {
201+
setInputConf(sparkSession, dataFrameReader, key, jobProps.getString(key));
202+
}
203+
}
204+
setInputConf(sparkSession, dataFrameReader, PUBSUB_SECURITY_PROTOCOL, PubSubSecurityProtocol.SSL.name());
205+
}
175206
}

‎clients/venice-push-job/src/main/java/com/linkedin/venice/spark/input/pubsub/SparkPubSubInputFormat.java‎

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,15 @@
11
package com.linkedin.venice.spark.input.pubsub;
22

33
import static com.linkedin.venice.vpj.VenicePushJobConstants.KAFKA_INPUT_TOPIC;
4+
import static com.linkedin.venice.vpj.VenicePushJobConstants.SSL_CONFIGURATOR_CLASS_CONFIG;
45

56
import com.linkedin.venice.annotation.VisibleForTesting;
7+
import com.linkedin.venice.hadoop.utils.VPJSSLUtils;
68
import com.linkedin.venice.utils.VeniceProperties;
79
import com.linkedin.venice.vpj.pubsub.input.PubSubPartitionSplit;
810
import com.linkedin.venice.vpj.pubsub.input.PubSubSplitPlanner;
911
import java.util.List;
12+
import java.util.Properties;
1013
import java.util.function.Supplier;
1114
import org.apache.logging.log4j.LogManager;
1215
import org.apache.logging.log4j.Logger;
@@ -35,8 +38,11 @@ public SparkPubSubInputFormat(VeniceProperties jobConfig) {
3538

3639
@Override
3740
public InputPartition[] planInputPartitions() {
41+
// Setup SSL on the driver side
42+
VeniceProperties configWithSsl = setupSSLForDriver(jobConfig);
43+
3844
PubSubSplitPlanner planner = plannerSupplier.get();
39-
List<PubSubPartitionSplit> planned = planner.plan(jobConfig);
45+
List<PubSubPartitionSplit> planned = planner.plan(configWithSsl);
4046
InputPartition[] partitions = new InputPartition[planned.size()];
4147

4248
int index = 0;
@@ -58,6 +64,27 @@ public PartitionReaderFactory createReaderFactory() {
5864
return new SparkPubSubPartitionReaderFactory(jobConfig);
5965
}
6066

67+
/**
68+
* Sets up SSL on the driver side before partition planning.
69+
*/
70+
private VeniceProperties setupSSLForDriver(VeniceProperties config) {
71+
if (!config.containsKey(SSL_CONFIGURATOR_CLASS_CONFIG)) {
72+
return config;
73+
}
74+
try {
75+
Properties sslProps = VPJSSLUtils.getSslProperties(config);
76+
Properties merged = config.toProperties();
77+
merged.putAll(sslProps);
78+
return new VeniceProperties(merged);
79+
} catch (Exception e) {
80+
String msg = "Failed to setup SSL for driver-side partition planning. "
81+
+ "Ensure the Hadoop token file is accessible and SSL certificates are valid. " + "SSL configurator class: "
82+
+ config.getString(SSL_CONFIGURATOR_CLASS_CONFIG);
83+
LOGGER.error(msg, e);
84+
throw new RuntimeException(msg, e);
85+
}
86+
}
87+
6188
@Override
6289
public StructType readSchema() {
6390
return null;

‎clients/venice-push-job/src/main/java/com/linkedin/venice/spark/input/pubsub/SparkPubSubPartitionReaderFactory.java‎

Lines changed: 33 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,11 @@
44
import static com.linkedin.venice.vpj.VenicePushJobConstants.KAFKA_INPUT_BROKER_URL;
55
import static com.linkedin.venice.vpj.VenicePushJobConstants.KAFKA_INPUT_FABRIC;
66
import static com.linkedin.venice.vpj.VenicePushJobConstants.PUBSUB_INPUT_SECONDARY_COMPARATOR_USE_LOCAL_LOGICAL_INDEX;
7+
import static com.linkedin.venice.vpj.VenicePushJobConstants.SSL_CONFIGURATOR_CLASS_CONFIG;
78

89
import com.linkedin.venice.chunking.ChunkKeyValueTransformer;
910
import com.linkedin.venice.chunking.ChunkKeyValueTransformerImpl;
11+
import com.linkedin.venice.hadoop.utils.VPJSSLUtils;
1012
import com.linkedin.venice.pubsub.PubSubClientsFactory;
1113
import com.linkedin.venice.pubsub.PubSubConsumerAdapterContext;
1214
import com.linkedin.venice.pubsub.PubSubPositionTypeRegistry;
@@ -17,6 +19,7 @@
1719
import com.linkedin.venice.utils.VeniceProperties;
1820
import com.linkedin.venice.vpj.VenicePushJobConstants;
1921
import com.linkedin.venice.vpj.pubsub.input.PubSubPartitionSplit;
22+
import java.util.Properties;
2023
import org.apache.avro.Schema;
2124
import org.apache.logging.log4j.LogManager;
2225
import org.apache.logging.log4j.Logger;
@@ -67,27 +70,30 @@ public PartitionReader<InternalRow> createReader(final InputPartition genericInp
6770
"SparkPubSubPartitionReaderFactory can only create readers for SparkPubSubInputPartitionReader");
6871
}
6972

73+
// Setup SSL on the executor side
74+
VeniceProperties configWithSsl = setupSSLForExecutor(jobConfig);
75+
7076
final SparkPubSubInputPartition inputPartition = (SparkPubSubInputPartition) genericInputPartition;
7177
final PubSubPartitionSplit partitionSplit = inputPartition.getPubSubPartitionSplit();
7278
final PubSubTopicPartition topicPartition = partitionSplit.getPubSubTopicPartition();
7379
final PubSubTopicRepository topicRepository = partitionSplit.getTopicRepository();
74-
final String inputRegionBroker = jobConfig.getString(KAFKA_INPUT_BROKER_URL);
75-
final String regionName = jobConfig.getString(KAFKA_INPUT_FABRIC, inputRegionBroker);
80+
final String inputRegionBroker = configWithSsl.getString(KAFKA_INPUT_BROKER_URL);
81+
final String regionName = configWithSsl.getString(KAFKA_INPUT_FABRIC, inputRegionBroker);
7682
final String consumerName = String.format("raw_kif_%s_%s", inputRegionBroker, topicPartition);
7783

7884
// Create consumer adapter with proper context
7985
final PubSubConsumerAdapterContext consumerContext =
8086
new PubSubConsumerAdapterContext.Builder().setPubSubBrokerAddress(inputRegionBroker)
81-
.setVeniceProperties(jobConfig)
87+
.setVeniceProperties(configWithSsl)
8288
.setPubSubTopicRepository(topicRepository)
8389
.setPubSubMessageDeserializer(PubSubMessageDeserializer.createOptimizedDeserializer())
84-
.setPubSubPositionTypeRegistry(PubSubPositionTypeRegistry.fromPropertiesOrDefault(jobConfig))
90+
.setPubSubPositionTypeRegistry(PubSubPositionTypeRegistry.fromPropertiesOrDefault(configWithSsl))
8591
.setConsumerName(consumerName)
8692
.build();
8793
final PubSubConsumerAdapter pubSubConsumer =
88-
PubSubClientsFactory.createConsumerFactory(jobConfig).create(consumerContext);
94+
PubSubClientsFactory.createConsumerFactory(configWithSsl).create(consumerContext);
8995

90-
boolean shouldUseLocallyBuiltIndexAsOffset = jobConfig.getBoolean(
96+
boolean shouldUseLocallyBuiltIndexAsOffset = configWithSsl.getBoolean(
9197
PUBSUB_INPUT_SECONDARY_COMPARATOR_USE_LOCAL_LOGICAL_INDEX,
9298
DEFAULT_PUBSUB_INPUT_SECONDARY_COMPARATOR_USE_LOCAL_LOGICAL_INDEX);
9399
SparkPubSubInputPartitionReader reader = new SparkPubSubInputPartitionReader(
@@ -106,6 +112,27 @@ public PartitionReader<InternalRow> createReader(final InputPartition genericInp
106112
return reader;
107113
}
108114

115+
/**
116+
* Sets up SSL on the executor side before creating the PubSub consumer.
117+
*/
118+
private VeniceProperties setupSSLForExecutor(VeniceProperties config) {
119+
if (!config.containsKey(SSL_CONFIGURATOR_CLASS_CONFIG)) {
120+
return config;
121+
}
122+
try {
123+
Properties sslProps = VPJSSLUtils.getSslProperties(config);
124+
Properties merged = config.toProperties();
125+
merged.putAll(sslProps);
126+
return new VeniceProperties(merged);
127+
} catch (Exception e) {
128+
String msg = "Failed to setup SSL for executor-side consumer creation. "
129+
+ "Ensure the Hadoop token file is accessible and SSL certificates are valid. " + "SSL configurator class: "
130+
+ config.getString(SSL_CONFIGURATOR_CLASS_CONFIG);
131+
LOGGER.error(msg, e);
132+
throw new RuntimeException(msg, e);
133+
}
134+
}
135+
109136
// Make it explicit that this reader does not support columnar reads.
110137
@Override
111138
public boolean supportColumnarReads(InputPartition partition) {

‎clients/venice-push-job/src/main/java/com/linkedin/venice/vpj/VenicePushJobConstants.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,13 @@ private VenicePushJobConstants() {
4242
// This is a temporary config used to rollout the native input format for Spark. This will be removed soon
4343
public static final String SPARK_NATIVE_INPUT_FORMAT_ENABLED = "spark.native.input.format.enabled";
4444

45+
/**
46+
* Feature flag to enable Spark for KIF (Kafka Input Format) repush jobs.
47+
* When false (default), KIF repush jobs fall back to MapReduce even when a Spark compute job class is configured.
48+
* This allows gradual rollout of Spark for KIF repush without affecting all jobs at once.
49+
*/
50+
public static final String SPARK_KIF_REPUSH_ENABLED = "spark.kif.repush.enabled";
51+
4552
// Vson input configs
4653
// Vson files store key/value schema on file header. key / value fields are optional
4754
// and should be specified only when key / value schema is the partial of the files.

0 commit comments

Comments
 (0)