You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by sj...@apache.org on 2020/09/14 17:48:37 UTC
[flink] 03/05: [FLINK-18735][table] Add support for more types to
DataGen source
This is an automated email from the ASF dual-hosted git repository.
sjwiesman pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit 8f7369eec0f3d1779271e9fa35718a4da9bff8cb
Author: Seth Wiesman <sj...@gmail.com>
AuthorDate: Tue Jul 28 11:51:27 2020 -0500
[FLINK-18735][table] Add support for more types to DataGen source
---
.../functions/source/datagen/RandomGenerator.java | 46 +++++
.../source/datagen/SequenceGenerator.java | 13 ++
.../table/factories/DataGenTableSourceFactory.java | 4 +-
.../factories/datagen/DataGenVisitorBase.java | 110 ++++++++++
.../factories/datagen/RandomGeneratorVisitor.java | 224 +++++++++++++++++++--
.../datagen/SequenceGeneratorVisitor.java | 24 ++-
.../factories/DataGenTableSourceFactoryTest.java | 128 ++++++++----
7 files changed, 482 insertions(+), 67 deletions(-)
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/functions/source/datagen/RandomGenerator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/functions/source/datagen/RandomGenerator.java
index 6bc795f..4f48a9c 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/functions/source/datagen/RandomGenerator.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/functions/source/datagen/RandomGenerator.java
@@ -24,6 +24,9 @@ import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.commons.math3.random.RandomDataGenerator;
+import java.util.HashMap;
+import java.util.Map;
+
/**
* Random generator.
*/
@@ -116,4 +119,47 @@ public abstract class RandomGenerator<T> implements DataGenerator<T> {
}
};
}
+
+ public static <T> RandomGenerator<T[]> arrayGenerator(DataGenerator<T> generator, int len) {
+ return new RandomGenerator<T[]>() {
+
+ @Override
+ public void open(String name, FunctionInitializationContext context, RuntimeContext runtimeContext) throws Exception {
+ generator.open(name, context, runtimeContext);
+ }
+
+ @Override
+ public T[] next() {
+ @SuppressWarnings("unchecked")
+ T[] array = (T[]) new Object[len];
+
+ for (int i = 0; i < len; i++) {
+ array[i] = generator.next();
+ }
+
+ return array;
+ }
+ };
+ }
+
+ public static <K, V> RandomGenerator<Map<K, V>> mapGenerator(DataGenerator<K> key, DataGenerator<V> value, int size) {
+ return new RandomGenerator<Map<K, V>>() {
+
+ @Override
+ public void open(String name, FunctionInitializationContext context, RuntimeContext runtimeContext) throws Exception {
+ key.open(name + ".key", context, runtimeContext);
+ value.open(name + ".value", context, runtimeContext);
+ }
+
+ @Override
+ public Map<K, V> next() {
+ Map<K, V> map = new HashMap<>(size);
+ for (int i = 0; i < size; i++) {
+ map.put(key.next(), value.next());
+ }
+
+ return map;
+ }
+ };
+ }
}
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/functions/source/datagen/SequenceGenerator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/functions/source/datagen/SequenceGenerator.java
index 855937e..50f4c13 100644
--- a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/functions/source/datagen/SequenceGenerator.java
+++ b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/functions/source/datagen/SequenceGenerator.java
@@ -27,6 +27,9 @@ import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.util.Preconditions;
+import java.math.BigDecimal;
+import java.math.MathContext;
+import java.math.RoundingMode;
import java.util.ArrayDeque;
import java.util.Deque;
@@ -166,6 +169,16 @@ public abstract class SequenceGenerator<T> implements DataGenerator<T> {
};
}
+ public static SequenceGenerator<BigDecimal> bigDecimalGenerator(int start, int end, int precision, int scale) {
+ return new SequenceGenerator<BigDecimal>(start, end) {
+ @Override
+ public BigDecimal next() {
+ BigDecimal decimal = new BigDecimal(valuesToEmit.poll().doubleValue(), new MathContext(precision));
+ return decimal.setScale(scale, RoundingMode.DOWN);
+ }
+ };
+ }
+
public static SequenceGenerator<String> stringGenerator(long start, long end) {
return new SequenceGenerator<String>(start, end) {
@Override
diff --git a/flink-table/flink-table-api-java-bridge/src/main/java/org/apache/flink/table/factories/DataGenTableSourceFactory.java b/flink-table/flink-table-api-java-bridge/src/main/java/org/apache/flink/table/factories/DataGenTableSourceFactory.java
index c3ad2b9..a2b9df3 100644
--- a/flink-table/flink-table-api-java-bridge/src/main/java/org/apache/flink/table/factories/DataGenTableSourceFactory.java
+++ b/flink-table/flink-table-api-java-bridge/src/main/java/org/apache/flink/table/factories/DataGenTableSourceFactory.java
@@ -125,9 +125,9 @@ public class DataGenTableSourceFactory implements DynamicTableSourceFactory {
String name, DataType type, String kind, ReadableConfig options) {
switch (kind) {
case RANDOM:
- return type.getLogicalType().accept(new RandomGeneratorVisitor(name, type, options));
+ return type.getLogicalType().accept(new RandomGeneratorVisitor(name, options));
case SEQUENCE:
- return type.getLogicalType().accept(new SequenceGeneratorVisitor(name, type, options));
+ return type.getLogicalType().accept(new SequenceGeneratorVisitor(name, options));
default:
throw new ValidationException("Unsupported generator kind: " + kind);
}
diff --git a/flink-table/flink-table-api-java-bridge/src/main/java/org/apache/flink/table/factories/datagen/DataGenVisitorBase.java b/flink-table/flink-table-api-java-bridge/src/main/java/org/apache/flink/table/factories/datagen/DataGenVisitorBase.java
new file mode 100644
index 0000000..a34c91f
--- /dev/null
+++ b/flink-table/flink-table-api-java-bridge/src/main/java/org/apache/flink/table/factories/datagen/DataGenVisitorBase.java
@@ -0,0 +1,110 @@
+/*
+ * 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.flink.table.factories.datagen;
+
+import org.apache.flink.api.common.functions.RuntimeContext;
+import org.apache.flink.configuration.ReadableConfig;
+import org.apache.flink.runtime.state.FunctionInitializationContext;
+import org.apache.flink.streaming.api.functions.source.datagen.DataGenerator;
+import org.apache.flink.table.api.ValidationException;
+import org.apache.flink.table.types.logical.DateType;
+import org.apache.flink.table.types.logical.LocalZonedTimestampType;
+import org.apache.flink.table.types.logical.LogicalType;
+import org.apache.flink.table.types.logical.TimeType;
+import org.apache.flink.table.types.logical.TimestampType;
+import org.apache.flink.table.types.logical.ZonedTimestampType;
+import org.apache.flink.table.types.logical.utils.LogicalTypeDefaultVisitor;
+
+import java.io.Serializable;
+import java.time.Instant;
+import java.time.LocalDate;
+import java.time.LocalDateTime;
+import java.time.LocalTime;
+import java.time.OffsetDateTime;
+import java.util.function.Supplier;
+
+/**
+ * Base class for translating {@link LogicalType LogicalTypes} to {@link DataGeneratorContainer}'s.
+ */
+public abstract class DataGenVisitorBase extends LogicalTypeDefaultVisitor<DataGeneratorContainer> {
+
+ protected final String name;
+
+ protected final ReadableConfig config;
+
+ protected DataGenVisitorBase(String name, ReadableConfig config) {
+ this.name = name;
+ this.config = config;
+ }
+
+ @Override
+ public DataGeneratorContainer visit(DateType dateType) {
+ return DataGeneratorContainer.of(TimeGenerator.of(LocalDate::now));
+ }
+
+ @Override
+ public DataGeneratorContainer visit(TimeType timeType) {
+ return DataGeneratorContainer.of(TimeGenerator.of(LocalTime::now));
+ }
+
+ @Override
+ public DataGeneratorContainer visit(TimestampType timestampType) {
+ return DataGeneratorContainer.of(TimeGenerator.of(LocalDateTime::now));
+ }
+
+ @Override
+ public DataGeneratorContainer visit(ZonedTimestampType zonedTimestampType) {
+ return DataGeneratorContainer.of(TimeGenerator.of(OffsetDateTime::now));
+ }
+
+ @Override
+ public DataGeneratorContainer visit(LocalZonedTimestampType localZonedTimestampType) {
+ return DataGeneratorContainer.of(TimeGenerator.of(Instant::now));
+ }
+
+ @Override
+ protected DataGeneratorContainer defaultMethod(LogicalType logicalType) {
+ throw new ValidationException("Unsupported type: " + logicalType);
+ }
+
+ private interface SerializableSupplier<T> extends Supplier<T>, Serializable { }
+
+ private abstract static class TimeGenerator<T> implements DataGenerator<T> {
+
+ public static <T> TimeGenerator<T> of(SerializableSupplier<T> supplier) {
+ return new TimeGenerator<T>() {
+ @Override
+ public T next() {
+ return supplier.get();
+ }
+ };
+ }
+
+ @Override
+ public void open(
+ String name,
+ FunctionInitializationContext context,
+ RuntimeContext runtimeContext) throws Exception { }
+
+ @Override
+ public boolean hasNext() {
+ return true;
+ }
+ }
+}
diff --git a/flink-table/flink-table-api-java-bridge/src/main/java/org/apache/flink/table/factories/datagen/RandomGeneratorVisitor.java b/flink-table/flink-table-api-java-bridge/src/main/java/org/apache/flink/table/factories/datagen/RandomGeneratorVisitor.java
index c203114..97fafe4 100644
--- a/flink-table/flink-table-api-java-bridge/src/main/java/org/apache/flink/table/factories/datagen/RandomGeneratorVisitor.java
+++ b/flink-table/flink-table-api-java-bridge/src/main/java/org/apache/flink/table/factories/datagen/RandomGeneratorVisitor.java
@@ -19,50 +19,65 @@
package org.apache.flink.table.factories.datagen;
import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.functions.RuntimeContext;
import org.apache.flink.configuration.ConfigOption;
import org.apache.flink.configuration.ConfigOptions;
import org.apache.flink.configuration.ReadableConfig;
+import org.apache.flink.runtime.state.FunctionInitializationContext;
+import org.apache.flink.streaming.api.functions.source.datagen.DataGenerator;
import org.apache.flink.streaming.api.functions.source.datagen.RandomGenerator;
import org.apache.flink.table.api.ValidationException;
import org.apache.flink.table.data.StringData;
-import org.apache.flink.table.types.DataType;
+import org.apache.flink.table.types.logical.ArrayType;
import org.apache.flink.table.types.logical.BigIntType;
import org.apache.flink.table.types.logical.BooleanType;
import org.apache.flink.table.types.logical.CharType;
+import org.apache.flink.table.types.logical.DayTimeIntervalType;
+import org.apache.flink.table.types.logical.DecimalType;
import org.apache.flink.table.types.logical.DoubleType;
import org.apache.flink.table.types.logical.FloatType;
import org.apache.flink.table.types.logical.IntType;
import org.apache.flink.table.types.logical.LogicalType;
+import org.apache.flink.table.types.logical.MapType;
+import org.apache.flink.table.types.logical.MultisetType;
+import org.apache.flink.table.types.logical.RowType;
import org.apache.flink.table.types.logical.SmallIntType;
import org.apache.flink.table.types.logical.TinyIntType;
import org.apache.flink.table.types.logical.VarCharType;
-import org.apache.flink.table.types.logical.utils.LogicalTypeDefaultVisitor;
+import org.apache.flink.table.types.logical.YearMonthIntervalType;
+import org.apache.flink.types.Row;
+
+import java.math.BigDecimal;
+import java.math.MathContext;
+import java.math.RoundingMode;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
import static org.apache.flink.configuration.ConfigOptions.key;
-import static org.apache.flink.table.factories.DataGenTableSourceFactory.*;
+import static org.apache.flink.table.factories.DataGenTableSourceFactory.FIELDS;
+import static org.apache.flink.table.factories.DataGenTableSourceFactory.LENGTH;
+import static org.apache.flink.table.factories.DataGenTableSourceFactory.MAX;
+import static org.apache.flink.table.factories.DataGenTableSourceFactory.MIN;
+
/**
* Creates a random {@link DataGeneratorContainer} for a particular logical type.
*/
@Internal
-public class RandomGeneratorVisitor extends LogicalTypeDefaultVisitor<DataGeneratorContainer> {
+@SuppressWarnings("unchecked")
+public class RandomGeneratorVisitor extends DataGenVisitorBase {
public static final int RANDOM_STRING_LENGTH_DEFAULT = 100;
- private final String name;
-
- private final DataType type;
-
- private final ReadableConfig config;
+ private static final int RANDOM_COLLECTION_LENGTH_DEFAULT = 3;
private final ConfigOptions.OptionBuilder minKey;
private final ConfigOptions.OptionBuilder maxKey;
- public RandomGeneratorVisitor(String name, DataType type, ReadableConfig config) {
- this.name = name;
- this.type = type;
- this.config = config;
+ public RandomGeneratorVisitor(String name, ReadableConfig config) {
+ super(name, config);
this.minKey = key(FIELDS + "." + name + "." + MIN);
this.maxKey = key(FIELDS + "." + name + "." + MAX);
@@ -151,8 +166,122 @@ public class RandomGeneratorVisitor extends LogicalTypeDefaultVisitor<DataGenera
}
@Override
+ public DataGeneratorContainer visit(DecimalType decimalType) {
+ ConfigOption<Double> min = minKey.doubleType().defaultValue(Double.MIN_VALUE);
+ ConfigOption<Double> max = maxKey.doubleType().defaultValue(Double.MAX_VALUE);
+ return DataGeneratorContainer.of(
+ new BigDecimalRandomGenerator(
+ RandomGenerator.doubleGenerator(config.get(min), config.get(max)),
+ decimalType.getPrecision(), decimalType.getScale()));
+ }
+
+ @Override
+ public DataGeneratorContainer visit(YearMonthIntervalType yearMonthIntervalType) {
+ ConfigOption<Integer> min = minKey.intType().defaultValue(0);
+ ConfigOption<Integer> max = maxKey.intType().defaultValue(120000); // Period max
+ return DataGeneratorContainer.of(
+ RandomGenerator.intGenerator(
+ config.get(min), config.get(max)),
+ min, max);
+ }
+
+ @Override
+ public DataGeneratorContainer visit(DayTimeIntervalType dayTimeIntervalType) {
+ ConfigOption<Long> min = minKey.longType().defaultValue(Long.MIN_VALUE);
+ ConfigOption<Long> max = maxKey.longType().defaultValue(Long.MAX_VALUE);
+ return DataGeneratorContainer.of(
+ RandomGenerator.longGenerator(
+ config.get(min), config.get(max)),
+ min, max);
+ }
+
+ @Override
+ public DataGeneratorContainer visit(ArrayType arrayType) {
+ ConfigOption<Integer> lenOption = key(FIELDS + "." + name + "." + LENGTH)
+ .intType()
+ .defaultValue(RANDOM_COLLECTION_LENGTH_DEFAULT);
+
+ String fieldName = name + "." + "element";
+ DataGeneratorContainer container = arrayType
+ .getElementType()
+ .accept(new RandomGeneratorVisitor(fieldName, config));
+
+ return DataGeneratorContainer.of(
+ RandomGenerator.arrayGenerator(container.getGenerator(), config.get(lenOption)),
+ container.getOptions().toArray(new ConfigOption<?>[0]));
+ }
+
+ @Override
+ public DataGeneratorContainer visit(MultisetType multisetType) {
+ ConfigOption<Integer> lenOption = key(FIELDS + "." + name + "." + LENGTH)
+ .intType()
+ .defaultValue(RANDOM_COLLECTION_LENGTH_DEFAULT);
+
+ String fieldName = name + "." + "element";
+ DataGeneratorContainer container = multisetType
+ .getElementType()
+ .accept(new RandomGeneratorVisitor(fieldName, config));
+
+ return DataGeneratorContainer.of(
+ RandomGenerator.mapGenerator(container.getGenerator(),
+ RandomGenerator.intGenerator(0, 10),
+ config.get(lenOption)),
+ container.getOptions().toArray(new ConfigOption<?>[0]));
+ }
+
+ @Override
+ public DataGeneratorContainer visit(MapType mapType) {
+ ConfigOption<Integer> lenOption = key(FIELDS + "." + name + "." + LENGTH)
+ .intType()
+ .defaultValue(RANDOM_COLLECTION_LENGTH_DEFAULT);
+
+ String keyName = name + "." + "key";
+ String valName = name + "." + "value";
+
+ DataGeneratorContainer keyContainer = mapType
+ .getKeyType()
+ .accept(new RandomGeneratorVisitor(keyName, config));
+
+ DataGeneratorContainer valContainer = mapType
+ .getValueType()
+ .accept(new RandomGeneratorVisitor(valName, config));
+
+ Set<ConfigOption<?>> options = keyContainer.getOptions();
+ options.addAll(valContainer.getOptions());
+
+ return DataGeneratorContainer.of(
+ RandomGenerator.mapGenerator(
+ keyContainer.getGenerator(),
+ valContainer.getGenerator(),
+ config.get(lenOption)),
+ options.toArray(new ConfigOption<?>[0]));
+ }
+
+ @Override
+ public DataGeneratorContainer visit(RowType rowType) {
+ List<DataGeneratorContainer> fieldContainers = rowType.getFields()
+ .stream()
+ .map(field -> {
+ String fieldName = name + "." + field.getName();
+ return field.getType().accept(new RandomGeneratorVisitor(fieldName, config));
+ }).collect(Collectors.toList());
+
+ ConfigOption<?>[] options = fieldContainers
+ .stream()
+ .flatMap(container -> container.getOptions().stream())
+ .toArray(ConfigOption[]::new);
+
+ List<DataGenerator> generators = fieldContainers
+ .stream()
+ .map(DataGeneratorContainer::getGenerator)
+ .collect(Collectors.toList());
+
+ return DataGeneratorContainer.of(new RowGenerator(name, rowType.getFieldNames(), generators), options);
+ }
+
+ @Override
protected DataGeneratorContainer defaultMethod(LogicalType logicalType) {
- throw new ValidationException("Unsupported type: " + type);
+ throw new ValidationException("Unsupported type: " + logicalType);
}
private static RandomGenerator<StringData> getRandomStringGenerator(int length) {
@@ -163,4 +292,71 @@ public class RandomGeneratorVisitor extends LogicalTypeDefaultVisitor<DataGenera
}
};
}
+
+ private static class RowGenerator implements DataGenerator<Row> {
+
+ private final String rowName;
+
+ private final List<String> fieldNames;
+
+ private final List<DataGenerator> generators;
+
+ private RowGenerator(String rowName, List<String> fieldNames, List<DataGenerator> generators) {
+ this.rowName = rowName;
+ this.fieldNames = fieldNames;
+ this.generators = generators;
+ }
+
+ @Override
+ public void open(String name, FunctionInitializationContext context, RuntimeContext runtimeContext) throws Exception {
+ for (int i = 0; i < generators.size(); i++) {
+ String fullName = rowName + "." + fieldNames.get(i);
+ generators.get(i).open(fullName, context, runtimeContext);
+ }
+ }
+
+ @Override
+ public boolean hasNext() {
+ return true;
+ }
+
+ @Override
+ public Row next() {
+ Row row = new Row(generators.size());
+
+ for (int i = 0; i < generators.size(); i++) {
+ row.setField(i, generators.get(i).next());
+ }
+
+ return row;
+ }
+ }
+
+ private static class BigDecimalRandomGenerator implements DataGenerator<BigDecimal> {
+ private final RandomGenerator<Double> generator;
+ private final int precision;
+ private final int scale;
+
+ public BigDecimalRandomGenerator(RandomGenerator<Double> generator, int precision, int scale) {
+ this.generator = generator;
+ this.precision = precision;
+ this.scale = scale;
+ }
+
+ @Override
+ public void open(String name, FunctionInitializationContext context, RuntimeContext runtimeContext) throws Exception {
+ generator.open(name, context, runtimeContext);
+ }
+
+ @Override
+ public boolean hasNext() {
+ return true;
+ }
+
+ @Override
+ public BigDecimal next() {
+ BigDecimal decimal = new BigDecimal(generator.next(), new MathContext(precision));
+ return decimal.setScale(scale, RoundingMode.DOWN);
+ }
+ }
}
diff --git a/flink-table/flink-table-api-java-bridge/src/main/java/org/apache/flink/table/factories/datagen/SequenceGeneratorVisitor.java b/flink-table/flink-table-api-java-bridge/src/main/java/org/apache/flink/table/factories/datagen/SequenceGeneratorVisitor.java
index 2fa00d7..9a2c655 100644
--- a/flink-table/flink-table-api-java-bridge/src/main/java/org/apache/flink/table/factories/datagen/SequenceGeneratorVisitor.java
+++ b/flink-table/flink-table-api-java-bridge/src/main/java/org/apache/flink/table/factories/datagen/SequenceGeneratorVisitor.java
@@ -26,29 +26,27 @@ import org.apache.flink.streaming.api.functions.source.datagen.RandomGenerator;
import org.apache.flink.streaming.api.functions.source.datagen.SequenceGenerator;
import org.apache.flink.table.api.ValidationException;
import org.apache.flink.table.data.StringData;
-import org.apache.flink.table.types.DataType;
import org.apache.flink.table.types.logical.BigIntType;
import org.apache.flink.table.types.logical.BooleanType;
import org.apache.flink.table.types.logical.CharType;
+import org.apache.flink.table.types.logical.DecimalType;
import org.apache.flink.table.types.logical.DoubleType;
import org.apache.flink.table.types.logical.FloatType;
import org.apache.flink.table.types.logical.IntType;
-import org.apache.flink.table.types.logical.LogicalType;
import org.apache.flink.table.types.logical.SmallIntType;
import org.apache.flink.table.types.logical.TinyIntType;
import org.apache.flink.table.types.logical.VarCharType;
-import org.apache.flink.table.types.logical.utils.LogicalTypeDefaultVisitor;
import static org.apache.flink.configuration.ConfigOptions.key;
-import static org.apache.flink.table.factories.DataGenTableSourceFactory.*;
+import static org.apache.flink.table.factories.DataGenTableSourceFactory.END;
+import static org.apache.flink.table.factories.DataGenTableSourceFactory.FIELDS;
+import static org.apache.flink.table.factories.DataGenTableSourceFactory.START;
/**
* Creates a sequential {@link DataGeneratorContainer} for a particular logical type.
*/
@Internal
-public class SequenceGeneratorVisitor extends LogicalTypeDefaultVisitor<DataGeneratorContainer> {
-
- private final DataType type;
+public class SequenceGeneratorVisitor extends DataGenVisitorBase {
private final ReadableConfig config;
@@ -64,8 +62,9 @@ public class SequenceGeneratorVisitor extends LogicalTypeDefaultVisitor<DataGene
private final ConfigOption<Long> longEnd;
- public SequenceGeneratorVisitor(String name, DataType type, ReadableConfig config) {
- this.type = type;
+ public SequenceGeneratorVisitor(String name, ReadableConfig config) {
+ super(name, config);
+
this.config = config;
this.startKeyStr = FIELDS + "." + name + "." + START;
@@ -160,8 +159,11 @@ public class SequenceGeneratorVisitor extends LogicalTypeDefaultVisitor<DataGene
}
@Override
- protected DataGeneratorContainer defaultMethod(LogicalType logicalType) {
- throw new ValidationException("Unsupported type: " + type);
+ public DataGeneratorContainer visit(DecimalType decimalType) {
+ return DataGeneratorContainer.of(
+ SequenceGenerator.bigDecimalGenerator(
+ config.get(intStart), config.get(intEnd), decimalType.getPrecision(), decimalType.getScale()),
+ intStart, intEnd);
}
private static SequenceGenerator<StringData> getSequenceStringGenerator(long start, long end) {
diff --git a/flink-table/flink-table-api-java-bridge/src/test/java/org/apache/flink/table/factories/DataGenTableSourceFactoryTest.java b/flink-table/flink-table-api-java-bridge/src/test/java/org/apache/flink/table/factories/DataGenTableSourceFactoryTest.java
index c5a547e..1be0ffc 100644
--- a/flink-table/flink-table-api-java-bridge/src/test/java/org/apache/flink/table/factories/DataGenTableSourceFactoryTest.java
+++ b/flink-table/flink-table-api-java-bridge/src/test/java/org/apache/flink/table/factories/DataGenTableSourceFactoryTest.java
@@ -53,6 +53,7 @@ import static org.apache.flink.table.factories.DataGenTableSourceFactory.KIND;
import static org.apache.flink.table.factories.DataGenTableSourceFactory.LENGTH;
import static org.apache.flink.table.factories.DataGenTableSourceFactory.MAX;
import static org.apache.flink.table.factories.DataGenTableSourceFactory.MIN;
+import static org.apache.flink.table.factories.DataGenTableSourceFactory.NUMBER_OF_ROWS;
import static org.apache.flink.table.factories.DataGenTableSourceFactory.RANDOM;
import static org.apache.flink.table.factories.DataGenTableSourceFactory.ROWS_PER_SECOND;
import static org.apache.flink.table.factories.DataGenTableSourceFactory.SEQUENCE;
@@ -71,6 +72,43 @@ public class DataGenTableSourceFactoryTest {
.build();
@Test
+ public void testDataTypeCoverage() throws Exception {
+ TableSchema schema = TableSchema.builder()
+ .field("f0", DataTypes.CHAR(1))
+ .field("f1", DataTypes.VARCHAR(10))
+ .field("f2", DataTypes.STRING())
+ .field("f3", DataTypes.BOOLEAN())
+ .field("f4", DataTypes.DECIMAL(32, 2))
+ .field("f5", DataTypes.TINYINT())
+ .field("f6", DataTypes.SMALLINT())
+ .field("f7", DataTypes.INT())
+ .field("f8", DataTypes.BIGINT())
+ .field("f9", DataTypes.FLOAT())
+ .field("f10", DataTypes.DOUBLE())
+ .field("f11", DataTypes.DATE())
+ .field("f12", DataTypes.TIME())
+ .field("f13", DataTypes.TIMESTAMP())
+ .field("f14", DataTypes.TIMESTAMP_WITH_TIME_ZONE())
+ .field("f15", DataTypes.TIMESTAMP_WITH_LOCAL_TIME_ZONE())
+ .field("f16", DataTypes.INTERVAL(DataTypes.DAY()))
+ .field("f17", DataTypes.ARRAY(DataTypes.INT()))
+ .field("f18", DataTypes.MAP(DataTypes.STRING(), DataTypes.DATE()))
+ .field("f19", DataTypes.MULTISET(DataTypes.DECIMAL(32, 2)))
+ .field("f20", DataTypes.ROW(
+ DataTypes.FIELD("a", DataTypes.BIGINT()),
+ DataTypes.FIELD("b", DataTypes.TIME()),
+ DataTypes.FIELD("c", DataTypes.ROW(DataTypes.FIELD("d", DataTypes.TIMESTAMP())))
+ )).build();
+
+ DescriptorProperties descriptor = new DescriptorProperties();
+ descriptor.putString(FactoryUtil.CONNECTOR.key(), "datagen");
+ descriptor.putString(NUMBER_OF_ROWS.key(), "10");
+
+ List<RowData> results = runGenerator(schema, descriptor);
+ Assert.assertEquals("Failed to generate all rows", 10, results.size());
+ }
+
+ @Test
public void testSource() throws Exception {
DescriptorProperties descriptor = new DescriptorProperties();
descriptor.putString(FactoryUtil.CONNECTOR.key(), "datagen");
@@ -87,8 +125,21 @@ public class DataGenTableSourceFactoryTest {
descriptor.putLong(FIELDS + ".f2." + START, 50);
descriptor.putLong(FIELDS + ".f2." + END, 60);
+ List<RowData> results = runGenerator(TEST_SCHEMA, descriptor);
+
+ Assert.assertEquals(11, results.size());
+ for (int i = 0; i < results.size(); i++) {
+ RowData row = results.get(i);
+ Assert.assertEquals(20, row.getString(0).toString().length());
+ long f1 = row.getLong(1);
+ Assert.assertTrue(f1 >= 10 && f1 <= 100);
+ Assert.assertEquals(i + 50, row.getLong(2));
+ }
+ }
+
+ private List<RowData> runGenerator(TableSchema schema, DescriptorProperties descriptor) throws Exception {
DynamicTableSource source = createSource(
- TEST_SCHEMA, descriptor.asMap());
+ schema, descriptor.asMap());
assertTrue(source instanceof DataGenTableSource);
@@ -103,47 +154,11 @@ public class DataGenTableSourceFactoryTest {
new AbstractStreamOperatorTestHarness<>(src, 1, 1, 0);
testHarness.open();
- List<RowData> results = new ArrayList<>();
-
- gen.run(new SourceFunction.SourceContext<RowData>() {
-
- private Object lock = new Object();
-
- @Override
- public void collect(RowData element) {
- results.add(element);
- }
-
- @Override
- public void collectWithTimestamp(RowData element, long timestamp) {
- }
-
- @Override
- public void emitWatermark(Watermark mark) {
- }
-
- @Override
- public void markAsTemporarilyIdle() {
- }
-
- @Override
- public Object getCheckpointLock() {
- return lock;
- }
+ TestContext ctx = new TestContext();
- @Override
- public void close() {
- }
- });
+ gen.run(ctx);
- Assert.assertEquals(11, results.size());
- for (int i = 0; i < results.size(); i++) {
- RowData row = results.get(i);
- Assert.assertEquals(20, row.getString(0).toString().length());
- long f1 = row.getLong(1);
- Assert.assertTrue(f1 >= 10 && f1 <= 100);
- Assert.assertEquals(i + 50, row.getLong(2));
- }
+ return ctx.results;
}
@Test
@@ -311,4 +326,37 @@ public class DataGenTableSourceFactoryTest {
new Configuration(),
Thread.currentThread().getContextClassLoader());
}
+
+ private static class TestContext implements SourceFunction.SourceContext<RowData> {
+
+ private final Object lock = new Object();
+
+ private final List<RowData> results = new ArrayList<>();
+
+ @Override
+ public void collect(RowData element) {
+ results.add(element);
+ }
+
+ @Override
+ public void collectWithTimestamp(RowData element, long timestamp) {
+ }
+
+ @Override
+ public void emitWatermark(Watermark mark) {
+ }
+
+ @Override
+ public void markAsTemporarilyIdle() {
+ }
+
+ @Override
+ public Object getCheckpointLock() {
+ return lock;
+ }
+
+ @Override
+ public void close() {
+ }
+ }
}