Skip to content

Commit db810bb

Browse files
ggjh-159ggjh-159
committed
feat(flink): add columnar StreamRecordTimestampInserter operator
1 parent fb646be commit db810bb

7 files changed

Lines changed: 275 additions & 4 deletions

File tree

‎gluten-flink/planner/src/main/java/org/apache/gluten/velox/FileSystemSinkFactory.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -134,9 +134,15 @@ public Transformation<RowData> buildVeloxSink(
134134
"FileSystemInsertTable");
135135
GlutenOneInputOperatorFactory<?, ?> operatorFactory =
136136
new GlutenOneInputOperatorFactory(onewInputOperator);
137+
// If rowtime transformation was applied (rowtimeFieldIndex != -1), a native
138+
// StreamRecordTimestampInserter sits at this position. Replace it with a Gluten columnar
139+
// inserter so the chain stays columnar end-to-end; otherwise this is a no-op.
140+
Transformation<RowData> veloxFileWriterInput =
141+
GlutenRowtimeInserterHelper.processTransformation(
142+
(Transformation<RowData>) fileWriterTransformation.getInputs().get(0));
137143
Transformation<RowData> veloxFileWriterTransformation =
138144
new OneInputTransformation(
139-
fileWriterTransformation.getInputs().get(0),
145+
veloxFileWriterInput,
140146
fileWriterTransformation.getName(),
141147
operatorFactory,
142148
fileWriterTransformation.getOutputType(),

‎gluten-flink/planner/src/main/java/org/apache/gluten/velox/FuzzerSourceSinkFactory.java‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -138,8 +138,7 @@ public Transformation<RowData> buildVeloxSink(
138138
RowData.class,
139139
"FuzzerSink"));
140140
DataStream<RowData> newInputStream =
141-
sinkTransformation
142-
.getInputStream()
141+
GlutenRowtimeInserterHelper.process(sinkTransformation.getInputStream())
143142
.transform("Writer", CommittableMessageTypeInfo.noOutput(), operatorFactory);
144143
return new SinkTransformation<RowData, RowData>(
145144
newInputStream,

‎gluten-flink/planner/src/main/java/org/apache/gluten/velox/GlutenRowtimeInserterHelper.java‎

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
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+
import org.apache.gluten.util.LogicalTypeConverter;
21+
import org.apache.gluten.util.PlanNodeIdGenerator;
22+
import org.apache.gluten.util.ReflectUtils;
23+
24+
import io.github.zhztheplayer.velox4j.expression.FieldAccessTypedExpr;
25+
import io.github.zhztheplayer.velox4j.expression.InputTypedExpr;
26+
import io.github.zhztheplayer.velox4j.plan.EmptyNode;
27+
import io.github.zhztheplayer.velox4j.plan.ProjectNode;
28+
import io.github.zhztheplayer.velox4j.plan.StatefulPlanNode;
29+
import io.github.zhztheplayer.velox4j.plan.StreamRecordTimestampInserterNode;
30+
import io.github.zhztheplayer.velox4j.type.RowType;
31+
32+
import org.apache.flink.api.dag.Transformation;
33+
import org.apache.flink.streaming.api.datastream.DataStream;
34+
import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
35+
import org.apache.flink.streaming.api.operators.SimpleOperatorFactory;
36+
import org.apache.flink.streaming.api.transformations.OneInputTransformation;
37+
import org.apache.flink.table.data.RowData;
38+
import org.apache.flink.table.runtime.operators.sink.StreamRecordTimestampInserter;
39+
import org.apache.flink.table.runtime.typeutils.InternalTypeInfo;
40+
41+
import java.util.List;
42+
import java.util.Map;
43+
44+
/**
45+
* Replaces a native Flink {@link StreamRecordTimestampInserter} (per-row timestamp) with a Gluten
46+
* columnar inserter (batch-max timestamp) so the columnar chain stays intact when the downstream
47+
* sink is offloaded to Velox.
48+
*
49+
* <p>The native inserter is constructed by {@code CommonExecSink.applyRowtimeTransformation} and
50+
* sits somewhere on the sink's input chain (position depends on the sink: for a simple sink like
51+
* Fuzzer/DiscardingSink it is the direct input; for FileSystem it sits below the
52+
* StreamingFileWriter and PartitionCommitter). Each sink factory's {@code buildVeloxSink} invokes
53+
* this helper on the inserter-bearing transformation it is about to replace; if no inserter is
54+
* present (e.g., {@code rowtimeFieldIndex == -1}), the helper is a no-op and returns the input
55+
* unchanged.
56+
*/
57+
public final class GlutenRowtimeInserterHelper {
58+
59+
private GlutenRowtimeInserterHelper() {}
60+
61+
/**
62+
* Convenience overload that accepts a {@link DataStream} (typical entry point for factories that
63+
* use {@code sinkTransformation.getInputStream()}). Inspects the underlying transformation; if it
64+
* is a native inserter, rebuilds it as a Gluten columnar inserter and returns a new DataStream
65+
* whose terminal node is the Gluten inserter. Otherwise returns the inputStream unchanged.
66+
*/
67+
public static DataStream<RowData> process(DataStream<RowData> inputStream) {
68+
Transformation<RowData> inputTrans = inputStream.getTransformation();
69+
Transformation<RowData> newTrans = processTransformation(inputTrans);
70+
if (newTrans == inputTrans) {
71+
return inputStream;
72+
}
73+
return new DataStream<>(inputStream.getExecutionEnvironment(), newTrans);
74+
}
75+
76+
/**
77+
* Inspect {@code inputTrans}; if it is a native {@link StreamRecordTimestampInserter}, rebuild it
78+
* as a Gluten columnar inserter whose input is the native inserter's upstream. Returns the new
79+
* transformation, or the original inputTrans when no replacement happened.
80+
*/
81+
public static Transformation<RowData> processTransformation(Transformation<RowData> inputTrans) {
82+
if (!(inputTrans instanceof OneInputTransformation)) {
83+
return inputTrans;
84+
}
85+
OneInputTransformation<?, ?> oneInput = (OneInputTransformation<?, ?>) inputTrans;
86+
if (!(oneInput.getOperatorFactory() instanceof SimpleOperatorFactory)) {
87+
return inputTrans;
88+
}
89+
@SuppressWarnings("rawtypes")
90+
Object op = ((SimpleOperatorFactory) oneInput.getOperatorFactory()).getOperator();
91+
if (!(op instanceof StreamRecordTimestampInserter)) {
92+
return inputTrans;
93+
}
94+
int rowtimeIndex =
95+
(int) ReflectUtils.getObjectField(StreamRecordTimestampInserter.class, op, "rowtimeIndex");
96+
List<Transformation<?>> inputs = oneInput.getInputs();
97+
if (inputs.isEmpty()) {
98+
return inputTrans;
99+
}
100+
@SuppressWarnings("unchecked")
101+
Transformation<RowData> aboveInserter = (Transformation<RowData>) inputs.get(0);
102+
return buildGlutenInserter(aboveInserter, rowtimeIndex, oneInput.getParallelism());
103+
}
104+
105+
private static Transformation<RowData> buildGlutenInserter(
106+
Transformation<RowData> aboveInserter, int rowtimeFieldIndex, int parallelism) {
107+
@SuppressWarnings("unchecked")
108+
InternalTypeInfo<RowData> internalTypeInfo =
109+
(InternalTypeInfo<RowData>) aboveInserter.getOutputType();
110+
final org.apache.flink.table.types.logical.RowType inputRowType =
111+
(org.apache.flink.table.types.logical.RowType) internalTypeInfo.toLogicalType();
112+
final RowType vlInputType = (RowType) LogicalTypeConverter.toVLType(inputRowType);
113+
final List<String> fieldNames = inputRowType.getFieldNames();
114+
final String rowtimeFieldName = fieldNames.get(rowtimeFieldIndex);
115+
final InputTypedExpr inputExpr = new InputTypedExpr(vlInputType);
116+
final ProjectNode project =
117+
new ProjectNode(
118+
PlanNodeIdGenerator.newId(),
119+
List.of(new EmptyNode(vlInputType)),
120+
List.of(rowtimeFieldName),
121+
List.of(FieldAccessTypedExpr.create(inputExpr, rowtimeFieldName)));
122+
final StreamRecordTimestampInserterNode inserterNode =
123+
new StreamRecordTimestampInserterNode(
124+
PlanNodeIdGenerator.newId(), null, project, rowtimeFieldIndex);
125+
final StatefulPlanNode statefulPlan = new StatefulPlanNode(inserterNode.getId(), inserterNode);
126+
127+
final GlutenStreamRecordTimestampInserter operator =
128+
new GlutenStreamRecordTimestampInserter(
129+
statefulPlan,
130+
PlanNodeIdGenerator.newId(),
131+
vlInputType,
132+
Map.of(inserterNode.getId(), vlInputType),
133+
rowtimeFieldIndex,
134+
"GlutenStreamRecordTimestampInserter");
135+
136+
@SuppressWarnings({"rawtypes", "unchecked"})
137+
final OneInputStreamOperator rawOperator = (OneInputStreamOperator) operator;
138+
return new OneInputTransformation<>(
139+
aboveInserter,
140+
"GlutenStreamRecordTimestampInserter",
141+
SimpleOperatorFactory.of(rawOperator),
142+
aboveInserter.getOutputType(),
143+
parallelism);
144+
}
145+
}

‎gluten-flink/planner/src/main/java/org/apache/gluten/velox/PrintSinkFactory.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,11 @@ public boolean isStdErr() {
108108
public Transformation buildVeloxSink(
109109
Transformation<RowData> transformation, Map<String, Object> parameters) {
110110
Transformation inputTrans = (Transformation) transformation.getInputs().get(0);
111+
// If the upstream chain contains a native StreamRecordTimestampInserter (e.g., when this
112+
// sink is reached via a rowtime-bearing path), replace it with a Gluten columnar inserter.
113+
// For the typical PrintSink path (SinkFunctionProvider), no inserter is present and this is
114+
// a no-op.
115+
inputTrans = GlutenRowtimeInserterHelper.processTransformation(inputTrans);
111116
InternalTypeInfo inputTypeInfo = (InternalTypeInfo) inputTrans.getOutputType();
112117

113118
PrintOptions printOpts = extractPrintOptions(transformation);

‎gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenOneInputOperator.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,7 @@ public class GlutenOneInputOperator<IN, OUT> extends TableStreamOperator<OUT>
6767
private final Class<IN> inClass;
6868
private final Class<OUT> outClass;
6969
private transient VectorInputBridge<IN> inputBridge;
70-
private transient VectorOutputBridge<OUT> outputBridge;
70+
protected transient VectorOutputBridge<OUT> outputBridge;
7171
private final GlutenMailboxHolder mailboxHolder = new GlutenMailboxHolder();
7272

7373
public GlutenOneInputOperator(

‎gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenStreamRecordTimestampInserter.java‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
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.table.runtime.operators;
18+
19+
import org.apache.gluten.util.TimestampSettingStatefulRecordOutputBridge;
20+
21+
import io.github.zhztheplayer.velox4j.plan.StatefulPlanNode;
22+
import io.github.zhztheplayer.velox4j.stateful.StatefulRecord;
23+
import io.github.zhztheplayer.velox4j.type.RowType;
24+
25+
import java.util.Map;
26+
27+
/**
28+
* Gluten equivalent of Flink's StreamRecordTimestampInserter that keeps data columnar end-to-end.
29+
* Both input and output classes are {@link StatefulRecord}, so the upstream chain (e.g. Calc) and
30+
* the downstream sink can stay in vector form without row-column conversion. The rowtime is
31+
* extracted in the native {@code StreamRecordTimestampInserter} operator (wrapped by the
32+
* StatefulPlanNode), and {@link TimestampSettingStatefulOutputBridge} only forwards it onto the
33+
* StreamRecord.
34+
*/
35+
public class GlutenStreamRecordTimestampInserter
36+
extends GlutenOneInputOperator<StatefulRecord, StatefulRecord> {
37+
38+
public GlutenStreamRecordTimestampInserter(
39+
StatefulPlanNode plan,
40+
String id,
41+
RowType inputType,
42+
Map<String, RowType> outputTypes,
43+
int rowtimeIndex,
44+
String description) {
45+
super(
46+
plan, id, inputType, outputTypes, StatefulRecord.class, StatefulRecord.class, description);
47+
this.outputBridge = new TimestampSettingStatefulRecordOutputBridge();
48+
}
49+
50+
@Override
51+
void initSession() {
52+
if (outputBridge == null) {
53+
outputBridge = new TimestampSettingStatefulRecordOutputBridge();
54+
}
55+
super.initSession();
56+
}
57+
}

‎gluten-flink/runtime/src/main/java/org/apache/gluten/util/TimestampSettingStatefulRecordOutputBridge.java‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
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.util;
18+
19+
import io.github.zhztheplayer.velox4j.stateful.StatefulRecord;
20+
import io.github.zhztheplayer.velox4j.type.RowType;
21+
22+
import org.apache.flink.streaming.api.operators.Output;
23+
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
24+
25+
import org.apache.arrow.memory.BufferAllocator;
26+
27+
/**
28+
* Columnar output bridge for {@code GlutenStreamRecordTimestampInserter}. The native
29+
* StreamRecordTimestampInserter operator has already annotated the StatefulRecord with the max
30+
* rowtime of the batch during its advance(); this bridge just forwards that timestamp onto the
31+
* Flink StreamRecord. No JNI helper call, no row-column conversion.
32+
*/
33+
public class TimestampSettingStatefulRecordOutputBridge
34+
implements VectorOutputBridge<StatefulRecord> {
35+
36+
private static final long serialVersionUID = 1L;
37+
private transient StreamRecord<StatefulRecord> outputElement;
38+
39+
public TimestampSettingStatefulRecordOutputBridge() {
40+
this.outputElement = new StreamRecord<>(null);
41+
}
42+
43+
@Override
44+
public void collect(
45+
Output<StreamRecord<StatefulRecord>> collector,
46+
StatefulRecord record,
47+
BufferAllocator allocator,
48+
RowType outputType) {
49+
final long ts = record.hasTimestamp() ? record.getTimestamp() : Long.MIN_VALUE;
50+
collector.collect(getOrCreateOutputElement().replace(record, ts));
51+
}
52+
53+
private StreamRecord<StatefulRecord> getOrCreateOutputElement() {
54+
if (outputElement == null) {
55+
outputElement = new StreamRecord<>(null);
56+
}
57+
return outputElement;
58+
}
59+
}

0 commit comments

Comments
 (0)