This is an automated email from the ASF dual-hosted git repository.
lgbo-ustc pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new 9202345c23 [GLUTEN-10246][FLINK] Support processing time
window(tumble) for nextmark q12 (#10898)
9202345c23 is described below
commit 9202345c23fec9e62f516bde51a07b09f4d58ef6
Author: kevinyhzou <[email protected]>
AuthorDate: Wed Jun 10 09:39:04 2026 +0800
[GLUTEN-10246][FLINK] Support processing time window(tumble) for nextmark
q12 (#10898)
* support proctime window
* remove useless changes
* remove useless changes
* support kafka data source
* fix reviews
* optimize code
* fix memory leak
* merge master
* xxx
* support proc time window
* remove useless code
* fix ut
* restore q5,q7,q8 tests
* Fix tz download failure
* fix ci
* update velox4j version
---
.github/workflows/flink.yml | 12 ++-
gluten-flink/docs/Flink.md | 2 +-
.../stream/StreamExecGlobalWindowAggregate.java | 85 ++++++----------
.../stream/StreamExecLocalWindowAggregate.java | 5 +-
.../exec/stream/StreamExecWindowAggregate.java | 42 +++++---
.../org/apache/gluten/rexnode/WindowUtils.java | 31 +++++-
.../gluten/client/OffloadedJobGraphGenerator.java | 10 +-
.../operators/GlutenOneInputOperatorFactory.java | 6 ++
.../streaming/api/operators/GlutenOperator.java | 45 +++++++++
.../runtime/tasks/GlutenOutputCollector.java | 3 -
.../table/runtime/config/VeloxQueryConfig.java | 3 +
.../runtime/operators/GlutenMailboxHolder.java} | 30 +++---
.../operators/GlutenMailboxOperatorHelper.java | 108 +++++++++++++++++++++
.../runtime/operators/GlutenOneInputOperator.java | 78 +++++++++------
...onResource.java => GlutenSessionResources.java} | 41 ++++++++
.../runtime/operators/GlutenSourceFunction.java | 3 +-
.../runtime/operators/GlutenTwoInputOperator.java | 41 ++++++--
.../table/runtime/operators/WindowAggOperator.java | 4 +-
.../typeutils/GlutenStatefulRecordSerializer.java | 56 ++++-------
19 files changed, 431 insertions(+), 174 deletions(-)
diff --git a/.github/workflows/flink.yml b/.github/workflows/flink.yml
index b53a8f62c5..cf607f5f30 100644
--- a/.github/workflows/flink.yml
+++ b/.github/workflows/flink.yml
@@ -59,10 +59,18 @@ jobs:
run: |
source /opt/rh/gcc-toolset-11/enable
sudo dnf install -y patchelf
- sudo yum install
https://mirror.stream.centos.org/9-stream/BaseOS/x86_64/os/Packages/tzdata-2025a-1.el9.noarch.rpm
-y
+ sudo yum install
https://mirror.stream.centos.org/9-stream/BaseOS/x86_64/os/Packages/tzdata-2026a-1.el9.noarch.rpm
-y
sudo .github/workflows/util/install-flink-deps.sh
+ # Remove system fmt v11 from the CI image to avoid header conflicts
with
+ # velox bundled fmt 10.x during folly/velox native compilation.
+ sudo rm -rf /usr/local/include/fmt
+ sudo rm -rf /usr/local/lib/cmake/fmt
+ sudo rm -f /usr/local/lib/libfmt* /usr/local/lib64/libfmt*
+ export VELOX_DEPENDENCY_SOURCE=BUNDLED
+ export fmt_SOURCE=BUNDLED
+ export folly_SOURCE=BUNDLED
git clone -b gluten-0530 https://github.com/bigo-sg/velox4j.git
- cd velox4j && git reset --hard
889bafcf2fa04e8c31a30edbdf40fe203ef58484
+ cd velox4j && git reset --hard
115edf79d265a61c30d45dfcc6ce932ad92378ca
git apply $GITHUB_WORKSPACE/gluten-flink/patches/fix-velox4j.patch
$GITHUB_WORKSPACE/build/mvn clean install -DskipTests -Dgpg.skip
-Dspotless.skip=true
cd ..
diff --git a/gluten-flink/docs/Flink.md b/gluten-flink/docs/Flink.md
index 1c800fb42d..ba8f4ec1dc 100644
--- a/gluten-flink/docs/Flink.md
+++ b/gluten-flink/docs/Flink.md
@@ -48,7 +48,7 @@ As some features have not been committed to upstream, you
have to use the follow
## fetch velox4j code
git clone -b gluten-0530 https://github.com/bigo-sg/velox4j.git
cd velox4j
-git reset --hard 889bafcf2fa04e8c31a30edbdf40fe203ef58484
+git reset --hard 115edf79d265a61c30d45dfcc6ce932ad92378ca
mvn clean install -DskipTests -Dgpg.skip -Dspotless.skip=true
```
**Get gluten**
diff --git
a/gluten-flink/planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/stream/StreamExecGlobalWindowAggregate.java
b/gluten-flink/planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/stream/StreamExecGlobalWindowAggregate.java
index 9cbc017232..e5dca4ee8d 100644
---
a/gluten-flink/planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/stream/StreamExecGlobalWindowAggregate.java
+++
b/gluten-flink/planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/stream/StreamExecGlobalWindowAggregate.java
@@ -19,7 +19,6 @@ package org.apache.flink.table.planner.plan.nodes.exec.stream;
import org.apache.gluten.rexnode.AggregateCallConverter;
import org.apache.gluten.rexnode.Utils;
import org.apache.gluten.rexnode.WindowUtils;
-import org.apache.gluten.table.runtime.operators.GlutenOneInputOperator;
import org.apache.gluten.util.LogicalTypeConverter;
import org.apache.gluten.util.PlanNodeIdGenerator;
@@ -37,14 +36,13 @@ import
io.github.zhztheplayer.velox4j.plan.StreamWindowPartitionFunctionSpec;
import org.apache.flink.FlinkVersion;
import org.apache.flink.api.dag.Transformation;
+import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple5;
import org.apache.flink.configuration.ReadableConfig;
import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
import org.apache.flink.streaming.api.operators.SimpleOperatorFactory;
import org.apache.flink.streaming.api.transformations.OneInputTransformation;
import org.apache.flink.table.data.RowData;
-import org.apache.flink.table.planner.codegen.CodeGeneratorContext;
-import org.apache.flink.table.planner.codegen.agg.AggsHandlerCodeGenerator;
import org.apache.flink.table.planner.delegation.PlannerBase;
import org.apache.flink.table.planner.plan.logical.WindowingStrategy;
import org.apache.flink.table.planner.plan.nodes.exec.ExecEdge;
@@ -55,24 +53,21 @@ import
org.apache.flink.table.planner.plan.nodes.exec.ExecNodeMetadata;
import org.apache.flink.table.planner.plan.nodes.exec.InputProperty;
import org.apache.flink.table.planner.plan.nodes.exec.utils.ExecNodeUtil;
import org.apache.flink.table.planner.plan.utils.AggregateInfoList;
+import org.apache.flink.table.planner.plan.utils.AggregateUtil;
import org.apache.flink.table.planner.plan.utils.KeySelectorUtil;
import org.apache.flink.table.planner.utils.JavaScalaConversionUtil;
import org.apache.flink.table.planner.utils.TableConfigUtils;
-import
org.apache.flink.table.runtime.generated.GeneratedNamespaceAggsHandleFunction;
import org.apache.flink.table.runtime.groupwindow.NamedWindowProperty;
-import org.apache.flink.table.runtime.groupwindow.WindowProperty;
import org.apache.flink.table.runtime.keyselector.RowDataKeySelector;
-import
org.apache.flink.table.runtime.operators.window.tvf.slicing.SliceAssigner;
import org.apache.flink.table.runtime.typeutils.InternalTypeInfo;
import org.apache.flink.table.runtime.util.TimeWindowUtil;
-import org.apache.flink.table.types.DataType;
+import org.apache.flink.table.types.logical.LogicalType;
import org.apache.flink.table.types.logical.RowType;
import
org.apache.flink.shaded.jackson2.com.fasterxml.jackson.annotation.JsonCreator;
import
org.apache.flink.shaded.jackson2.com.fasterxml.jackson.annotation.JsonProperty;
import org.apache.calcite.rel.core.AggregateCall;
-import org.apache.calcite.tools.RelBuilder;
import org.apache.commons.math3.util.ArithmeticUtils;
import javax.annotation.Nullable;
@@ -184,6 +179,14 @@ public class StreamExecGlobalWindowAggregate extends
StreamExecWindowAggregateBa
final ZoneId shiftTimeZone =
TimeWindowUtil.getShiftTimeZone(
windowing.getTimeAttributeType(),
TableConfigUtils.getLocalTimeZone(config));
+ final AggregateInfoList globalAggInfoList =
+ AggregateUtil.deriveStreamWindowAggregateInfoList(
+ planner.getTypeFactory(),
+ localAggInputRowType, // should use original input here
+ JavaScalaConversionUtil.toScala(Arrays.asList(aggCalls)),
+ needRetraction,
+ windowing.getWindow(),
+ true);
// --- Begin Gluten-specific code changes ---
// TODO: velox window not equal to flink window.
@@ -205,11 +208,15 @@ public class StreamExecGlobalWindowAggregate extends
StreamExecWindowAggregateBa
// TODO: support more window types.
Tuple5<Long, Long, Long, Integer, Integer> windowSpecParams =
WindowUtils.extractWindowParameters(windowing);
+ Tuple2<Integer, Integer> windowStartAndEndIndexes =
+ WindowUtils.getWindowStartAndEndIndexes(namedWindowProperties,
(RowType) getOutputType());
long size = windowSpecParams.f0;
long slide = windowSpecParams.f1;
long offset = windowSpecParams.f2;
int rowtimeIndex = windowSpecParams.f3;
int windowType = windowSpecParams.f4;
+ int windowStartIndex = windowStartAndEndIndexes.f0;
+ int windowEndIndex = windowStartAndEndIndexes.f1;
PartitionFunctionSpec sliceAssignerSpec =
new StreamWindowPartitionFunctionSpec(
inputType, rowtimeIndex, size, slide, offset, windowType);
@@ -252,24 +259,30 @@ public class StreamExecGlobalWindowAggregate extends
StreamExecWindowAggregateBa
offset,
windowType,
outputType,
- rowtimeIndex);
- final OneInputStreamOperator windowOperator =
- new GlutenOneInputOperator(
+ true,
+ rowtimeIndex,
+ windowStartIndex,
+ windowEndIndex);
+ final LogicalType[] accTypes =
convertToLogicalTypes(globalAggInfoList.getAccTypes());
+ final RowDataKeySelector selector =
+ KeySelectorUtil.getRowDataSelector(
+ planner.getFlinkContext().getClassLoader(),
+ grouping,
+ InternalTypeInfo.of(inputRowType));
+ final OneInputStreamOperator<RowData, RowData> windowOperator =
+ new
org.apache.gluten.table.runtime.operators.WindowAggOperator<RowData, RowData,
Long>(
new StatefulPlanNode(windowAgg.getId(), windowAgg),
PlanNodeIdGenerator.newId(),
inputType,
Map.of(windowAgg.getId(), outputType),
RowData.class,
RowData.class,
- "StreamExecGlobalWindowAggregate");
+ "StreamExecWindowAggregate",
+ selector.getProducedType(),
+ globalAggInfoList.getAggNames(),
+ accTypes);
// --- End Gluten-specific code changes ---
- final RowDataKeySelector selector =
- KeySelectorUtil.getRowDataSelector(
- planner.getFlinkContext().getClassLoader(),
- grouping,
- InternalTypeInfo.of(inputRowType));
-
final OneInputTransformation<RowData, RowData> transform =
ExecNodeUtil.createOneInputTransformation(
inputTransform,
@@ -285,40 +298,4 @@ public class StreamExecGlobalWindowAggregate extends
StreamExecWindowAggregateBa
transform.setStateKeyType(selector.getProducedType());
return transform;
}
-
- private GeneratedNamespaceAggsHandleFunction<Long> createAggsHandler(
- String name,
- SliceAssigner sliceAssigner,
- AggregateInfoList aggInfoList,
- int mergedAccOffset,
- boolean mergedAccIsOnHeap,
- DataType[] mergedAccExternalTypes,
- ExecNodeConfig config,
- ClassLoader classLoader,
- RelBuilder relBuilder,
- ZoneId shifTimeZone) {
- final AggsHandlerCodeGenerator generator =
- new AggsHandlerCodeGenerator(
- new CodeGeneratorContext(config, classLoader),
- relBuilder,
-
JavaScalaConversionUtil.toScala(localAggInputRowType.getChildren()),
- true) // copyInputField
- .needAccumulate()
- .needMerge(mergedAccOffset, mergedAccIsOnHeap,
mergedAccExternalTypes);
-
- final List<WindowProperty> windowProperties =
- Arrays.asList(
- Arrays.stream(namedWindowProperties)
- .map(NamedWindowProperty::getProperty)
- .toArray(WindowProperty[]::new));
-
- return generator.generateNamespaceAggsHandler(
- name,
- aggInfoList,
- JavaScalaConversionUtil.toScala(windowProperties),
- sliceAssigner,
- // we use window end timestamp to indicate a slicing window, see
SliceAssigner
- Long.class,
- shifTimeZone);
- }
}
diff --git
a/gluten-flink/planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/stream/StreamExecLocalWindowAggregate.java
b/gluten-flink/planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/stream/StreamExecLocalWindowAggregate.java
index fd0cc04a30..d67d61709e 100644
---
a/gluten-flink/planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/stream/StreamExecLocalWindowAggregate.java
+++
b/gluten-flink/planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/stream/StreamExecLocalWindowAggregate.java
@@ -220,7 +220,10 @@ public class StreamExecLocalWindowAggregate extends
StreamExecWindowAggregateBas
offset,
windowType,
outputType,
- rowtimeIndex);
+ false,
+ rowtimeIndex,
+ -1,
+ -1);
final OneInputStreamOperator localAggOperator =
new GlutenOneInputOperator(
new StatefulPlanNode(windowAgg.getId(), windowAgg),
diff --git
a/gluten-flink/planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/stream/StreamExecWindowAggregate.java
b/gluten-flink/planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/stream/StreamExecWindowAggregate.java
index 93e0703986..28b78687f6 100644
---
a/gluten-flink/planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/stream/StreamExecWindowAggregate.java
+++
b/gluten-flink/planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/stream/StreamExecWindowAggregate.java
@@ -36,6 +36,7 @@ import
io.github.zhztheplayer.velox4j.plan.StreamWindowPartitionFunctionSpec;
import org.apache.flink.FlinkVersion;
import org.apache.flink.api.dag.Transformation;
+import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple5;
import org.apache.flink.configuration.ReadableConfig;
import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
@@ -208,11 +209,16 @@ public class StreamExecWindowAggregate extends
StreamExecWindowAggregateBase {
// TODO: support more window types.
Tuple5<Long, Long, Long, Integer, Integer> windowSpecParams =
WindowUtils.extractWindowParameters(windowing);
+ Tuple2<Integer, Integer> windowStartAndEndIndexes =
+ WindowUtils.getWindowStartAndEndIndexes(namedWindowProperties,
(RowType) getOutputType());
long size = windowSpecParams.f0;
long slide = windowSpecParams.f1;
long offset = windowSpecParams.f2;
int rowtimeIndex = windowSpecParams.f3;
int windowType = windowSpecParams.f4;
+ boolean isRowTime = windowing.isRowtime();
+ int windowStartIndex = windowStartAndEndIndexes.f0;
+ int windowEndIndex = windowStartAndEndIndexes.f1;
PartitionFunctionSpec sliceAssignerSpec =
new StreamWindowPartitionFunctionSpec(
inputType, rowtimeIndex, size, slide, offset, windowType);
@@ -221,25 +227,29 @@ public class StreamExecWindowAggregate extends
StreamExecWindowAggregateBase {
PlanNodeIdGenerator.newId(),
AggregateStep.SINGLE,
groupingKeys,
- groupingKeys,
+ isRowTime ? groupingKeys : List.of(),
aggNames,
aggregates,
false,
List.of(new EmptyNode(inputType)),
null,
List.of());
+ // processing time window can not apply to local-global aggregate
optimization, so here we need
+ // to set local aggregtate as null when it is not event time window.
PlanNode localAgg =
- new AggregationNode(
- PlanNodeIdGenerator.newId(),
- AggregateStep.SINGLE,
- groupingKeys,
- groupingKeys,
- aggNames,
- aggregates,
- false,
- List.of(new EmptyNode(inputType)),
- null,
- List.of());
+ isRowTime
+ ? new AggregationNode(
+ PlanNodeIdGenerator.newId(),
+ AggregateStep.SINGLE,
+ groupingKeys,
+ groupingKeys,
+ aggNames,
+ aggregates,
+ false,
+ List.of(new EmptyNode(inputType)),
+ null,
+ List.of())
+ : null;
PlanNode windowAgg =
new StreamWindowAggregationNode(
PlanNodeIdGenerator.newId(),
@@ -255,7 +265,10 @@ public class StreamExecWindowAggregate extends
StreamExecWindowAggregateBase {
offset,
windowType,
outputType,
- rowtimeIndex);
+ windowing.isRowtime(),
+ rowtimeIndex,
+ windowStartIndex,
+ windowEndIndex);
final RowDataKeySelector selector =
KeySelectorUtil.getRowDataSelector(
planner.getFlinkContext().getClassLoader(),
@@ -266,6 +279,9 @@ public class StreamExecWindowAggregate extends
StreamExecWindowAggregateBase {
.map(x -> x.getLogicalType())
.collect(Collectors.toList())
.toArray(new LogicalType[] {});
+ // For TVF windows (Tumbling, Hopping, Cumulative, Session), the window
namespace
+ // is identified by the window end timestamp (Long). If count-based
windows are
+ // supported in the future, a different serializer may be needed.
final OneInputStreamOperator<RowData, RowData> windowOperator =
new
org.apache.gluten.table.runtime.operators.WindowAggOperator<RowData, RowData,
Long>(
new StatefulPlanNode(windowAgg.getId(), windowAgg),
diff --git
a/gluten-flink/planner/src/main/java/org/apache/gluten/rexnode/WindowUtils.java
b/gluten-flink/planner/src/main/java/org/apache/gluten/rexnode/WindowUtils.java
index 6f7a64f6f4..021969f2a2 100644
---
a/gluten-flink/planner/src/main/java/org/apache/gluten/rexnode/WindowUtils.java
+++
b/gluten-flink/planner/src/main/java/org/apache/gluten/rexnode/WindowUtils.java
@@ -16,16 +16,24 @@
*/
package org.apache.gluten.rexnode;
+import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple5;
+import org.apache.flink.table.planner.plan.logical.CumulativeWindowSpec;
import org.apache.flink.table.planner.plan.logical.HoppingWindowSpec;
+import org.apache.flink.table.planner.plan.logical.SessionWindowSpec;
import
org.apache.flink.table.planner.plan.logical.SliceAttachedWindowingStrategy;
import
org.apache.flink.table.planner.plan.logical.TimeAttributeWindowingStrategy;
import org.apache.flink.table.planner.plan.logical.TumblingWindowSpec;
import
org.apache.flink.table.planner.plan.logical.WindowAttachedWindowingStrategy;
import org.apache.flink.table.planner.plan.logical.WindowSpec;
import org.apache.flink.table.planner.plan.logical.WindowingStrategy;
+import org.apache.flink.table.runtime.groupwindow.NamedWindowProperty;
+import org.apache.flink.table.runtime.groupwindow.WindowEnd;
+import org.apache.flink.table.runtime.groupwindow.WindowStart;
+import org.apache.flink.table.types.logical.RowType;
import java.time.Duration;
+import java.util.List;
/** Utility to store some useful functions. */
public class WindowUtils {
@@ -53,24 +61,27 @@ public class WindowUtils {
if (windowOffset != null) {
offset = windowOffset.toMillis();
}
+ windowType = 0;
} else if (windowSpec instanceof TumblingWindowSpec) {
size = ((TumblingWindowSpec) windowSpec).getSize().toMillis();
Duration windowOffset = ((TumblingWindowSpec) windowSpec).getOffset();
if (windowOffset != null) {
offset = windowOffset.toMillis();
}
+ windowType = 1;
+ } else if (windowSpec instanceof CumulativeWindowSpec) {
+ windowType = 2;
+ } else if (windowSpec instanceof SessionWindowSpec) {
+ windowType = 3;
} else {
throw new RuntimeException("Not support window spec " + windowSpec);
}
-
if (windowing instanceof TimeAttributeWindowingStrategy) {
if (windowing.isRowtime()) {
rowtimeIndex = ((TimeAttributeWindowingStrategy)
windowing).getTimeAttributeIndex();
}
- windowType = 0;
} else if (windowing instanceof WindowAttachedWindowingStrategy) {
rowtimeIndex = ((WindowAttachedWindowingStrategy)
windowing).getWindowEnd();
- windowType = 1;
} else if (windowing instanceof SliceAttachedWindowingStrategy) {
rowtimeIndex = ((SliceAttachedWindowingStrategy)
windowing).getSliceEnd();
} else {
@@ -79,4 +90,18 @@ public class WindowUtils {
return new Tuple5<Long, Long, Long, Integer, Integer>(
size, slide, offset, rowtimeIndex, windowType);
}
+
+ public static Tuple2<Integer, Integer> getWindowStartAndEndIndexes(
+ NamedWindowProperty[] props, RowType outputType) {
+ int windowStartIndex = -1, windowEndIndex = -1;
+ List<String> outputNames = outputType.getFieldNames();
+ for (NamedWindowProperty prop : props) {
+ if (prop.getProperty() instanceof WindowStart) {
+ windowStartIndex = outputNames.indexOf(prop.getName());
+ } else if (prop.getProperty() instanceof WindowEnd) {
+ windowEndIndex = outputNames.indexOf(prop.getName());
+ }
+ }
+ return new Tuple2<Integer, Integer>(windowStartIndex, windowEndIndex);
+ }
}
diff --git
a/gluten-flink/runtime/src/main/java/org/apache/gluten/client/OffloadedJobGraphGenerator.java
b/gluten-flink/runtime/src/main/java/org/apache/gluten/client/OffloadedJobGraphGenerator.java
index 42784fce28..27dc0b1933 100644
---
a/gluten-flink/runtime/src/main/java/org/apache/gluten/client/OffloadedJobGraphGenerator.java
+++
b/gluten-flink/runtime/src/main/java/org/apache/gluten/client/OffloadedJobGraphGenerator.java
@@ -240,7 +240,7 @@ public class OffloadedJobGraphGenerator {
Class<?> inClass = supportsVectorInput ? StatefulRecord.class :
RowData.class;
Class<?> outClass = supportsVectorOutput ? StatefulRecord.class :
RowData.class;
GlutenOneInputOperator<?, ?> newOneInputOp =
- sourceOperator.cloneWithInputOutputClasses(inClass, outClass);
+ sourceOperator.cloneWithInputOutputClasses(planNode, inClass,
outClass);
offloadedOpConfig.setStreamOperator(newOneInputOp);
if (supportsVectorOutput) {
setOffloadedOutputSerializer(offloadedOpConfig, sourceOperator);
@@ -292,19 +292,19 @@ public class OffloadedJobGraphGenerator {
private void setOffloadedOutputSerializer(StreamConfig opConfig,
GlutenOperator operator) {
RowType rowType =
operator.getOutputTypes().entrySet().iterator().next().getValue();
- opConfig.setTypeSerializerOut(new GlutenStatefulRecordSerializer(rowType,
operator.getId()));
+ opConfig.setTypeSerializerOut(new GlutenStatefulRecordSerializer(rowType,
operator));
}
private void setOffloadedInputSerializer(StreamConfig opConfig,
GlutenOperator operator) {
opConfig.setupNetworkInputs(
- new GlutenStatefulRecordSerializer(operator.getInputType(),
operator.getId()));
+ new GlutenStatefulRecordSerializer(operator.getInputType(), operator));
}
private void setOffloadedInputSerializersForTwoInputOperator(
StreamConfig opConfig, GlutenTwoInputOperator<?, ?> operator) {
opConfig.setupNetworkInputs(
- new GlutenStatefulRecordSerializer(operator.getLeftInputType(),
operator.getId()),
- new GlutenStatefulRecordSerializer(operator.getRightInputType(),
operator.getId()));
+ new GlutenStatefulRecordSerializer(operator.getLeftInputType(),
operator),
+ new GlutenStatefulRecordSerializer(operator.getRightInputType(),
operator));
}
private void setOffloadedStatePartitioner(
diff --git
a/gluten-flink/runtime/src/main/java/org/apache/gluten/streaming/api/operators/GlutenOneInputOperatorFactory.java
b/gluten-flink/runtime/src/main/java/org/apache/gluten/streaming/api/operators/GlutenOneInputOperatorFactory.java
index 831dfde66f..41522731a3 100644
---
a/gluten-flink/runtime/src/main/java/org/apache/gluten/streaming/api/operators/GlutenOneInputOperatorFactory.java
+++
b/gluten-flink/runtime/src/main/java/org/apache/gluten/streaming/api/operators/GlutenOneInputOperatorFactory.java
@@ -16,6 +16,8 @@
*/
package org.apache.gluten.streaming.api.operators;
+import org.apache.gluten.table.runtime.operators.GlutenMailboxOperatorHelper;
+
import io.github.zhztheplayer.velox4j.serde.Serde;
import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
@@ -59,6 +61,10 @@ public class GlutenOneInputOperatorFactory<IN, OUT> extends
AbstractStreamOperat
.setup(
parameters.getContainingTask(), parameters.getStreamConfig(),
parameters.getOutput());
}
+ if (operator instanceof GlutenOperator) {
+ GlutenMailboxOperatorHelper.bindAtTaskStartup(
+ ((GlutenOperator) operator).mailboxHolder(), parameters);
+ }
return (T) operator;
}
diff --git
a/gluten-flink/runtime/src/main/java/org/apache/gluten/streaming/api/operators/GlutenOperator.java
b/gluten-flink/runtime/src/main/java/org/apache/gluten/streaming/api/operators/GlutenOperator.java
index 9720551967..bb6bc3bf54 100644
---
a/gluten-flink/runtime/src/main/java/org/apache/gluten/streaming/api/operators/GlutenOperator.java
+++
b/gluten-flink/runtime/src/main/java/org/apache/gluten/streaming/api/operators/GlutenOperator.java
@@ -16,9 +16,15 @@
*/
package org.apache.gluten.streaming.api.operators;
+import org.apache.gluten.table.runtime.operators.GlutenMailboxHolder;
+import org.apache.gluten.table.runtime.operators.GlutenSessionResources;
+
import io.github.zhztheplayer.velox4j.plan.StatefulPlanNode;
import io.github.zhztheplayer.velox4j.type.RowType;
+import org.apache.flink.api.common.operators.MailboxExecutor;
+import org.apache.flink.streaming.runtime.tasks.StreamTask;
+
import java.util.Map;
/** Interface for all gluten operators. */
@@ -34,4 +40,43 @@ public interface GlutenOperator {
public default String getDescription() {
return "";
}
+
+ /** Mailbox drain helper holder; must be a non-transient field on the
concrete operator. */
+ default GlutenMailboxHolder mailboxHolder() {
+ return new GlutenMailboxHolder();
+ }
+
+ default void processElementInternal() {}
+
+ default void bindMailboxExecutor(MailboxExecutor mailboxExecutor) {
+ mailboxHolder().get().bindMailboxExecutor(mailboxExecutor);
+ }
+
+ default void ensureMailboxInitialized(StreamTask<?, ?> containingTask) {
+ mailboxHolder().get().ensureMailboxInitialized(containingTask);
+ }
+
+ default void drainOutput(Runnable drainAction) {
+ mailboxHolder().get().runDrain(drainAction);
+ }
+
+ default void scheduleDrainOnMailbox(Runnable drainAction) {
+ mailboxHolder().get().scheduleDrain(drainAction);
+ }
+
+ /**
+ * Called from native Velox code to drain operator output. Drain is always
scheduled on the Flink
+ * task mailbox thread.
+ */
+ static void processElementByJni(String operatorId) {
+ GlutenOperator operator =
+
GlutenSessionResources.getInstance().getOperator(operatorId).orElse(null);
+ if (operator == null) {
+ throw new IllegalArgumentException("Operator not found: " + operatorId);
+ }
+ operator.scheduleProcessElementOnMailbox();
+ }
+
+ /** Schedules native output drain on the mailbox thread. Implemented by
concrete operators. */
+ default void scheduleProcessElementOnMailbox() {}
}
diff --git
a/gluten-flink/runtime/src/main/java/org/apache/gluten/streaming/runtime/tasks/GlutenOutputCollector.java
b/gluten-flink/runtime/src/main/java/org/apache/gluten/streaming/runtime/tasks/GlutenOutputCollector.java
index 60dc8b62c6..8da53faa93 100644
---
a/gluten-flink/runtime/src/main/java/org/apache/gluten/streaming/runtime/tasks/GlutenOutputCollector.java
+++
b/gluten-flink/runtime/src/main/java/org/apache/gluten/streaming/runtime/tasks/GlutenOutputCollector.java
@@ -29,10 +29,8 @@ import
org.apache.flink.streaming.runtime.tasks.OutputWithChainingCheck;
import org.apache.flink.streaming.runtime.tasks.WatermarkGaugeExposingOutput;
import org.apache.flink.streaming.runtime.watermarkstatus.WatermarkStatus;
import org.apache.flink.util.OutputTag;
-import org.apache.flink.util.XORShiftRandom;
import java.util.Map;
-import java.util.Random;
/**
* Collector for gluten, it may contain several outputs, dispatch the record
according to node id.
@@ -40,7 +38,6 @@ import java.util.Random;
public class GlutenOutputCollector<T> implements
WatermarkGaugeExposingOutput<StreamRecord<T>> {
protected final Map<String, OutputWithChainingCheck<StreamRecord<T>>>
outputs;
- private final Random random = new XORShiftRandom();
private final WatermarkGauge watermarkGauge = new WatermarkGauge();
protected final Counter numRecordsOutForTask;
diff --git
a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/config/VeloxQueryConfig.java
b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/config/VeloxQueryConfig.java
index 788ecb4a68..a1026cfe49 100644
---
a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/config/VeloxQueryConfig.java
+++
b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/config/VeloxQueryConfig.java
@@ -31,6 +31,8 @@ public class VeloxQueryConfig {
private static final String keyVeloxAdjustTimestampToSessionTimeZone =
"adjust_timestamp_to_session_timezone";
private static final String keyVeloxSessionTimezone = "session_timezone";
+ private static final String kStreamingAggregationMinOutputBatchRows =
+ "streaming_aggregation_min_output_batch_rows";
public static Config getConfig(RuntimeContext context) {
if (!(context instanceof StreamingRuntimeContext)) {
@@ -47,6 +49,7 @@ public class VeloxQueryConfig {
} else {
configMap.put(keyVeloxSessionTimezone, localTimeZone);
}
+ configMap.put(kStreamingAggregationMinOutputBatchRows, String.valueOf(1));
return Config.create(configMap);
}
}
diff --git
a/gluten-flink/runtime/src/main/java/org/apache/gluten/streaming/api/operators/GlutenOperator.java
b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenMailboxHolder.java
similarity index 57%
copy from
gluten-flink/runtime/src/main/java/org/apache/gluten/streaming/api/operators/GlutenOperator.java
copy to
gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenMailboxHolder.java
index 9720551967..9415fcfe36 100644
---
a/gluten-flink/runtime/src/main/java/org/apache/gluten/streaming/api/operators/GlutenOperator.java
+++
b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenMailboxHolder.java
@@ -14,24 +14,26 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.gluten.streaming.api.operators;
+package org.apache.gluten.table.runtime.operators;
-import io.github.zhztheplayer.velox4j.plan.StatefulPlanNode;
-import io.github.zhztheplayer.velox4j.type.RowType;
+import java.io.Serializable;
-import java.util.Map;
-
-/** Interface for all gluten operators. */
-public interface GlutenOperator {
- public StatefulPlanNode getPlanNode();
-
- public RowType getInputType();
+/**
+ * Serializable holder for {@link GlutenMailboxOperatorHelper}.
+ *
+ * <p>The holder survives Flink operator deserialization; the helper is lazily
recreated because it
+ * is {@code transient}.
+ */
+public final class GlutenMailboxHolder implements Serializable {
- public Map<String, RowType> getOutputTypes();
+ private static final long serialVersionUID = 1L;
- public String getId();
+ private transient GlutenMailboxOperatorHelper helper;
- public default String getDescription() {
- return "";
+ public GlutenMailboxOperatorHelper get() {
+ if (helper == null) {
+ helper = new GlutenMailboxOperatorHelper();
+ }
+ return helper;
}
}
diff --git
a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenMailboxOperatorHelper.java
b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenMailboxOperatorHelper.java
new file mode 100644
index 0000000000..ecec25ede6
--- /dev/null
+++
b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenMailboxOperatorHelper.java
@@ -0,0 +1,108 @@
+/*
+ * 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.gluten.table.runtime.operators;
+
+import org.apache.flink.api.common.operators.MailboxExecutor;
+import org.apache.flink.streaming.api.operators.StreamOperatorParameters;
+import org.apache.flink.streaming.runtime.tasks.StreamTask;
+
+/**
+ * Schedules Velox output drain onto the Flink task mailbox thread.
+ *
+ * <p>JNI callbacks from Velox must not call {@code SerialTask.advance()} or
emit records directly.
+ * When a callback arrives while a drain is already in progress on the mailbox
thread, a follow-up
+ * drain is requested via {@code pendingMailboxDrain} instead of nesting
another drain.
+ */
+public final class GlutenMailboxOperatorHelper {
+
+ private transient MailboxExecutor mailboxExecutor;
+ private transient boolean draining;
+ private transient boolean pendingMailboxDrain;
+ private transient boolean mailboxDrainScheduled;
+
+ public void bindMailboxExecutor(MailboxExecutor mailboxExecutor) {
+ this.mailboxExecutor = mailboxExecutor;
+ }
+
+ public void ensureMailboxInitialized(StreamTask<?, ?> containingTask) {
+ if (mailboxExecutor == null) {
+ mailboxExecutor =
containingTask.getMailboxExecutorFactory().createExecutor(0);
+ }
+ }
+
+ public boolean isMailboxBound() {
+ return mailboxExecutor != null;
+ }
+
+ public void runDrain(Runnable drainAction) {
+ draining = true;
+ try {
+ boolean repeat;
+ do {
+ pendingMailboxDrain = false;
+ drainAction.run();
+ repeat = pendingMailboxDrain;
+ } while (repeat);
+ } finally {
+ draining = false;
+ }
+ }
+
+ public void scheduleDrain(Runnable drainAction) {
+ if (mailboxExecutor == null) {
+ runDrain(drainAction);
+ return;
+ }
+ if (draining) {
+ pendingMailboxDrain = true;
+ return;
+ }
+ synchronized (this) {
+ if (mailboxDrainScheduled) {
+ return;
+ }
+ mailboxDrainScheduled = true;
+ }
+ try {
+ mailboxExecutor.submit(
+ () -> {
+ try {
+ runDrain(drainAction);
+ } finally {
+ synchronized (GlutenMailboxOperatorHelper.this) {
+ mailboxDrainScheduled = false;
+ }
+ }
+ },
+ "gluten-drain-output");
+ } catch (RuntimeException e) {
+ synchronized (this) {
+ mailboxDrainScheduled = false;
+ }
+ throw e;
+ }
+ }
+
+ /** Binds mailbox from {@link StreamOperatorParameters} during operator
factory startup on TM. */
+ public static void bindAtTaskStartup(
+ GlutenMailboxHolder mailboxHolder, StreamOperatorParameters<?>
parameters) {
+ mailboxHolder
+ .get()
+ .bindMailboxExecutor(
+
parameters.getContainingTask().getMailboxExecutorFactory().createExecutor(0));
+ }
+}
diff --git
a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenOneInputOperator.java
b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenOneInputOperator.java
index a2733c2438..09b059f3bc 100644
---
a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenOneInputOperator.java
+++
b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenOneInputOperator.java
@@ -36,6 +36,7 @@ import io.github.zhztheplayer.velox4j.stateful.StatefulRecord;
import io.github.zhztheplayer.velox4j.stateful.StatefulWatermark;
import io.github.zhztheplayer.velox4j.type.RowType;
+import org.apache.flink.contrib.streaming.state.RocksDBKeyedStateBackend;
import org.apache.flink.runtime.state.StateInitializationContext;
import org.apache.flink.runtime.state.StateSnapshotContext;
import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
@@ -43,9 +44,6 @@ import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.table.runtime.operators.TableStreamOperator;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
import java.util.List;
import java.util.Map;
@@ -53,8 +51,6 @@ import java.util.Map;
public class GlutenOneInputOperator<IN, OUT> extends TableStreamOperator<OUT>
implements OneInputStreamOperator<IN, OUT>, GlutenOperator {
- private static final Logger LOG =
LoggerFactory.getLogger(GlutenOneInputOperator.class);
-
private final StatefulPlanNode glutenPlan;
private final String id;
private final RowType inputType;
@@ -70,6 +66,7 @@ public class GlutenOneInputOperator<IN, OUT> extends
TableStreamOperator<OUT>
private final Class<OUT> outClass;
private transient VectorInputBridge<IN> inputBridge;
private transient VectorOutputBridge<OUT> outputBridge;
+ private final GlutenMailboxHolder mailboxHolder = new GlutenMailboxHolder();
public GlutenOneInputOperator(
StatefulPlanNode plan,
@@ -121,6 +118,8 @@ public class GlutenOneInputOperator<IN, OUT> extends
TableStreamOperator<OUT>
outputBridge = VectorOutputBridge.Factory.create(outClass);
}
sessionResource = new GlutenSessionResource();
+ GlutenSessionResources.getInstance().addSessionResource(id,
sessionResource);
+
GlutenSessionResources.getInstance().addOperator(this.getClass().getSimpleName(),
this);
inputQueue =
sessionResource.getSession().externalStreamOps().newBlockingQueue();
// add a mock input as velox not allow the source is empty.
if (inputType == null) {
@@ -152,6 +151,9 @@ public class GlutenOneInputOperator<IN, OUT> extends
TableStreamOperator<OUT>
@Override
public void open() throws Exception {
super.open();
+ if (!mailboxHolder().get().isMailboxBound()) {
+ ensureMailboxInitialized(getContainingTask());
+ }
initSession();
}
@@ -163,28 +165,48 @@ public class GlutenOneInputOperator<IN, OUT> extends
TableStreamOperator<OUT>
StatefulRecord statefulRecord =
inputBridge.convertToStatefulRecord(
element, sessionResource.getAllocator(),
sessionResource.getSession(), inputType);
- inputQueue.put(statefulRecord.getRowVector());
-
- // Only the rowvectors generated by this operator should be closed here.
- if (getId().equals(statefulRecord.getNodeId())) {
- statefulRecord.close();
+ try {
+ inputQueue.put(statefulRecord.getRowVector());
+ } finally {
+ // Only the rowvectors generated by this operator should be closed here.
+ if (getId().equals(statefulRecord.getNodeId())) {
+ statefulRecord.close();
+ }
}
processElementInternal();
}
- private void processElementInternal() {
+ @Override
+ public GlutenMailboxHolder mailboxHolder() {
+ return mailboxHolder;
+ }
+
+ @Override
+ public void scheduleProcessElementOnMailbox() {
+ scheduleDrainOnMailbox(this::drainTaskOutput);
+ }
+
+ @Override
+ public void processElementInternal() {
+ drainOutput(this::drainTaskOutput);
+ }
+
+ private void drainTaskOutput() {
while (true) {
UpIterator.State state = task.advance();
if (state == UpIterator.State.AVAILABLE) {
final StatefulElement statefulElement = task.statefulGet();
- if (statefulElement.isWatermark()) {
- StatefulWatermark watermark = statefulElement.asWatermark();
- output.emitWatermark(new Watermark(watermark.getTimestamp()));
- } else {
- outputBridge.collect(
- output, statefulElement.asRecord(),
sessionResource.getAllocator(), outputType);
+ try {
+ if (statefulElement.isWatermark()) {
+ StatefulWatermark watermark = statefulElement.asWatermark();
+ output.emitWatermark(new Watermark(watermark.getTimestamp()));
+ } else {
+ outputBridge.collect(
+ output, statefulElement.asRecord(),
sessionResource.getAllocator(), outputType);
+ }
+ } finally {
+ statefulElement.close();
}
- statefulElement.close();
} else {
break;
}
@@ -192,15 +214,9 @@ public class GlutenOneInputOperator<IN, OUT> extends
TableStreamOperator<OUT>
}
public <NIN, NOUT> GlutenOneInputOperator<NIN, NOUT>
cloneWithInputOutputClasses(
- Class<NIN> newInClass, Class<NOUT> newOutClass) {
+ StatefulPlanNode plan, Class<NIN> newInClass, Class<NOUT> newOutClass) {
return new GlutenOneInputOperator<>(
- this.glutenPlan,
- this.id,
- this.inputType,
- this.outputTypes,
- newInClass,
- newOutClass,
- this.description);
+ plan, this.id, this.inputType, this.outputTypes, newInClass,
newOutClass, this.description);
}
@Override
@@ -222,13 +238,13 @@ public class GlutenOneInputOperator<IN, OUT> extends
TableStreamOperator<OUT>
@Override
public void close() throws Exception {
+ if (task != null) {
+ task.close();
+ }
if (inputQueue != null) {
inputQueue.noMoreInput();
inputQueue.close();
}
- if (task != null) {
- task.close();
- }
if (sessionResource != null) {
sessionResource.close();
}
@@ -272,7 +288,9 @@ public class GlutenOneInputOperator<IN, OUT> extends
TableStreamOperator<OUT>
if (task == null) {
initSession();
}
- task.initializeState(0, null);
+ if (!(getKeyedStateBackend() instanceof RocksDBKeyedStateBackend)) {
+ task.initializeState(0, null);
+ }
super.initializeState(context);
}
diff --git
a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenSessionResource.java
b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenSessionResources.java
similarity index 67%
rename from
gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenSessionResource.java
rename to
gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenSessionResources.java
index ea38229e95..fac0784a24 100644
---
a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenSessionResource.java
+++
b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenSessionResources.java
@@ -16,6 +16,8 @@
*/
package org.apache.gluten.table.runtime.operators;
+import org.apache.gluten.streaming.api.operators.GlutenOperator;
+
import io.github.zhztheplayer.velox4j.Velox4j;
import io.github.zhztheplayer.velox4j.memory.AllocationListener;
import io.github.zhztheplayer.velox4j.memory.MemoryManager;
@@ -26,6 +28,10 @@ import org.apache.flink.runtime.state.KeyedStateBackend;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+
// Manage the session and resource for Velox.
class GlutenSessionResource {
private Session session;
@@ -74,3 +80,38 @@ class GlutenSessionResource {
this.keyedStateBackend = keyedStateBackend;
}
}
+
+public class GlutenSessionResources {
+ private static final GlutenSessionResources instance = new
GlutenSessionResources();
+ private Map<String, GlutenSessionResource> sessionResources = new
HashMap<>();
+ private Map<String, GlutenOperator> operators = new HashMap<>();
+
+ private GlutenSessionResources() {}
+
+ public static GlutenSessionResources getInstance() {
+ return instance;
+ }
+
+ public GlutenSessionResource getSessionResource(String id) {
+ return sessionResources.get(id);
+ }
+
+ public void addSessionResource(String id, GlutenSessionResource
sessionResource) {
+ sessionResources.put(id, sessionResource);
+ }
+
+ public Session getSession(String id) {
+ return sessionResources.get(id).getSession();
+ }
+
+ public void addOperator(String id, GlutenOperator operator) {
+ operators.put(id, operator);
+ }
+
+ public Optional<GlutenOperator> getOperator(String id) {
+ if (operators.containsKey(id)) {
+ return Optional.of(operators.get(id));
+ }
+ return Optional.empty();
+ }
+}
diff --git
a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunction.java
b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunction.java
index ea0ddcbc7c..53f36fcf67 100644
---
a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunction.java
+++
b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunction.java
@@ -109,10 +109,8 @@ public class GlutenSourceFunction<OUT> extends
RichParallelSourceFunction<OUT>
processAvailableElement(sourceContext);
break;
case BLOCKED:
- LOG.debug("Get empty row");
break;
default:
- LOG.info("Velox task finished");
return;
}
taskMetrics.updateMetrics(task, id);
@@ -233,6 +231,7 @@ public class GlutenSourceFunction<OUT> extends
RichParallelSourceFunction<OUT>
return;
}
sessionResource = new GlutenSessionResource();
+ GlutenSessionResources.getInstance().addSessionResource(id,
sessionResource);
Session session = sessionResource.getSession();
query =
new Query(
diff --git
a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenTwoInputOperator.java
b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenTwoInputOperator.java
index 2352d74943..3f73ad4bbd 100644
---
a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenTwoInputOperator.java
+++
b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenTwoInputOperator.java
@@ -72,6 +72,7 @@ public class GlutenTwoInputOperator<IN, OUT> extends
AbstractStreamOperator<OUT>
private VectorInputBridge<IN> inputBridge;
private VectorOutputBridge<OUT> outputBridge;
private String description;
+ private final GlutenMailboxHolder mailboxHolder = new GlutenMailboxHolder();
public GlutenTwoInputOperator(
StatefulPlanNode plan,
@@ -117,6 +118,9 @@ public class GlutenTwoInputOperator<IN, OUT> extends
AbstractStreamOperator<OUT>
@Override
public void open() throws Exception {
super.open();
+ if (!mailboxHolder().get().isMailboxBound()) {
+ ensureMailboxInitialized(getContainingTask());
+ }
initSession();
}
@@ -125,6 +129,16 @@ public class GlutenTwoInputOperator<IN, OUT> extends
AbstractStreamOperator<OUT>
return glutenPlan.getId();
}
+ @Override
+ public GlutenMailboxHolder mailboxHolder() {
+ return mailboxHolder;
+ }
+
+ @Override
+ public void scheduleProcessElementOnMailbox() {
+ scheduleDrainOnMailbox(this::drainTaskOutput);
+ }
+
@Override
public void processElement1(StreamRecord<IN> element) {
StatefulRecord statefulRecord =
@@ -152,19 +166,27 @@ public class GlutenTwoInputOperator<IN, OUT> extends
AbstractStreamOperator<OUT>
processElementInternal();
}
- private void processElementInternal() {
+ @Override
+ public void processElementInternal() {
+ drainOutput(this::drainTaskOutput);
+ }
+
+ private void drainTaskOutput() {
while (true) {
UpIterator.State state = task.advance();
if (state == UpIterator.State.AVAILABLE) {
final StatefulElement element = task.statefulGet();
- if (element.isWatermark()) {
- StatefulWatermark watermark = element.asWatermark();
- output.emitWatermark(new Watermark(watermark.getTimestamp()));
- } else {
- outputBridge.collect(
- output, element.asRecord(), sessionResource.getAllocator(),
outputType);
+ try {
+ if (element.isWatermark()) {
+ StatefulWatermark watermark = element.asWatermark();
+ output.emitWatermark(new Watermark(watermark.getTimestamp()));
+ } else {
+ outputBridge.collect(
+ output, element.asRecord(), sessionResource.getAllocator(),
outputType);
+ }
+ } finally {
+ element.close();
}
- element.close();
} else {
break;
}
@@ -263,7 +285,8 @@ public class GlutenTwoInputOperator<IN, OUT> extends
AbstractStreamOperator<OUT>
}
sessionResource = new GlutenSessionResource();
-
+ GlutenSessionResources.getInstance().addSessionResource(getId(),
sessionResource);
+ GlutenSessionResources.getInstance().addOperator(getId(), this);
leftInputQueue =
sessionResource.getSession().externalStreamOps().newBlockingQueue();
rightInputQueue =
sessionResource.getSession().externalStreamOps().newBlockingQueue();
diff --git
a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/WindowAggOperator.java
b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/WindowAggOperator.java
index 50d91f2c73..f15d4c8c39 100644
---
a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/WindowAggOperator.java
+++
b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/WindowAggOperator.java
@@ -129,12 +129,14 @@ public class WindowAggOperator<IN, OUT, W> extends
GlutenOneInputOperator<IN, OU
new
org.apache.flink.table.types.logical.RowType(accFields))),
Map.of(windowStateName, new BigIntType()));
task.initializeState(0, parameters);
+ } else {
+ task.initializeState(0, null);
}
}
@Override
public <NIN, NOUT> WindowAggOperator<NIN, NOUT, W>
cloneWithInputOutputClasses(
- Class<NIN> newInClass, Class<NOUT> newOutClass) {
+ StatefulPlanNode plan, Class<NIN> newInClass, Class<NOUT> newOutClass) {
return new WindowAggOperator<>(
getPlanNode(),
getId(),
diff --git
a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/typeutils/GlutenStatefulRecordSerializer.java
b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/typeutils/GlutenStatefulRecordSerializer.java
index 1040e776c1..66fdd7bc51 100644
---
a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/typeutils/GlutenStatefulRecordSerializer.java
+++
b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/typeutils/GlutenStatefulRecordSerializer.java
@@ -16,11 +16,10 @@
*/
package org.apache.gluten.table.runtime.typeutils;
-import io.github.zhztheplayer.velox4j.Velox4j;
+import org.apache.gluten.streaming.api.operators.GlutenOperator;
+import org.apache.gluten.table.runtime.operators.GlutenSessionResources;
+
import io.github.zhztheplayer.velox4j.data.RowVector;
-import io.github.zhztheplayer.velox4j.memory.AllocationListener;
-import io.github.zhztheplayer.velox4j.memory.MemoryManager;
-import io.github.zhztheplayer.velox4j.session.Session;
import io.github.zhztheplayer.velox4j.stateful.StatefulRecord;
import io.github.zhztheplayer.velox4j.type.RowType;
@@ -31,31 +30,23 @@ import
org.apache.flink.api.common.typeutils.TypeSerializerSnapshot;
import org.apache.flink.core.memory.DataInputView;
import org.apache.flink.core.memory.DataOutputView;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.io.Closeable;
import java.io.IOException;
/** Serializer for {@link RowVector}. */
@Internal
-public class GlutenStatefulRecordSerializer extends
TypeSerializer<StatefulRecord>
- implements Closeable {
- private static final Logger LOG =
LoggerFactory.getLogger(GlutenStatefulRecordSerializer.class);
+public class GlutenStatefulRecordSerializer extends
TypeSerializer<StatefulRecord> {
private static final long serialVersionUID = 1L;
private final RowType rowType;
- private transient MemoryManager memoryManager;
- private transient Session session;
- private final String nodeId;
+ private final GlutenOperator operator;
- public GlutenStatefulRecordSerializer(RowType rowType, String nodeId) {
- this.nodeId = nodeId;
+ public GlutenStatefulRecordSerializer(RowType rowType, GlutenOperator
operator) {
this.rowType = rowType;
+ this.operator = operator;
}
@Override
public TypeSerializer<StatefulRecord> duplicate() {
- return new GlutenStatefulRecordSerializer(rowType, nodeId);
+ return new GlutenStatefulRecordSerializer(rowType, operator);
}
@Override
@@ -72,15 +63,16 @@ public class GlutenStatefulRecordSerializer extends
TypeSerializer<StatefulRecor
@Override
public StatefulRecord deserialize(DataInputView source) throws IOException {
- if (memoryManager == null) {
- memoryManager = MemoryManager.create(AllocationListener.NOOP);
- session = Velox4j.newSession(memoryManager);
- }
int len = source.readInt();
byte[] str = new byte[len];
source.readFully(str);
- RowVector rowVector = session.baseVectorOps().deserializeOne(new
String(str)).asRowVector();
- StatefulRecord record = new StatefulRecord(nodeId, 0, 0, false, -1);
+ RowVector rowVector =
+ GlutenSessionResources.getInstance()
+ .getSession(operator.getId())
+ .baseVectorOps()
+ .deserializeOne(new String(str))
+ .asRowVector();
+ StatefulRecord record = new StatefulRecord(operator.getId(),
rowVector.id(), 0, false, -1);
record.setRowVector(rowVector);
return record;
}
@@ -138,15 +130,7 @@ public class GlutenStatefulRecordSerializer extends
TypeSerializer<StatefulRecor
@Override
public TypeSerializerSnapshot<StatefulRecord> snapshotConfiguration() {
- return new RowVectorSerializerSnapshot(rowType, nodeId);
- }
-
- @Override
- public void close() {
- if (memoryManager != null) {
- memoryManager.close();
- session.close();
- }
+ return new RowVectorSerializerSnapshot(rowType, operator);
}
/** {@link TypeSerializerSnapshot} for Gluten RowVector.. */
@@ -155,16 +139,16 @@ public class GlutenStatefulRecordSerializer extends
TypeSerializer<StatefulRecor
private static final int CURRENT_VERSION = 1;
private RowType rowType;
- private String nodeId;
+ private GlutenOperator operator;
@SuppressWarnings("unused")
public RowVectorSerializerSnapshot() {
// this constructor is used when restoring from a checkpoint/savepoint.
}
- RowVectorSerializerSnapshot(RowType rowType, String nodeId) {
- this.nodeId = nodeId;
+ RowVectorSerializerSnapshot(RowType rowType, GlutenOperator operator) {
this.rowType = rowType;
+ this.operator = operator;
}
@Override
@@ -181,7 +165,7 @@ public class GlutenStatefulRecordSerializer extends
TypeSerializer<StatefulRecor
@Override
public GlutenStatefulRecordSerializer restoreSerializer() {
- return new GlutenStatefulRecordSerializer(rowType, nodeId);
+ return new GlutenStatefulRecordSerializer(rowType, operator);
}
@Override
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
