Skip to content

Commit 83451d7

Browse files
tkaymakclaude
andcommitted
feat: add Spark 4 structured streaming runner source
Add the Spark 4 structured streaming runner implementation and tests. Most files are adapted from the Spark 3 structured streaming runner with targeted changes for Spark 4 / Scala 2.13 API compatibility. Key Spark 4-specific changes (diff against runners/spark/3/src/): EncoderFactory — Replaced the direct ExpressionEncoder constructor (removed in Spark 4) with BeamAgnosticEncoder, a named class implementing both AgnosticExpressionPathEncoder (for expression delegation via toCatalyst/fromCatalyst) and AgnosticEncoders .StructEncoder (so Dataset.select(TypedColumn) creates an N-attribute plan, preventing FIELD_NUMBER_MISMATCH). The toCatalyst/fromCatalyst methods substitute the provided input expression via transformUp, enabling correct nesting inside composite encoders like Encoders.tuple(). EncoderHelpers — Added toExpressionEncoder() helper to handle Spark 4 built-in encoders that are AgnosticEncoder subclasses rather than ExpressionEncoder. GroupByKeyTranslatorBatch — Migrated from internal catalyst Expression API (CreateNamedStruct, Literal$) to public Column API (struct(), lit(), array()), as required by Spark 4. BoundedDatasetFactory — Use classic.Dataset$.MODULE$.ofRows() as Dataset moved to org.apache.spark.sql.classic in Spark 4. ScalaInterop — Replace WrappedArray.ofRef (removed in Scala 2.13) with JavaConverters.asScalaBuffer().toList() in seqOf(). GroupByKeyHelpers, CombinePerKeyTranslatorBatch — Replace TraversableOnce with IterableOnce (Scala 2.13 rename). SparkStructuredStreamingPipelineResult — Replace sparkproject.guava with Beam's vendored Guava. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 46e2dba commit 83451d7

73 files changed

Lines changed: 10030 additions & 0 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
package org.apache.beam.runners.spark.structuredstreaming;
19+
20+
import org.apache.beam.runners.spark.SparkCommonPipelineOptions;
21+
import org.apache.beam.sdk.options.Default;
22+
import org.apache.beam.sdk.options.Description;
23+
import org.apache.beam.sdk.options.PipelineOptions;
24+
25+
/**
26+
* Spark runner {@link PipelineOptions} handles Spark execution-related configurations, such as the
27+
* master address, and other user-related knobs.
28+
*/
29+
public interface SparkStructuredStreamingPipelineOptions extends SparkCommonPipelineOptions {
30+
31+
/** Set to true to run the job in test mode. */
32+
@Default.Boolean(false)
33+
boolean getTestMode();
34+
35+
void setTestMode(boolean testMode);
36+
37+
@Description("Enable if the runner should use the currently active Spark session.")
38+
@Default.Boolean(false)
39+
boolean getUseActiveSparkSession();
40+
41+
void setUseActiveSparkSession(boolean value);
42+
}
Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,131 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
package org.apache.beam.runners.spark.structuredstreaming;
19+
20+
import static org.apache.beam.runners.core.metrics.MetricsContainerStepMap.asAttemptedOnlyMetricResults;
21+
import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects.firstNonNull;
22+
23+
import java.io.IOException;
24+
import java.util.concurrent.ExecutionException;
25+
import java.util.concurrent.Future;
26+
import java.util.concurrent.TimeUnit;
27+
import java.util.concurrent.TimeoutException;
28+
import javax.annotation.Nullable;
29+
import org.apache.beam.runners.spark.structuredstreaming.metrics.MetricsAccumulator;
30+
import org.apache.beam.sdk.Pipeline;
31+
import org.apache.beam.sdk.PipelineResult;
32+
import org.apache.beam.sdk.metrics.MetricResults;
33+
import org.apache.beam.sdk.util.UserCodeException;
34+
import org.apache.spark.SparkException;
35+
import org.joda.time.Duration;
36+
37+
public class SparkStructuredStreamingPipelineResult implements PipelineResult {
38+
39+
private final Future<?> pipelineExecution;
40+
private final MetricsAccumulator metrics;
41+
private @Nullable final Runnable onTerminalState;
42+
private PipelineResult.State state;
43+
44+
SparkStructuredStreamingPipelineResult(
45+
Future<?> pipelineExecution,
46+
MetricsAccumulator metrics,
47+
@Nullable final Runnable onTerminalState) {
48+
this.pipelineExecution = pipelineExecution;
49+
this.metrics = metrics;
50+
this.onTerminalState = onTerminalState;
51+
// pipelineExecution is expected to have started executing eagerly.
52+
this.state = State.RUNNING;
53+
}
54+
55+
private static RuntimeException runtimeExceptionFrom(final Throwable e) {
56+
return (e instanceof RuntimeException) ? (RuntimeException) e : new RuntimeException(e);
57+
}
58+
59+
/**
60+
* Unwrap cause of SparkException or UserCodeException as PipelineExecutionException. Otherwise,
61+
* return {@code exception} as RuntimeException.
62+
*/
63+
private static RuntimeException unwrapCause(Throwable exception) {
64+
Throwable next = exception;
65+
while (next != null && (next instanceof SparkException || next instanceof UserCodeException)) {
66+
exception = next;
67+
next = next.getCause();
68+
}
69+
return exception == next
70+
? runtimeExceptionFrom(exception)
71+
: new Pipeline.PipelineExecutionException(firstNonNull(next, exception));
72+
}
73+
74+
private State awaitTermination(Duration duration)
75+
throws TimeoutException, ExecutionException, InterruptedException {
76+
pipelineExecution.get(duration.getMillis(), TimeUnit.MILLISECONDS);
77+
// Throws an exception if the job is not finished successfully in the given time.
78+
return PipelineResult.State.DONE;
79+
}
80+
81+
@Override
82+
public PipelineResult.State getState() {
83+
return state;
84+
}
85+
86+
@Override
87+
public PipelineResult.State waitUntilFinish() {
88+
return waitUntilFinish(Duration.millis(Long.MAX_VALUE));
89+
}
90+
91+
@Override
92+
public State waitUntilFinish(final Duration duration) {
93+
try {
94+
State finishState = awaitTermination(duration);
95+
offerNewState(finishState);
96+
} catch (final TimeoutException e) {
97+
// ignore.
98+
} catch (final ExecutionException e) {
99+
offerNewState(PipelineResult.State.FAILED);
100+
throw unwrapCause(firstNonNull(e.getCause(), e));
101+
} catch (final Exception e) {
102+
offerNewState(PipelineResult.State.FAILED);
103+
throw unwrapCause(e);
104+
}
105+
106+
return state;
107+
}
108+
109+
@Override
110+
public MetricResults metrics() {
111+
return asAttemptedOnlyMetricResults(metrics.value());
112+
}
113+
114+
@Override
115+
public PipelineResult.State cancel() throws IOException {
116+
offerNewState(PipelineResult.State.CANCELLED);
117+
return state;
118+
}
119+
120+
private void offerNewState(State newState) {
121+
State oldState = this.state;
122+
this.state = newState;
123+
if (!oldState.isTerminal() && newState.isTerminal() && onTerminalState != null) {
124+
try {
125+
onTerminalState.run();
126+
} catch (Exception e) {
127+
throw unwrapCause(e);
128+
}
129+
}
130+
}
131+
}

0 commit comments

Comments
 (0)