This is an automated email from the ASF dual-hosted git repository.
robertwb pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new ccba2e3a90f Add YamlTransform to the Java SDK. (#34596)
ccba2e3a90f is described below
commit ccba2e3a90fb5fc746f1f8752f9e614ebec357bc
Author: Robert Bradshaw <[email protected]>
AuthorDate: Fri Apr 11 08:14:04 2025 -0700
Add YamlTransform to the Java SDK. (#34596)
---
build.gradle.kts | 1 +
.../extensions/python/PythonExternalTransform.java | 3 +
sdks/java/extensions/yaml/build.gradle | 39 ++++
.../beam/sdk/extensions/yaml/YamlTransform.java | 210 +++++++++++++++++++++
.../beam/sdk/extensions/yaml/package-info.java | 20 ++
.../sdk/extensions/yaml/YamlTransformTest.java | 143 ++++++++++++++
settings.gradle.kts | 1 +
7 files changed, 417 insertions(+)
diff --git a/build.gradle.kts b/build.gradle.kts
index b901715f085..9cd63258a4a 100644
--- a/build.gradle.kts
+++ b/build.gradle.kts
@@ -303,6 +303,7 @@ tasks.register("javaPreCommit") {
dependsOn(":sdks:java:extensions:sketching:build")
dependsOn(":sdks:java:extensions:sorter:build")
dependsOn(":sdks:java:extensions:timeseries:build")
+ dependsOn(":sdks:java:extensions:yaml:build")
dependsOn(":sdks:java:extensions:zetasketch:build")
dependsOn(":sdks:java:harness:build")
dependsOn(":sdks:java:harness:jmh:build")
diff --git
a/sdks/java/extensions/python/src/main/java/org/apache/beam/sdk/extensions/python/PythonExternalTransform.java
b/sdks/java/extensions/python/src/main/java/org/apache/beam/sdk/extensions/python/PythonExternalTransform.java
index 50782060d0f..fb44f827215 100644
---
a/sdks/java/extensions/python/src/main/java/org/apache/beam/sdk/extensions/python/PythonExternalTransform.java
+++
b/sdks/java/extensions/python/src/main/java/org/apache/beam/sdk/extensions/python/PythonExternalTransform.java
@@ -54,6 +54,7 @@ import org.apache.beam.sdk.util.ReleaseInfo;
import org.apache.beam.sdk.util.construction.External;
import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionRowTuple;
import org.apache.beam.sdk.values.PCollectionTuple;
import org.apache.beam.sdk.values.PInput;
import org.apache.beam.sdk.values.POutput;
@@ -586,6 +587,8 @@ public class PythonExternalTransform<InputT extends PInput,
OutputT extends POut
outputs = ((PCollection<?>) input).apply(transform);
} else if (input instanceof PCollectionTuple) {
outputs = ((PCollectionTuple) input).apply(transform);
+ } else if (input instanceof PCollectionRowTuple) {
+ outputs = ((PCollectionRowTuple) input).apply(transform);
} else if (input instanceof PBegin) {
outputs = ((PBegin) input).apply(transform);
} else {
diff --git a/sdks/java/extensions/yaml/build.gradle
b/sdks/java/extensions/yaml/build.gradle
new file mode 100644
index 00000000000..5b86d323922
--- /dev/null
+++ b/sdks/java/extensions/yaml/build.gradle
@@ -0,0 +1,39 @@
+/*
+ * 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.
+ */
+
+plugins { id 'org.apache.beam.module' }
+applyJavaNature( automaticModuleName: 'org.apache.beam.sdk.extensions.yaml')
+
+description = "Apache Beam :: SDKs :: Java :: Extensions :: Yaml"
+
+evaluationDependsOn(":sdks:java:core")
+
+dependencies {
+ implementation library.java.vendored_guava_32_1_2_jre
+ implementation project(path: ":sdks:java:core", configuration: "shadow")
+ implementation project(path: ":sdks:java:extensions:python")
+ testImplementation library.java.junit
+ testImplementation library.java.hamcrest
+ testImplementation project(":sdks:java:core").sourceSets.test.output
+ testImplementation project(path: ":runners:portability:java")
+ testRuntimeOnly library.java.slf4j_simple
+}
+
+tasks.test {
+ systemProperty "testRunner.jobEndpoint",
System.getProperty("testRunner.jobEndpoint")
+}
diff --git
a/sdks/java/extensions/yaml/src/main/java/org/apache/beam/sdk/extensions/yaml/YamlTransform.java
b/sdks/java/extensions/yaml/src/main/java/org/apache/beam/sdk/extensions/yaml/YamlTransform.java
new file mode 100644
index 00000000000..11245c4b4ca
--- /dev/null
+++
b/sdks/java/extensions/yaml/src/main/java/org/apache/beam/sdk/extensions/yaml/YamlTransform.java
@@ -0,0 +1,210 @@
+/*
+ * 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.beam.sdk.extensions.yaml;
+
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
+import org.apache.beam.sdk.extensions.python.PythonExternalTransform;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.values.PBegin;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionRowTuple;
+import org.apache.beam.sdk.values.PInput;
+import org.apache.beam.sdk.values.POutput;
+import org.apache.beam.sdk.values.PValue;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Joiner;
+import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
+import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * Allows one to invoke <a
href="https://beam.apache.org/documentation/sdks/yaml/";>Beam YAML</a>
+ * transforms from Java.
+ *
+ * <p>This leverages Beam's cross-langauge transforms. Although python is
required to parse and
+ * expand the given transforms, the actual implementation may still be in Java.
+ *
+ * @param <InputT> the type of the input to this PTransform
+ * @param <OutputT> the type of the output to this PTransform
+ */
+public class YamlTransform<InputT extends PInput, OutputT extends POutput>
+ extends PTransform<InputT, OutputT> {
+
+ /** The YAML definition of this transform. */
+ private final String yamlDefinition;
+ /**
+ * If non-null, the set of input tags that are expected to be passed to this
transform.
+ *
+ * <p>If null, a {@literal PCollection<Row>} or PBegin is expected.
+ */
+ private final @Nullable Set<String> inputTags;
+
+ /**
+ * If non-null, the set of output tags that are expected to be produced by
this transform.
+ *
+ * <p>If null, exactly one output is expected and will be returned as a
{@literal
+ * PCollection<Row>}.
+ */
+ private final @Nullable Set<String> outputTags;
+
+ private YamlTransform(
+ String yamlDefinition,
+ @Nullable Iterable<String> inputTags,
+ @Nullable Iterable<String> outputTags) {
+ this.yamlDefinition = yamlDefinition;
+ this.inputTags = inputTags == null ? null : ImmutableSet.copyOf(inputTags);
+ this.outputTags = outputTags == null ? null :
ImmutableSet.copyOf(outputTags);
+ }
+
+ /**
+ * Creates a new YamlTransform mapping a single input {@literal
PCollection<Row>} to a single
+ * {@literal PCollection<Row>} output.
+ *
+ * <p>Use {@link #withMultipleInputs} or {@link #withMultipleOutputs} to
indicate that this
+ * transform has multiple inputs and/or outputs.
+ *
+ * @param yamlDefinition a YAML string defining this transform.
+ * @return a PTransform that applies this YAML to its inputs.
+ */
+ public static YamlTransform<PCollection<Row>, PCollection<Row>> of(String
yamlDefinition) {
+ return new YamlTransform<PCollection<Row>,
PCollection<Row>>(yamlDefinition, null, null);
+ }
+
+ /**
+ * Creates a new YamlTransform PBegin a single {@literal PCollection<Row>}
output.
+ *
+ * @param yamlDefinition a YAML string defining this source.
+ * @return a PTransform that applies this YAML as a root transform.
+ */
+ public static YamlTransform<PBegin, PCollection<Row>> source(String
yamlDefinition) {
+ return new YamlTransform<PBegin, PCollection<Row>>(yamlDefinition, null,
null);
+ }
+
+ /**
+ * Creates a new YamlTransform mapping a single input {@literal
PCollection<Row>} to a single
+ * {@literal PCollection<Row>} output.
+ *
+ * <p>Use {@link #withMultipleOutputs} to indicate that this sink has
multiple (or no) or outputs.
+ *
+ * @param yamlDefinition a YAML string defining this sink.
+ * @return a PTransform that applies this YAML to its inputs.
+ */
+ public static YamlTransform<PCollection<Row>, PCollection<Row>> sink(String
yamlDefinition) {
+ return of(yamlDefinition);
+ }
+
+ /**
+ * Indicates that this YamlTransform expects multiple, named inputs.
+ *
+ * @param inputTags the set of expected input tags to this transform
+ * @return a PTransform like this but with a {@link PCollectionRowTuple}
input type.
+ */
+ public YamlTransform<PCollectionRowTuple, OutputT>
withMultipleInputs(String... inputTags) {
+ return new YamlTransform<PCollectionRowTuple, OutputT>(
+ yamlDefinition, ImmutableSet.copyOf(inputTags), outputTags);
+ }
+
+ /**
+ * Indicates that this YamlTransform expects multiple, named outputs.
+ *
+ * @param outputTags the set of expected output tags to this transform
+ * @return a PTransform like this but with a {@link PCollectionRowTuple}
output type.
+ */
+ public YamlTransform<InputT, PCollectionRowTuple>
withMultipleOutputs(String... outputTags) {
+ return new YamlTransform<InputT, PCollectionRowTuple>(
+ yamlDefinition, inputTags, ImmutableSet.copyOf(outputTags));
+ }
+
+ @Override
+ public OutputT expand(InputT input) {
+ if (inputTags != null) {
+ Set<String> actualInputTags =
+ input.expand().keySet().stream()
+ .map(TupleTag::getId)
+ .collect(Collectors.toCollection(HashSet::new));
+ if (!inputTags.equals(actualInputTags)) {
+ throw new IllegalArgumentException(
+ "Input has tags "
+ + Joiner.on(", ").join(actualInputTags)
+ + " but expected input tags "
+ + Joiner.on(", ").join(inputTags));
+ }
+ }
+
+ // There is no generic apply...
+ POutput output;
+ @SuppressWarnings("rawtypes")
+ PTransform externalTransform =
+
PythonExternalTransform.from("apache_beam.yaml.yaml_transform.YamlTransform")
+ .withArgs(yamlDefinition)
+ .withExtraPackages(ImmutableList.of("jinja2", "pyyaml",
"virtualenv-clone"));
+ if (input instanceof PBegin) {
+ output = ((PBegin) input).apply(externalTransform);
+ } else if (input instanceof PCollection) {
+ output = ((PCollection<?>) input).apply(externalTransform);
+ } else if (input instanceof PCollection) {
+ output = ((PCollection<?>) input).apply(externalTransform);
+ } else if (input instanceof PCollectionRowTuple) {
+ output = ((PCollectionRowTuple) input).apply(externalTransform);
+ } else {
+ throw new IllegalArgumentException("Unrecognized input type: " + input);
+ }
+
+ if (outputTags == null) {
+ if (!(output instanceof PCollection)) {
+ throw new IllegalArgumentException(
+ "Expected a single PCollection output, but got "
+ + output
+ + ". Perhaps withMultipleOutputs() needs to be specified?");
+ }
+ return (OutputT) output;
+ } else {
+ if (output instanceof PCollection) {
+ // ExternalPythonTransform always returns single outputs as
PCollections.
+ if (outputTags.size() != 1) {
+ throw new IllegalArgumentException(
+ "Expected " + outputTags.size() + " outputs, but got exactly
one.");
+ }
+ return (OutputT)
+ PCollectionRowTuple.of(outputTags.iterator().next(),
(PCollection<Row>) output);
+ } else {
+ Map<TupleTag<?>, PValue> expandedOutputs = output.expand();
+ Set<String> actualOutputTags =
+ expandedOutputs.keySet().stream()
+ .map(TupleTag::getId)
+ .collect(Collectors.toCollection(HashSet::new));
+ if (!outputTags.equals(actualOutputTags)) {
+ throw new IllegalArgumentException(
+ "Output has tags "
+ + Joiner.on(", ").join(actualOutputTags)
+ + " but expected output tags "
+ + Joiner.on(", ").join(outputTags));
+ }
+ PCollectionRowTuple result =
PCollectionRowTuple.empty(input.getPipeline());
+ for (Map.Entry<TupleTag<?>, PValue> subOutput :
expandedOutputs.entrySet()) {
+ result = result.and(subOutput.getKey().getId(), (PCollection<Row>)
subOutput.getValue());
+ }
+ return (OutputT) result;
+ }
+ }
+ }
+}
diff --git
a/sdks/java/extensions/yaml/src/main/java/org/apache/beam/sdk/extensions/yaml/package-info.java
b/sdks/java/extensions/yaml/src/main/java/org/apache/beam/sdk/extensions/yaml/package-info.java
new file mode 100644
index 00000000000..d50f2f2a91c
--- /dev/null
+++
b/sdks/java/extensions/yaml/src/main/java/org/apache/beam/sdk/extensions/yaml/package-info.java
@@ -0,0 +1,20 @@
+/*
+ * 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.
+ */
+
+/** Extensions for invoking Beam YAML transforms from the Beam Java SDK. */
+package org.apache.beam.sdk.extensions.yaml;
diff --git
a/sdks/java/extensions/yaml/src/test/java/org/apache/beam/sdk/extensions/yaml/YamlTransformTest.java
b/sdks/java/extensions/yaml/src/test/java/org/apache/beam/sdk/extensions/yaml/YamlTransformTest.java
new file mode 100644
index 00000000000..50e3dfda6a3
--- /dev/null
+++
b/sdks/java/extensions/yaml/src/test/java/org/apache/beam/sdk/extensions/yaml/YamlTransformTest.java
@@ -0,0 +1,143 @@
+/*
+ * 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.beam.sdk.extensions.yaml;
+
+import static org.junit.Assume.assumeTrue;
+
+import java.io.Serializable;
+import org.apache.beam.runners.portability.PortableRunner;
+import org.apache.beam.sdk.options.PipelineOptions;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.options.PortablePipelineOptions;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.testing.PAssert;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionRowTuple;
+import org.apache.beam.sdk.values.Row;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class YamlTransformTest implements Serializable {
+
+ // See build.gradle for test task configuration.
+ private static final String PORTABLE_RUNNER_PROPERTY_NAME =
"testRunner.jobEndpoint";
+
+ private static final Schema INPUT_SCHEMA =
+ Schema.of(
+ Schema.Field.of("name", Schema.FieldType.STRING),
+ Schema.Field.of("email", Schema.FieldType.STRING));
+ private static final Row ALICE =
+
Row.withSchema(INPUT_SCHEMA).addValue("alice").addValue("[email protected]").build();
+ private static final Row BOB =
+
Row.withSchema(INPUT_SCHEMA).addValue("bob").addValue("[email protected]").build();
+
+ @Rule public TestPipeline pipeline = TestPipeline.fromOptions(options());
+
+ @Test
+ public void simpleYamlTransform() {
+ PCollection<Row> output =
+ pipeline
+ .apply(Create.of(ALICE, BOB))
+ .setRowSchema(INPUT_SCHEMA)
+ .apply(
+ YamlTransform.of(
+ "{type: Filter, config: {language: python, keep:
'len(name) < 4'}}"));
+ PAssert.that(output).containsInAnyOrder(BOB);
+
+ pipeline.run().waitUntilFinish();
+ }
+
+ @Test
+ public void multipleOutputYamlTransform() {
+ PCollectionRowTuple output =
+ pipeline
+ .apply(Create.of(ALICE, BOB))
+ .setRowSchema(INPUT_SCHEMA)
+ .apply(
+ // TODO(https://github.com/apache/beam/pull/34595): Use
Partition transform.
+ // This needs to wait until the above change is in a released
version of Beam.
+ // YamlTransform.of(
+ // "{type: Partition, "
+ // + " config: {language: python, by: 'name[0]',
outputs: [a, b]}}")
+ // .withMultipleOutputs("a", "b"));
+ YamlTransform.of(
+ "type: composite \n"
+ + "input: input \n"
+ + "transforms: [ \n"
+ + " {type: Filter, name: A, input: input, "
+ + " config: {language: python, keep: 'name[0]
== \"a\"'}}, \n"
+ + " {type: Filter, name: B, input: input, "
+ + " config: {language: python, keep: 'name[0]
== \"b\"'}} \n"
+ + "] \n"
+ + "output: {a: A, b: B} \n"
+ + "")
+ .withMultipleOutputs("a", "b"));
+ PAssert.that(output.get("a")).containsInAnyOrder(ALICE);
+ PAssert.that(output.get("b")).containsInAnyOrder(BOB);
+
+ pipeline.run().waitUntilFinish();
+ }
+
+ @Test
+ public void yamlSource() {
+ PCollection<Row> output =
+ pipeline.apply(
+ YamlTransform.source(
+ "{type: Create, "
+ + " config: {elements: [{name: alice, email:
[email protected]}]}}"));
+ PAssert.that(output).containsInAnyOrder(ALICE);
+
+ // TODO: Run this on a multi-language supporting runner.
+ pipeline.run().waitUntilFinish();
+ }
+
+ private static PipelineOptions options() {
+ PortablePipelineOptions opts =
+ PipelineOptionsFactory.create().as(PortablePipelineOptions.class);
+
+ opts.setRunner(PortableRunner.class);
+ // TODO(https://github.com/apache/beam/issues/34594):Use LOOPBACK here.
+ opts.setDefaultEnvironmentType("DOCKER");
+ opts.setJobEndpoint(getLocalRunnerAddress());
+
+ return opts;
+ }
+
+ /**
+ * Drives ignoring of tests via checking {@link org.junit.Assume#assumeTrue}
that the {@link
+ * System#getProperty} for {@link #PORTABLE_RUNNER_PROPERTY_NAME} is not
null or empty.
+ */
+ private static String getLocalRunnerAddress() {
+ String address = System.getProperty(PORTABLE_RUNNER_PROPERTY_NAME);
+ assumeTrue(
+ "System property: "
+ + PORTABLE_RUNNER_PROPERTY_NAME
+ + " is not set; start a local runner and pass "
+ + "-D"
+ + PORTABLE_RUNNER_PROPERTY_NAME
+ + "=localhost:port",
+ !Strings.isNullOrEmpty(address));
+ return address;
+ }
+}
diff --git a/settings.gradle.kts b/settings.gradle.kts
index 4ca8d32d9e9..4e71a330856 100644
--- a/settings.gradle.kts
+++ b/settings.gradle.kts
@@ -205,6 +205,7 @@ include(":sdks:java:extensions:sql:expansion-service")
include(":sdks:java:extensions:sql:udf")
include(":sdks:java:extensions:sql:udf-test-provider")
include(":sdks:java:extensions:timeseries")
+include(":sdks:java:extensions:yaml")
include(":sdks:java:extensions:zetasketch")
include(":sdks:java:harness")
include(":sdks:java:harness:jmh")
