Skip to content

Commit 47c85b5

Browse files
committed
test(flink): add GlutenRowtimeInserterHelper unit tests
1 parent db810bb commit 47c85b5

1 file changed

Lines changed: 117 additions & 0 deletions

File tree

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package org.apache.gluten.velox;
18+
19+
import org.apache.gluten.table.runtime.operators.GlutenStreamRecordTimestampInserter;
20+
21+
import org.apache.flink.api.common.functions.MapFunction;
22+
import org.apache.flink.api.dag.Transformation;
23+
import org.apache.flink.streaming.api.operators.SimpleOperatorFactory;
24+
import org.apache.flink.streaming.api.operators.StreamMap;
25+
import org.apache.flink.streaming.api.transformations.OneInputTransformation;
26+
import org.apache.flink.table.data.RowData;
27+
import org.apache.flink.table.runtime.operators.sink.StreamRecordTimestampInserter;
28+
import org.apache.flink.table.runtime.typeutils.InternalTypeInfo;
29+
import org.apache.flink.table.types.logical.BigIntType;
30+
import org.apache.flink.table.types.logical.RowType;
31+
32+
import org.junit.jupiter.api.Test;
33+
34+
import java.util.Collections;
35+
import java.util.List;
36+
37+
import static org.junit.jupiter.api.Assertions.assertNotSame;
38+
import static org.junit.jupiter.api.Assertions.assertSame;
39+
import static org.junit.jupiter.api.Assertions.assertTrue;
40+
41+
class GlutenRowtimeInserterHelperTest {
42+
43+
private static final RowType UPSTREAM_ROW_TYPE = RowType.of(new BigIntType());
44+
45+
private static Transformation<RowData> newUpstream() {
46+
return new StubTransformation("upstream", InternalTypeInfo.of(UPSTREAM_ROW_TYPE));
47+
}
48+
49+
private static OneInputTransformation<RowData, RowData> newNativeInserterTx(
50+
Transformation<RowData> upstream, int rowtimeIndex) {
51+
StreamRecordTimestampInserter op = new StreamRecordTimestampInserter(rowtimeIndex);
52+
return new OneInputTransformation<>(
53+
upstream, "native-inserter", op, upstream.getOutputType(), 1);
54+
}
55+
56+
private static OneInputTransformation<RowData, RowData> newOtherOperatorTx(
57+
Transformation<RowData> upstream) {
58+
StreamMap<RowData, RowData> other = new StreamMap<>(new IdentityMapFunction());
59+
return new OneInputTransformation<>(upstream, "other-op", other, upstream.getOutputType(), 1);
60+
}
61+
62+
@Test
63+
void testNoOpForNonOneInputTransformation() {
64+
Transformation<RowData> stub =
65+
new StubTransformation("stub", InternalTypeInfo.of(UPSTREAM_ROW_TYPE));
66+
assertSame(stub, GlutenRowtimeInserterHelper.processTransformation(stub));
67+
}
68+
69+
@Test
70+
void testNoOpForNonInserterOperator() {
71+
Transformation<RowData> upstream = newUpstream();
72+
OneInputTransformation<RowData, RowData> tx = newOtherOperatorTx(upstream);
73+
assertSame(tx, GlutenRowtimeInserterHelper.processTransformation(tx));
74+
}
75+
76+
@Test
77+
void testReplacesNativeInserter() {
78+
Transformation<RowData> upstream = newUpstream();
79+
OneInputTransformation<RowData, RowData> nativeTx = newNativeInserterTx(upstream, 0);
80+
81+
Transformation<RowData> result = GlutenRowtimeInserterHelper.processTransformation(nativeTx);
82+
83+
assertNotSame(nativeTx, result);
84+
assertTrue(result instanceof OneInputTransformation);
85+
@SuppressWarnings("unchecked")
86+
OneInputTransformation<RowData, RowData> out =
87+
(OneInputTransformation<RowData, RowData>) result;
88+
assertTrue(out.getOperatorFactory() instanceof SimpleOperatorFactory);
89+
Object op = ((SimpleOperatorFactory<?>) out.getOperatorFactory()).getOperator();
90+
assertTrue(op instanceof GlutenStreamRecordTimestampInserter);
91+
assertSame(upstream, out.getInputs().get(0));
92+
assertSame(1, out.getParallelism());
93+
}
94+
95+
private static final class StubTransformation extends Transformation<RowData> {
96+
StubTransformation(String name, InternalTypeInfo<RowData> typeInfo) {
97+
super(name, typeInfo, 1);
98+
}
99+
100+
@Override
101+
public List<Transformation<?>> getInputs() {
102+
return Collections.emptyList();
103+
}
104+
105+
@Override
106+
protected List<Transformation<?>> getTransitivePredecessorsInternal() {
107+
return Collections.emptyList();
108+
}
109+
}
110+
111+
private static final class IdentityMapFunction implements MapFunction<RowData, RowData> {
112+
@Override
113+
public RowData map(RowData value) {
114+
return value;
115+
}
116+
}
117+
}

0 commit comments

Comments
 (0)