robertwb commented on a change in pull request #17101:
URL: https://github.com/apache/beam/pull/17101#discussion_r840092488
##########
File path:
sdks/java/extensions/python/src/main/java/org/apache/beam/sdk/extensions/python/ExternalPythonTransform.java
##########
@@ -33,25 +43,199 @@
import org.apache.beam.sdk.values.POutput;
import org.apache.beam.sdk.values.Row;
import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TypeDescriptor;
import org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString;
+import
org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
import
org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
+import org.checkerframework.checker.nullness.qual.Nullable;
/** Wrapper for invoking external Python transforms. */
public class ExternalPythonTransform<InputT extends PInput, OutputT extends
POutput>
extends PTransform<InputT, OutputT> {
- private final String fullyQualifiedName;
- private final Row args;
- private final Row kwargs;
- public ExternalPythonTransform(String fullyQualifiedName, Row args, Row
kwargs) {
+ private static final SchemaRegistry SCHEMA_REGISTRY =
SchemaRegistry.createDefault();
+ private String fullyQualifiedName;
+
+ // We preseve the order here since Schema's care about order of fields but
the order will not
+ // matter when applying kwargs at the Python side.
+ private SortedMap<String, Object> kwargsMap;
+
+ private Object[] argsArray;
+ private Row providedKwargsRow = null;
+
+ private ExternalPythonTransform(String fullyQualifiedName) {
this.fullyQualifiedName = fullyQualifiedName;
- this.args = args;
- this.kwargs = kwargs;
+ this.kwargsMap = new TreeMap<>();
+ this.argsArray = new Object[] {};
+ }
+
+ /**
+ * Instantiates a cross-language wrapper for a Python transform with a given
transform name.
+ *
+ * @param tranformName fully qualified transform name.
+ * @param <InputT> Input {@link PCollection} type
+ * @param <OutputT> Output {@link PCollection} type
+ * @return A {@link ExternalPythonTransform} for the given transform name.
+ */
+ public static <InputT extends PInput, OutputT extends POutput>
+ ExternalPythonTransform<InputT, OutputT> from(String tranformName) {
+ return new ExternalPythonTransform<InputT, OutputT>(tranformName);
+ }
+
+ /**
+ * Positional arguments for the Python cross-language transform. If invoked
more than once, new
+ * arguments will be appended to the previously specified arguments.
+ *
+ * @param args list of arguments.
+ * @return updated wrapper for the cross-language transform.
+ */
+ public ExternalPythonTransform<InputT, OutputT> withArgs(Object... args) {
+ Object[] result = Arrays.copyOf(this.argsArray, this.argsArray.length +
args.length);
+ System.arraycopy(args, 0, result, this.argsArray.length, args.length);
+ this.argsArray = result;
+ return this;
+ }
+
+ /**
+ * Specifies a single keyword argument for the Python cross-language
transform. This may be
+ * invoked multiple times to add more than one keyword argument.
+ *
+ * @param name argument name.
+ * @param value argument value
+ * @return updated wrapper for the cross-language transform.
+ */
+ public ExternalPythonTransform<InputT, OutputT> withKwarg(String name,
Object value) {
+ if (providedKwargsRow != null) {
+ throw new IllegalArgumentException("Kwargs were specified both directly
and as a Row object");
+ }
+ kwargsMap.put(name, value);
+ return this;
+ }
+
+ /**
+ * Specifies keyword arguments for the Python cross-language transform. If
invoked more than once,
+ * new keyword arguments map will be added to the previously prided keyword
arguments.
+ *
+ * @return updated wrapper for the cross-language transform.
+ */
+ public ExternalPythonTransform<InputT, OutputT> withKwargs(Map<String,
Object> kwargs) {
+ kwargsMap.putAll(kwargs);
+ return this;
+ }
+
+ /**
+ * Specifies keyword arguments as a Row objects.
+ *
+ * @param kwargs keyword arguments as a {@link Row} objects. An empty Row
represents zero keyword
+ * arguments.
+ * @return updated wrapper for the cross-language transform.
+ */
+ public ExternalPythonTransform<InputT, OutputT> withKwargs(Row kwargs) {
+ if (this.kwargsMap.size() > 0) {
+ throw new IllegalArgumentException("Kwargs were specified both directly
and as a Row object");
+ }
+ this.providedKwargsRow = kwargs;
+ return this;
+ }
+
+ @VisibleForTesting
+ Row buildOrGetKvargsRow() {
+ if (providedKwargsRow != null) {
+ return providedKwargsRow;
+ } else if (kwargsMap.size() == 0) {
+ return Row.nullRow(Schema.of());
+ } else {
+ Schema schema =
+ generateSchemaFromFieldValues(
+ kwargsMap.values().toArray(), kwargsMap.keySet().toArray(new
String[] {}));
+ return Row.withSchema(schema)
+ .addValues(convertComplexTypesToRows(kwargsMap.values().toArray()))
+ .build();
+ }
+ }
+
+ private static boolean isCustomType(java.lang.Class<?> type) {
+ boolean val =
+ !(ClassUtils.isPrimitiveOrWrapper(type)
+ || type == String.class
+ || Row.class.isAssignableFrom(type));
+ return val;
+ }
+
+ // If the custom type has a registered schema, we use that. OTherwise we try
to register it using
+ // 'JavaFieldSchema'.
+ private Row convertCustomValue(Object value) {
+ SerializableFunction<Object, Row> toRowFunc;
+ try {
+ toRowFunc =
+ (SerializableFunction<Object, Row>)
SCHEMA_REGISTRY.getToRowFunction(value.getClass());
+ } catch (NoSuchSchemaException e) {
+ SCHEMA_REGISTRY.registerSchemaProvider(value.getClass(), new
JavaFieldSchema());
+ try {
+ toRowFunc =
+ (SerializableFunction<Object, Row>)
SCHEMA_REGISTRY.getToRowFunction(value.getClass());
+ } catch (NoSuchSchemaException e1) {
+ throw new RuntimeException(e1);
+ }
+ }
+ return toRowFunc.apply(value);
+ }
+
+ private Object[] convertComplexTypesToRows(Object[] values) {
+ Object[] converted = new Object[values.length];
+ for (int i = 0; i < values.length; i++) {
+ Object value = values[i];
+ converted[i] = isCustomType(value.getClass()) ?
convertCustomValue(value) : value;
+ }
+ return converted;
+ }
+
+ @VisibleForTesting
+ Row buildOrGetArgsRow() {
+ if (argsArray.length == 0) {
Review comment:
Ideally generateSchemaFromFieldValues should handle the degenerate empty
case just fine, so we don't need to add a special case. (Similarly above for
the kwargs.)
##########
File path:
sdks/java/extensions/python/src/main/java/org/apache/beam/sdk/extensions/python/ExternalPythonTransform.java
##########
@@ -33,25 +43,199 @@
import org.apache.beam.sdk.values.POutput;
import org.apache.beam.sdk.values.Row;
import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TypeDescriptor;
import org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString;
+import
org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
import
org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
+import org.checkerframework.checker.nullness.qual.Nullable;
/** Wrapper for invoking external Python transforms. */
public class ExternalPythonTransform<InputT extends PInput, OutputT extends
POutput>
extends PTransform<InputT, OutputT> {
- private final String fullyQualifiedName;
- private final Row args;
- private final Row kwargs;
- public ExternalPythonTransform(String fullyQualifiedName, Row args, Row
kwargs) {
+ private static final SchemaRegistry SCHEMA_REGISTRY =
SchemaRegistry.createDefault();
+ private String fullyQualifiedName;
+
+ // We preseve the order here since Schema's care about order of fields but
the order will not
+ // matter when applying kwargs at the Python side.
+ private SortedMap<String, Object> kwargsMap;
+
+ private Object[] argsArray;
+ private Row providedKwargsRow = null;
+
+ private ExternalPythonTransform(String fullyQualifiedName) {
this.fullyQualifiedName = fullyQualifiedName;
- this.args = args;
- this.kwargs = kwargs;
+ this.kwargsMap = new TreeMap<>();
+ this.argsArray = new Object[] {};
+ }
+
+ /**
+ * Instantiates a cross-language wrapper for a Python transform with a given
transform name.
+ *
+ * @param tranformName fully qualified transform name.
+ * @param <InputT> Input {@link PCollection} type
+ * @param <OutputT> Output {@link PCollection} type
+ * @return A {@link ExternalPythonTransform} for the given transform name.
+ */
+ public static <InputT extends PInput, OutputT extends POutput>
+ ExternalPythonTransform<InputT, OutputT> from(String tranformName) {
+ return new ExternalPythonTransform<InputT, OutputT>(tranformName);
+ }
+
+ /**
+ * Positional arguments for the Python cross-language transform. If invoked
more than once, new
+ * arguments will be appended to the previously specified arguments.
+ *
+ * @param args list of arguments.
+ * @return updated wrapper for the cross-language transform.
+ */
+ public ExternalPythonTransform<InputT, OutputT> withArgs(Object... args) {
+ Object[] result = Arrays.copyOf(this.argsArray, this.argsArray.length +
args.length);
+ System.arraycopy(args, 0, result, this.argsArray.length, args.length);
+ this.argsArray = result;
+ return this;
+ }
+
+ /**
+ * Specifies a single keyword argument for the Python cross-language
transform. This may be
+ * invoked multiple times to add more than one keyword argument.
+ *
+ * @param name argument name.
+ * @param value argument value
+ * @return updated wrapper for the cross-language transform.
+ */
+ public ExternalPythonTransform<InputT, OutputT> withKwarg(String name,
Object value) {
+ if (providedKwargsRow != null) {
+ throw new IllegalArgumentException("Kwargs were specified both directly
and as a Row object");
+ }
+ kwargsMap.put(name, value);
+ return this;
+ }
+
+ /**
+ * Specifies keyword arguments for the Python cross-language transform. If
invoked more than once,
+ * new keyword arguments map will be added to the previously prided keyword
arguments.
+ *
+ * @return updated wrapper for the cross-language transform.
+ */
+ public ExternalPythonTransform<InputT, OutputT> withKwargs(Map<String,
Object> kwargs) {
+ kwargsMap.putAll(kwargs);
+ return this;
+ }
+
+ /**
+ * Specifies keyword arguments as a Row objects.
+ *
+ * @param kwargs keyword arguments as a {@link Row} objects. An empty Row
represents zero keyword
+ * arguments.
+ * @return updated wrapper for the cross-language transform.
+ */
+ public ExternalPythonTransform<InputT, OutputT> withKwargs(Row kwargs) {
+ if (this.kwargsMap.size() > 0) {
+ throw new IllegalArgumentException("Kwargs were specified both directly
and as a Row object");
+ }
+ this.providedKwargsRow = kwargs;
+ return this;
+ }
+
+ @VisibleForTesting
+ Row buildOrGetKvargsRow() {
+ if (providedKwargsRow != null) {
+ return providedKwargsRow;
+ } else if (kwargsMap.size() == 0) {
+ return Row.nullRow(Schema.of());
+ } else {
+ Schema schema =
+ generateSchemaFromFieldValues(
+ kwargsMap.values().toArray(), kwargsMap.keySet().toArray(new
String[] {}));
+ return Row.withSchema(schema)
+ .addValues(convertComplexTypesToRows(kwargsMap.values().toArray()))
+ .build();
+ }
+ }
+
+ private static boolean isCustomType(java.lang.Class<?> type) {
Review comment:
I'm not sure what is meant by "custom" here.
##########
File path:
sdks/java/extensions/python/src/main/java/org/apache/beam/sdk/extensions/python/ExternalPythonTransform.java
##########
@@ -33,25 +43,199 @@
import org.apache.beam.sdk.values.POutput;
import org.apache.beam.sdk.values.Row;
import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TypeDescriptor;
import org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString;
+import
org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
import
org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
+import org.checkerframework.checker.nullness.qual.Nullable;
/** Wrapper for invoking external Python transforms. */
public class ExternalPythonTransform<InputT extends PInput, OutputT extends
POutput>
extends PTransform<InputT, OutputT> {
- private final String fullyQualifiedName;
- private final Row args;
- private final Row kwargs;
- public ExternalPythonTransform(String fullyQualifiedName, Row args, Row
kwargs) {
+ private static final SchemaRegistry SCHEMA_REGISTRY =
SchemaRegistry.createDefault();
+ private String fullyQualifiedName;
+
+ // We preseve the order here since Schema's care about order of fields but
the order will not
+ // matter when applying kwargs at the Python side.
+ private SortedMap<String, Object> kwargsMap;
+
+ private Object[] argsArray;
+ private Row providedKwargsRow = null;
+
+ private ExternalPythonTransform(String fullyQualifiedName) {
this.fullyQualifiedName = fullyQualifiedName;
- this.args = args;
- this.kwargs = kwargs;
+ this.kwargsMap = new TreeMap<>();
+ this.argsArray = new Object[] {};
+ }
+
+ /**
+ * Instantiates a cross-language wrapper for a Python transform with a given
transform name.
+ *
+ * @param tranformName fully qualified transform name.
+ * @param <InputT> Input {@link PCollection} type
+ * @param <OutputT> Output {@link PCollection} type
+ * @return A {@link ExternalPythonTransform} for the given transform name.
+ */
+ public static <InputT extends PInput, OutputT extends POutput>
+ ExternalPythonTransform<InputT, OutputT> from(String tranformName) {
+ return new ExternalPythonTransform<InputT, OutputT>(tranformName);
+ }
+
+ /**
+ * Positional arguments for the Python cross-language transform. If invoked
more than once, new
+ * arguments will be appended to the previously specified arguments.
+ *
+ * @param args list of arguments.
+ * @return updated wrapper for the cross-language transform.
+ */
+ public ExternalPythonTransform<InputT, OutputT> withArgs(Object... args) {
+ Object[] result = Arrays.copyOf(this.argsArray, this.argsArray.length +
args.length);
+ System.arraycopy(args, 0, result, this.argsArray.length, args.length);
+ this.argsArray = result;
+ return this;
+ }
+
+ /**
+ * Specifies a single keyword argument for the Python cross-language
transform. This may be
+ * invoked multiple times to add more than one keyword argument.
+ *
+ * @param name argument name.
+ * @param value argument value
+ * @return updated wrapper for the cross-language transform.
+ */
+ public ExternalPythonTransform<InputT, OutputT> withKwarg(String name,
Object value) {
+ if (providedKwargsRow != null) {
+ throw new IllegalArgumentException("Kwargs were specified both directly
and as a Row object");
+ }
+ kwargsMap.put(name, value);
+ return this;
+ }
+
+ /**
+ * Specifies keyword arguments for the Python cross-language transform. If
invoked more than once,
+ * new keyword arguments map will be added to the previously prided keyword
arguments.
+ *
+ * @return updated wrapper for the cross-language transform.
+ */
+ public ExternalPythonTransform<InputT, OutputT> withKwargs(Map<String,
Object> kwargs) {
+ kwargsMap.putAll(kwargs);
+ return this;
+ }
+
+ /**
+ * Specifies keyword arguments as a Row objects.
+ *
+ * @param kwargs keyword arguments as a {@link Row} objects. An empty Row
represents zero keyword
+ * arguments.
+ * @return updated wrapper for the cross-language transform.
+ */
+ public ExternalPythonTransform<InputT, OutputT> withKwargs(Row kwargs) {
+ if (this.kwargsMap.size() > 0) {
+ throw new IllegalArgumentException("Kwargs were specified both directly
and as a Row object");
+ }
+ this.providedKwargsRow = kwargs;
+ return this;
+ }
+
+ @VisibleForTesting
+ Row buildOrGetKvargsRow() {
Review comment:
kvargs -> kwargs (v -> w)
##########
File path:
sdks/java/extensions/python/src/main/java/org/apache/beam/sdk/extensions/python/ExternalPythonTransform.java
##########
@@ -33,25 +43,199 @@
import org.apache.beam.sdk.values.POutput;
import org.apache.beam.sdk.values.Row;
import org.apache.beam.sdk.values.TupleTag;
+import org.apache.beam.sdk.values.TypeDescriptor;
import org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString;
+import
org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
import
org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
+import org.checkerframework.checker.nullness.qual.Nullable;
/** Wrapper for invoking external Python transforms. */
public class ExternalPythonTransform<InputT extends PInput, OutputT extends
POutput>
extends PTransform<InputT, OutputT> {
- private final String fullyQualifiedName;
- private final Row args;
- private final Row kwargs;
- public ExternalPythonTransform(String fullyQualifiedName, Row args, Row
kwargs) {
+ private static final SchemaRegistry SCHEMA_REGISTRY =
SchemaRegistry.createDefault();
+ private String fullyQualifiedName;
+
+ // We preseve the order here since Schema's care about order of fields but
the order will not
+ // matter when applying kwargs at the Python side.
+ private SortedMap<String, Object> kwargsMap;
+
+ private Object[] argsArray;
+ private Row providedKwargsRow = null;
+
+ private ExternalPythonTransform(String fullyQualifiedName) {
this.fullyQualifiedName = fullyQualifiedName;
- this.args = args;
- this.kwargs = kwargs;
+ this.kwargsMap = new TreeMap<>();
+ this.argsArray = new Object[] {};
+ }
+
+ /**
+ * Instantiates a cross-language wrapper for a Python transform with a given
transform name.
+ *
+ * @param tranformName fully qualified transform name.
+ * @param <InputT> Input {@link PCollection} type
+ * @param <OutputT> Output {@link PCollection} type
+ * @return A {@link ExternalPythonTransform} for the given transform name.
+ */
+ public static <InputT extends PInput, OutputT extends POutput>
+ ExternalPythonTransform<InputT, OutputT> from(String tranformName) {
+ return new ExternalPythonTransform<InputT, OutputT>(tranformName);
+ }
+
+ /**
+ * Positional arguments for the Python cross-language transform. If invoked
more than once, new
+ * arguments will be appended to the previously specified arguments.
+ *
+ * @param args list of arguments.
+ * @return updated wrapper for the cross-language transform.
+ */
+ public ExternalPythonTransform<InputT, OutputT> withArgs(Object... args) {
+ Object[] result = Arrays.copyOf(this.argsArray, this.argsArray.length +
args.length);
+ System.arraycopy(args, 0, result, this.argsArray.length, args.length);
+ this.argsArray = result;
+ return this;
+ }
+
+ /**
+ * Specifies a single keyword argument for the Python cross-language
transform. This may be
+ * invoked multiple times to add more than one keyword argument.
+ *
+ * @param name argument name.
+ * @param value argument value
+ * @return updated wrapper for the cross-language transform.
+ */
+ public ExternalPythonTransform<InputT, OutputT> withKwarg(String name,
Object value) {
+ if (providedKwargsRow != null) {
+ throw new IllegalArgumentException("Kwargs were specified both directly
and as a Row object");
+ }
+ kwargsMap.put(name, value);
+ return this;
+ }
+
+ /**
+ * Specifies keyword arguments for the Python cross-language transform. If
invoked more than once,
+ * new keyword arguments map will be added to the previously prided keyword
arguments.
+ *
+ * @return updated wrapper for the cross-language transform.
+ */
+ public ExternalPythonTransform<InputT, OutputT> withKwargs(Map<String,
Object> kwargs) {
+ kwargsMap.putAll(kwargs);
Review comment:
Should we check that `providedKwargsRow != null` here too?
##########
File path:
sdks/java/extensions/python/src/test/java/org/apache/beam/sdk/extensions/python/ExternalPythonTransformTest.java
##########
@@ -27,23 +30,218 @@
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.Row;
import org.apache.beam.sdk.values.TypeDescriptors;
+import
org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
@RunWith(JUnit4.class)
public class ExternalPythonTransformTest implements Serializable {
+
@Test
public void trivialPythonTransform() {
Pipeline p = Pipeline.create();
PCollection<String> output =
p.apply(Create.of(KV.of("A", "x"), KV.of("A", "y"), KV.of("B", "z")))
.apply(
- new ExternalPythonTransform<
- PCollection<KV<String, String>>, PCollection<KV<String,
Iterable<String>>>>(
- "apache_beam.GroupByKey", Row.nullRow(Schema.of()),
Row.nullRow(Schema.of())))
+ ExternalPythonTransform
+ .<PCollection<KV<String, String>>, PCollection<KV<String,
Iterable<String>>>>
+ from("apache_beam.GroupByKey"))
.apply(MapElements.into(TypeDescriptors.strings()).via(kv ->
kv.getKey()));
PAssert.that(output).containsInAnyOrder("A", "B");
// TODO: Run this on a multi-language supporting runner.
}
+
+ @Test
+ public void generateArgsWithPrimitives() {
+ ExternalPythonTransform<?, ?> transform =
+ ExternalPythonTransform
+ .<PCollection<KV<String, String>>, PCollection<KV<String,
Iterable<String>>>>from(
+ "DummyTransform")
+ .withArgs("aaa", "bbb", 11, 12L, 15.6, true);
+
+ Schema expectedSchema =
+ Schema.builder()
+ .addStringField("field0")
+ .addStringField("field1")
+ .addInt32Field("field2")
+ .addInt64Field("field3")
+ .addDoubleField("field4")
+ .addBooleanField("field5")
+ .build();
+ Row expectedRow =
+ Row.withSchema(expectedSchema).addValues("aaa", "bbb", 11, 12L, 15.6,
true).build();
+
+ Row receivedRow = transform.buildOrGetArgsRow();
+ assertEquals(expectedRow, receivedRow);
+ }
+
+ @Test
+ public void generateArgsWithRow() {
+ Schema subRowSchema1 =
+
Schema.builder().addStringField("field0").addInt32Field("field1").build();
+ Row rowField1 = Row.withSchema(subRowSchema1).addValues("xxx",
123).build();
+ Schema subRowSchema2 =
+ Schema.builder()
+ .addDoubleField("field0")
+ .addBooleanField("field1")
+ .addStringField("field2")
+ .build();
+ Row rowField2 = Row.withSchema(subRowSchema2).addValues(12.5, true,
"yyy").build();
+
+ ExternalPythonTransform<?, ?> transform =
+ ExternalPythonTransform
+ .<PCollection<KV<String, String>>, PCollection<KV<String,
Iterable<String>>>>from(
+ "DummyTransform")
+ .withArgs(rowField1, rowField2);
+
+ Schema expectedSchema =
+ Schema.builder()
+ .addRowField("field0", subRowSchema1)
+ .addRowField("field1", subRowSchema2)
+ .build();
+ Row expectedRow = Row.withSchema(expectedSchema).addValues(rowField1,
rowField2).build();
+
+ Row receivedRow = transform.buildOrGetArgsRow();
+ assertEquals(expectedRow, receivedRow);
+ }
+
+ static class CustomType {
+
Review comment:
Remove newline (unless spotless insists).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
