Skip to content

Commit aeb09c5

Browse files
talatuyarertalatuyarer
authored
BeamSQL Datagen Table Provider (#35572)
* Initial version of SQL Datagen Table * Missing package-info.java file of SQL Datagen Table * Changed RandomStringUtils import from common lang3 to vendor calcite's common lang * Fixed unnecessary whitespace in the doc. * Aligned kinds with BeamSQL's Data type rather than Schema Types to reduce confusion. Added event time watermark support for Datagen Table. * Fixed spotbugs issue * Mentioned after which version it is available * Addressed PR comments * I fixed flaky test
1 parent 52e42d7 commit aeb09c5

8 files changed

Lines changed: 1048 additions & 0 deletions

File tree

‎sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/datagen/AdvancingTimestampFn.java‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
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.sdk.extensions.sql.meta.provider.datagen;
19+
20+
import java.util.concurrent.ThreadLocalRandom;
21+
import org.apache.beam.sdk.transforms.SerializableFunction;
22+
import org.joda.time.Duration;
23+
import org.joda.time.Instant;
24+
25+
class AdvancingTimestampFn implements SerializableFunction<Long, Instant> {
26+
private final long maxOutOfOrdernessMs;
27+
private final Instant baseTime = Instant.now();
28+
29+
AdvancingTimestampFn(long maxOutOfOrdernessMs) {
30+
this.maxOutOfOrdernessMs = maxOutOfOrdernessMs;
31+
}
32+
33+
@Override
34+
public Instant apply(Long index) {
35+
long delay = (long) (ThreadLocalRandom.current().nextDouble() * maxOutOfOrdernessMs);
36+
return baseTime.plus(Duration.millis(index * 1000)).minus(Duration.millis(delay));
37+
}
38+
}

‎sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/datagen/DataGeneratorPTransform.java‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
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.sdk.extensions.sql.meta.provider.datagen;
19+
20+
import com.fasterxml.jackson.databind.JsonNode;
21+
import com.fasterxml.jackson.databind.node.ObjectNode;
22+
import javax.annotation.Nullable;
23+
import org.apache.beam.sdk.io.GenerateSequence;
24+
import org.apache.beam.sdk.schemas.Schema;
25+
import org.apache.beam.sdk.transforms.PTransform;
26+
import org.apache.beam.sdk.transforms.ParDo;
27+
import org.apache.beam.sdk.values.PBegin;
28+
import org.apache.beam.sdk.values.PCollection;
29+
import org.apache.beam.sdk.values.Row;
30+
import org.joda.time.Duration;
31+
32+
/** The main PTransform that encapsulates the data generation logic. */
33+
public class DataGeneratorPTransform extends PTransform<PBegin, PCollection<Row>> {
34+
private final Schema schema;
35+
private final ObjectNode properties;
36+
37+
public DataGeneratorPTransform(Schema schema, ObjectNode properties) {
38+
this.schema = schema;
39+
this.properties = properties;
40+
}
41+
42+
@Override
43+
public PCollection<Row> expand(PBegin input) {
44+
GenerateSequence generator;
45+
JsonNode rpsNode = properties.path("rows-per-second");
46+
JsonNode numRowsNode = properties.path("number-of-rows");
47+
48+
if (!rpsNode.isMissingNode()) {
49+
generator = GenerateSequence.from(0).withRate(rpsNode.asLong(), Duration.standardSeconds(1));
50+
} else if (!numRowsNode.isMissingNode()) {
51+
generator = GenerateSequence.from(0).to(numRowsNode.asLong());
52+
} else {
53+
throw new IllegalArgumentException(
54+
"A 'datagen' table requires either 'rows-per-second' (for unbounded) or "
55+
+ "'number-of-rows' (for bounded) in TBLPROPERTIES.");
56+
}
57+
58+
String behavior = properties.path("timestamp.behavior").asText("processing-time");
59+
@Nullable String eventTimeColumn = null;
60+
61+
if ("event-time".equalsIgnoreCase(behavior)) {
62+
JsonNode columnNode = properties.path("event-time.timestamp-column");
63+
64+
if (columnNode.isMissingNode() || columnNode.isNull()) {
65+
throw new IllegalArgumentException(
66+
"For 'event-time' behavior, 'event-time.timestamp-column' must be specified.");
67+
}
68+
eventTimeColumn = columnNode.asText();
69+
70+
// Validate that the specified column exists and is of type TIMESTAMP.
71+
if (!schema.hasField(eventTimeColumn)) {
72+
throw new IllegalArgumentException(
73+
String.format(
74+
"The specified 'event-time.timestamp-column' ('%s') does not exist in the table schema.",
75+
eventTimeColumn));
76+
}
77+
78+
Schema.Field eventTimeField = schema.getField(eventTimeColumn);
79+
if (!Schema.TypeName.DATETIME.equals(eventTimeField.getType().getTypeName())) {
80+
throw new IllegalArgumentException(
81+
String.format(
82+
"The specified 'event-time.timestamp-column' ('%s') must be of type TIMESTAMP, but was '%s'.",
83+
eventTimeColumn, eventTimeField.getType()));
84+
}
85+
86+
long maxOutOfOrdernessMs = properties.path("event_time.max-out-of-orderness").asLong(0L);
87+
generator = generator.withTimestampFn(new AdvancingTimestampFn(maxOutOfOrdernessMs));
88+
}
89+
90+
return input
91+
.getPipeline()
92+
.apply("GenerateSequence", generator)
93+
.apply(
94+
"GenerateRows", ParDo.of(new DataGeneratorRowFn(schema, properties, eventTimeColumn)))
95+
.setRowSchema(schema);
96+
}
97+
}

‎sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/datagen/DataGeneratorRowFn.java‎

Lines changed: 200 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,200 @@
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.sdk.extensions.sql.meta.provider.datagen;
19+
20+
import com.fasterxml.jackson.databind.JsonNode;
21+
import com.fasterxml.jackson.databind.node.ObjectNode;
22+
import java.io.Serializable;
23+
import java.math.BigDecimal;
24+
import java.util.HashMap;
25+
import java.util.Map;
26+
import java.util.concurrent.ThreadLocalRandom;
27+
import javax.annotation.Nullable;
28+
import org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils;
29+
import org.apache.beam.sdk.schemas.Schema;
30+
import org.apache.beam.sdk.transforms.DoFn;
31+
import org.apache.beam.sdk.values.Row;
32+
import org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.type.SqlTypeName;
33+
import org.apache.beam.vendor.calcite.v1_28_0.org.apache.commons.lang.RandomStringUtils;
34+
import org.joda.time.Duration;
35+
import org.joda.time.Instant;
36+
37+
/** A stateful DoFn that converts a sequence of Longs into structured Rows. */
38+
public class DataGeneratorRowFn extends DoFn<Long, Row> {
39+
private final Schema schema;
40+
private final ObjectNode properties;
41+
private final @Nullable String primaryTimestampField;
42+
43+
private transient Map<String, FieldGenerator> fieldGenerators;
44+
45+
@SuppressWarnings("initialization")
46+
public DataGeneratorRowFn(
47+
Schema schema, ObjectNode properties, @Nullable String primaryTimestampField) {
48+
this.schema = schema;
49+
this.properties = properties;
50+
this.primaryTimestampField = primaryTimestampField;
51+
}
52+
53+
@Setup
54+
public void setup() {
55+
this.fieldGenerators = new HashMap<>();
56+
57+
for (Schema.Field field : schema.getFields()) {
58+
fieldGenerators.put(field.getName(), createGeneratorForField(field));
59+
}
60+
}
61+
62+
@ProcessElement
63+
public void processElement(
64+
@Element Long index, @Timestamp Instant timestamp, OutputReceiver<Row> out) {
65+
Row.Builder rowBuilder = Row.withSchema(schema);
66+
for (Schema.Field field : schema.getFields()) {
67+
Object value;
68+
if (field.getName().equals(this.primaryTimestampField)) {
69+
value = timestamp.toDateTime();
70+
} else {
71+
FieldGenerator generator = fieldGenerators.get(field.getName());
72+
if (generator == null) {
73+
throw new IllegalStateException("No generator found for field: " + field.getName());
74+
}
75+
value = generator.generate(index);
76+
}
77+
rowBuilder.addValue(value);
78+
}
79+
out.output(rowBuilder.build());
80+
}
81+
82+
@FunctionalInterface
83+
private interface FieldGenerator extends Serializable {
84+
@Nullable
85+
Object generate(long index);
86+
}
87+
88+
private FieldGenerator createGeneratorForField(Schema.Field field) {
89+
String fieldName = field.getName();
90+
FieldGenerator valueGenerator = createValueGeneratorForField(field);
91+
double nullRate = properties.path("fields." + fieldName + ".null-rate").asDouble(0.0);
92+
93+
if (nullRate > 0) {
94+
return (index) ->
95+
ThreadLocalRandom.current().nextDouble() < nullRate
96+
? null
97+
: valueGenerator.generate(index);
98+
}
99+
return valueGenerator;
100+
}
101+
102+
private FieldGenerator createValueGeneratorForField(Schema.Field field) {
103+
String fieldName = field.getName();
104+
String kind = properties.path("fields." + fieldName + ".kind").asText("random");
105+
106+
final SqlTypeName sqlTypeName = CalciteUtils.toSqlTypeName(field.getType());
107+
if (sqlTypeName == null) {
108+
throw new UnsupportedOperationException(
109+
"Data generator requires a defined SQL type. Beam type '"
110+
+ field.getType().getTypeName()
111+
+ "' on field '"
112+
+ field.getName()
113+
+ "' is not supported.");
114+
}
115+
116+
if ("sequence".equalsIgnoreCase(kind)) {
117+
if (!SqlTypeName.INT_TYPES.contains(sqlTypeName)) {
118+
throw new IllegalArgumentException(
119+
String.format(
120+
"The 'sequence' generator for integers only supports integer types, but field '%s' is of type '%s'.",
121+
field.getName(), sqlTypeName));
122+
}
123+
124+
JsonNode startNode = properties.path("fields." + fieldName + ".start");
125+
JsonNode endNode = properties.path("fields." + fieldName + ".end");
126+
127+
if (startNode.isMissingNode() && endNode.isMissingNode()) {
128+
return (index) -> index;
129+
}
130+
131+
if (startNode.isMissingNode() || endNode.isMissingNode()) {
132+
throw new IllegalArgumentException(
133+
"For a cycling sequence generator, both 'start' and 'end' must be specified.");
134+
}
135+
136+
long start = startNode.asLong(0L);
137+
long end = endNode.asLong(Long.MAX_VALUE);
138+
139+
if (start > end) {
140+
throw new IllegalArgumentException(
141+
String.format(
142+
"For sequence generator, 'start' (%d) cannot be greater than 'end' (%d).",
143+
start, end));
144+
}
145+
long cycleLength = end - start + 1;
146+
switch (sqlTypeName) {
147+
case INTEGER:
148+
return (index) -> (int) (start + (index % cycleLength));
149+
case SMALLINT:
150+
return (index) -> (short) (start + (index % cycleLength));
151+
case TINYINT:
152+
return (index) -> (byte) (start + (index % cycleLength));
153+
default: // BIGINT
154+
return (index) -> start + (index % cycleLength);
155+
}
156+
}
157+
158+
switch (sqlTypeName) {
159+
case CHAR:
160+
case VARCHAR:
161+
int length = properties.path("fields." + fieldName + ".length").asInt(10);
162+
return (index) -> RandomStringUtils.randomAlphanumeric(length);
163+
case BOOLEAN:
164+
return (index) -> ThreadLocalRandom.current().nextBoolean();
165+
case FLOAT:
166+
case DOUBLE:
167+
double minD = properties.path("fields." + fieldName + ".min").asDouble(0.0);
168+
double maxD = properties.path("fields." + fieldName + ".max").asDouble(1.0);
169+
return (index) -> minD + (maxD - minD) * ThreadLocalRandom.current().nextDouble();
170+
case TINYINT:
171+
case SMALLINT:
172+
case INTEGER:
173+
case BIGINT:
174+
long minL = properties.path("fields." + fieldName + ".min").asLong(0L);
175+
long maxL = properties.path("fields." + fieldName + ".max").asLong(Long.MAX_VALUE);
176+
return (index) -> minL + (long) (ThreadLocalRandom.current().nextDouble() * (maxL - minL));
177+
case DECIMAL:
178+
double minBd = properties.path("fields." + fieldName + ".min").asDouble(0.0);
179+
double maxBd = properties.path("fields." + fieldName + ".max").asDouble(1000.0);
180+
return (index) ->
181+
BigDecimal.valueOf(minBd + (maxBd - minBd) * ThreadLocalRandom.current().nextDouble());
182+
case TIMESTAMP:
183+
JsonNode maxPastNode = properties.path("fields." + fieldName + ".max-past");
184+
if (!maxPastNode.isMissingNode()) {
185+
long maxPastMs = maxPastNode.asLong();
186+
if (maxPastMs <= 0) {
187+
throw new IllegalArgumentException("'max-past' must be a positive long value.");
188+
}
189+
return (index) ->
190+
Instant.now()
191+
.minus(
192+
Duration.millis(
193+
(long) (ThreadLocalRandom.current().nextDouble() * maxPastMs)));
194+
}
195+
return (index) -> Instant.now();
196+
default:
197+
throw new UnsupportedOperationException("Unsupported SQL type for datagen: " + sqlTypeName);
198+
}
199+
}
200+
}

0 commit comments

Comments
 (0)