This is an automated email from the ASF dual-hosted git repository. lgbo-ustc pushed a commit to branch feature/idle-source-handling in repository https://gitbox.apache.org/repos/asf/gluten.git
commit 5d08f10fb349c223af88b714f211ea8a7c021b6c Author: lgbo-ustc <[email protected]> AuthorDate: Tue Jul 7 14:39:54 2026 +0800 fix: Add shouldCallNoMoreSplits option to GlutenSourceFunction for unbounded test scenarios --- .../runtime/operators/GlutenSourceFunction.java | 16 +++++++- ...GlutenSourceFunctionWatermarkStatusE2ETest.java | 43 ++++++++++++++++++++++ 2 files changed, 58 insertions(+), 1 deletion(-) 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 1217cebdf7..a7ef34f32a 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 @@ -62,6 +62,13 @@ public class GlutenSourceFunction<OUT> extends RichParallelSourceFunction<OUT> private final ConnectorSplit split; private volatile boolean isRunning = true; + /** + * If true, {@code noMoreSplits()} is called after adding the initial split in {@link + * #initSession()}. Set to false for unbounded streaming test scenarios where the task should stay + * alive to detect idleness. + */ + private boolean shouldCallNoMoreSplits = true; + private GlutenSessionResource sessionResource; private Query query; private SerialTask task; @@ -81,6 +88,11 @@ public class GlutenSourceFunction<OUT> extends RichParallelSourceFunction<OUT> this.outClass = outClass; } + /** Sets whether noMoreSplits() should be called after adding the initial split. */ + public void setShouldCallNoMoreSplits(boolean value) { + this.shouldCallNoMoreSplits = value; + } + public StatefulPlanNode getPlanNode() { return planNode; } @@ -264,7 +276,9 @@ public class GlutenSourceFunction<OUT> extends RichParallelSourceFunction<OUT> VeloxConnectorConfig.getConfig(getRuntimeContext())); task = session.queryOps().execute(query); task.addSplit(id, activeSplit); - task.noMoreSplits(id); + if (shouldCallNoMoreSplits) { + task.noMoreSplits(id); + } taskMetrics = new SourceTaskMetrics(getRuntimeContext().getMetricGroup()); } } diff --git a/gluten-flink/ut/src/test/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunctionWatermarkStatusE2ETest.java b/gluten-flink/ut/src/test/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunctionWatermarkStatusE2ETest.java new file mode 100644 index 0000000000..9ae81ed71e --- /dev/null +++ b/gluten-flink/ut/src/test/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunctionWatermarkStatusE2ETest.java @@ -0,0 +1,43 @@ +/* + * 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.gluten.table.runtime.stream.common.Velox4jEnvironment; + +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +import java.lang.reflect.Method; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests that GlutenSourceFunction correctly handles WatermarkStatus elements. */ +class GlutenSourceFunctionWatermarkStatusE2ETest { + + @BeforeAll + static void setupGluten() { + Velox4jEnvironment.initializeOnce(); + } + + @Test + void testGlutenSourceFunctionClassHasWatermarkStatusHandling() throws Exception { + // Verify GlutenSourceFunction has the setShouldCallNoMoreSplits method + Method setter = + GlutenSourceFunction.class.getMethod("setShouldCallNoMoreSplits", boolean.class); + assertThat(setter).isNotNull(); + } +} --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
