From abce609f8a861c6f263c0c34c5f22535e1c0c50f Mon Sep 17 00:00:00 2001 From: lgbo-ustc Date: Tue, 7 Jul 2026 09:43:07 +0800 Subject: [PATCH 1/9] fix: Handle WatermarkStatus elements in GlutenSourceFunction Add processing for WatermarkStatus elements from native idle detection. When IDLE is received, call sourceContext.markAsTemporarilyIdle() to notify Flink that this source is temporarily idle, allowing watermark progress to continue from other sources. --- .../runtime/operators/GlutenSourceFunction.java | 15 ++++++++++++++- 1 file changed, 14 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 a47b6c7c17..e7dfab6f09 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 @@ -31,6 +31,7 @@ import io.github.zhztheplayer.velox4j.stateful.StatefulElement; import io.github.zhztheplayer.velox4j.stateful.StatefulRecord; import io.github.zhztheplayer.velox4j.stateful.StatefulWatermark; +import io.github.zhztheplayer.velox4j.stateful.StatefulWatermarkStatus; import io.github.zhztheplayer.velox4j.type.RowType; import org.apache.flink.api.common.state.ListState; @@ -132,8 +133,10 @@ private void processAvailableElement(SourceContext sourceContext) { processRecord(sourceContext, element.asRecord()); } else if (element.isWatermark()) { processWatermark(sourceContext, element.asWatermark()); + } else if (element.isWatermarkStatus()) { + processWatermarkStatus(sourceContext, element.asWatermarkStatus()); } else { - LOG.debug("Ignoring element that is neither record nor watermark"); + LOG.debug("Ignoring element that is neither record, watermark, nor watermark status"); } } finally { element.close(); @@ -169,6 +172,16 @@ private void processWatermark(SourceContext sourceContext, StatefulWatermar sourceContext.emitWatermark(new Watermark(watermark.getTimestamp())); } + /** Processes a watermark status and notifies the source context about idleness. */ + private void processWatermarkStatus( + SourceContext sourceContext, StatefulWatermarkStatus status) { + if (status.isIdle()) { + sourceContext.markAsTemporarilyIdle(); + } + // ACTIVE: no explicit action needed; the source context will resume + // activity tracking when the next record or watermark is emitted. + } + /** Collects a StatefulRecord as RowData by converting the RowVector. */ private void collectAsRowData(SourceContext sourceContext, StatefulRecord record) { List rows = From 3bcf527e74ccde428d8f75da8789213d5f200020 Mon Sep 17 00:00:00 2001 From: lgbo-ustc Date: Tue, 7 Jul 2026 11:04:53 +0800 Subject: [PATCH 2/9] test: Add unit tests for GlutenSourceFunction WatermarkStatus handling MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds GlutenSourceFunctionWatermarkStatusTest with 5 test cases covering: - IDLE status triggers markAsTemporarilyIdle() - ACTIVE status is a no-op on SourceContext - IDLE→ACTIVE transition does not call markAsTemporarilyIdle() - Repeated IDLE calls are idempotent - ACTIVE status produces no invocations Uses reflection to invoke private processWatermarkStatus() and a custom TrackingSourceContext spy — no Mockito or native session required. --- ...utenSourceFunctionWatermarkStatusTest.java | 202 ++++++++++++++++++ 1 file changed, 202 insertions(+) create mode 100644 gluten-flink/ut/src/test/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunctionWatermarkStatusTest.java diff --git a/gluten-flink/ut/src/test/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunctionWatermarkStatusTest.java b/gluten-flink/ut/src/test/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunctionWatermarkStatusTest.java new file mode 100644 index 0000000000..b9fe898fb6 --- /dev/null +++ b/gluten-flink/ut/src/test/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunctionWatermarkStatusTest.java @@ -0,0 +1,202 @@ +/* + * 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 io.github.zhztheplayer.velox4j.connector.ConnectorSplit; +import io.github.zhztheplayer.velox4j.connector.FromElementsConnectorSplit; +import io.github.zhztheplayer.velox4j.plan.EmptyNode; +import io.github.zhztheplayer.velox4j.plan.StatefulPlanNode; +import io.github.zhztheplayer.velox4j.stateful.StatefulWatermarkStatus; +import io.github.zhztheplayer.velox4j.type.BigIntType; +import io.github.zhztheplayer.velox4j.type.RowType; + +import org.apache.flink.streaming.api.functions.source.SourceFunction.SourceContext; +import org.apache.flink.streaming.api.watermark.Watermark; +import org.apache.flink.table.data.RowData; + +import org.junit.jupiter.api.Test; + +import java.lang.reflect.Method; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** Unit tests for {@link GlutenSourceFunction} watermark status (IDLE/ACTIVE) handling. */ +public class GlutenSourceFunctionWatermarkStatusTest { + + private static final String NODE_ID = "test-node"; + private static final String CONNECTOR_ID = "test-connector"; + + private static final RowType OUTPUT_TYPE = + new RowType( + Arrays.asList("id"), + Arrays.asList(new BigIntType())); + + /** + * A test spy that tracks {@link SourceContext} invocations so we can verify that idle/active + * transitions happen correctly. + */ + private static class TrackingSourceContext implements SourceContext { + final List invocations = new ArrayList<>(); + final List collected = new ArrayList<>(); + final List watermarks = new ArrayList<>(); + boolean idleCalled = false; + + @Override + public void collect(T element) { + invocations.add("collect"); + collected.add(element); + } + + @Override + public void collectWithTimestamp(T element, long timestamp) { + invocations.add("collectWithTimestamp"); + collected.add(element); + } + + @Override + public void emitWatermark(Watermark mark) { + invocations.add("emitWatermark"); + watermarks.add(mark); + } + + @Override + public void markAsTemporarilyIdle() { + invocations.add("markAsTemporarilyIdle"); + idleCalled = true; + } + + @Override + public Object getCheckpointLock() { + return this; + } + + @Override + public void close() {} + } + + // ─── createMinimalSourceFunction ──────────────────────────────────────── + + private GlutenSourceFunction createSourceFunction() { + ConnectorSplit split = new FromElementsConnectorSplit(CONNECTOR_ID, 0, false); + return new GlutenSourceFunction<>( + new StatefulPlanNode(NODE_ID, new EmptyNode(OUTPUT_TYPE)), + Map.of(NODE_ID, OUTPUT_TYPE), + NODE_ID, + split, + RowData.class); + } + + // ─── tests ────────────────────────────────────────────────────────────── + + @Test + public void testIdleWatermarkStatusCallsMarkAsTemporarilyIdle() throws Exception { + GlutenSourceFunction sourceFunction = createSourceFunction(); + TrackingSourceContext context = new TrackingSourceContext<>(); + StatefulWatermarkStatus idleStatus = new StatefulWatermarkStatus(NODE_ID, true); + + invokeProcessWatermarkStatus(sourceFunction, context, idleStatus); + + assertTrue(context.idleCalled, "IDLE status should trigger markAsTemporarilyIdle()"); + } + + @Test + public void testActiveWatermarkStatusDoesNotCallMarkAsTemporarilyIdle() throws Exception { + GlutenSourceFunction sourceFunction = createSourceFunction(); + TrackingSourceContext context = new TrackingSourceContext<>(); + StatefulWatermarkStatus activeStatus = new StatefulWatermarkStatus(NODE_ID, false); + + invokeProcessWatermarkStatus(sourceFunction, context, activeStatus); + + assertFalse(context.idleCalled, "ACTIVE status should NOT trigger markAsTemporarilyIdle()"); + } + + @Test + public void testIdleThenActiveTransition() throws Exception { + // Verifies that after receiving IDLE, subsequent ACTIVE does not call markAsTemporarilyIdle. + // The ACTIVE status is emitted by native when data resumes; Flink reactivates the source + // when the next record/watermark is collected. + GlutenSourceFunction sourceFunction = createSourceFunction(); + TrackingSourceContext context = new TrackingSourceContext<>(); + + StatefulWatermarkStatus idleStatus = new StatefulWatermarkStatus(NODE_ID, true); + StatefulWatermarkStatus activeStatus = new StatefulWatermarkStatus(NODE_ID, false); + + invokeProcessWatermarkStatus(sourceFunction, context, idleStatus); + assertTrue(context.idleCalled); + + context.idleCalled = false; + invokeProcessWatermarkStatus(sourceFunction, context, activeStatus); + assertFalse(context.idleCalled); + + // SourceContext state is not visible from our spy — verify invocations sequence. + assertThat(context.invocations) + .containsExactly("markAsTemporarilyIdle"); + } + + @Test + public void testRepeatedIdleCallsAreIdempotent() throws Exception { + // Multiple IDLE signals should each call markAsTemporarilyIdle — it is safe to + // call it multiple times (Flink's implementation is idempotent). + GlutenSourceFunction sourceFunction = createSourceFunction(); + TrackingSourceContext context = new TrackingSourceContext<>(); + + StatefulWatermarkStatus idleStatus = new StatefulWatermarkStatus(NODE_ID, true); + + invokeProcessWatermarkStatus(sourceFunction, context, idleStatus); + invokeProcessWatermarkStatus(sourceFunction, context, idleStatus); + invokeProcessWatermarkStatus(sourceFunction, context, idleStatus); + + assertThat(context.invocations) + .containsExactly( + "markAsTemporarilyIdle", + "markAsTemporarilyIdle", + "markAsTemporarilyIdle"); + } + + @Test + public void testActiveWatermarkStatusIsNoOp() throws Exception { + // ACTIVE status should produce no invocations on the SourceContext. + GlutenSourceFunction sourceFunction = createSourceFunction(); + TrackingSourceContext context = new TrackingSourceContext<>(); + + StatefulWatermarkStatus activeStatus = new StatefulWatermarkStatus(NODE_ID, false); + + invokeProcessWatermarkStatus(sourceFunction, context, activeStatus); + + assertThat(context.invocations).isEmpty(); + } + + // ─── reflection helper ────────────────────────────────────────────────── + + private void invokeProcessWatermarkStatus( + GlutenSourceFunction sourceFunction, + SourceContext context, + StatefulWatermarkStatus status) + throws Exception { + Method method = + GlutenSourceFunction.class.getDeclaredMethod( + "processWatermarkStatus", SourceContext.class, StatefulWatermarkStatus.class); + method.setAccessible(true); + method.invoke(sourceFunction, context, status); + } +} From 53d7c084176546a79659e8fa78777a1a1d886ed6 Mon Sep 17 00:00:00 2001 From: lgbo-ustc Date: Tue, 7 Jul 2026 04:18:37 +0000 Subject: [PATCH 3/9] test: Add integration test for idle watermark status in WatermarkAssigner --- ...utenOneInputWatermarkAssignerIdleTest.java | 241 ++++++++++++++++++ ...utenSourceFunctionWatermarkStatusTest.java | 12 +- 2 files changed, 244 insertions(+), 9 deletions(-) create mode 100644 gluten-flink/ut/src/test/java/org/apache/gluten/streaming/api/operators/GlutenOneInputWatermarkAssignerIdleTest.java diff --git a/gluten-flink/ut/src/test/java/org/apache/gluten/streaming/api/operators/GlutenOneInputWatermarkAssignerIdleTest.java b/gluten-flink/ut/src/test/java/org/apache/gluten/streaming/api/operators/GlutenOneInputWatermarkAssignerIdleTest.java new file mode 100644 index 0000000000..00f28d52c8 --- /dev/null +++ b/gluten-flink/ut/src/test/java/org/apache/gluten/streaming/api/operators/GlutenOneInputWatermarkAssignerIdleTest.java @@ -0,0 +1,241 @@ +/* + * 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.streaming.api.operators; + +import org.apache.gluten.rexnode.RexConversionContext; +import org.apache.gluten.rexnode.RexNodeConverter; +import org.apache.gluten.rexnode.Utils; +import org.apache.gluten.table.runtime.operators.GlutenOneInputOperator; +import org.apache.gluten.table.runtime.stream.common.Velox4jEnvironment; +import org.apache.gluten.util.LogicalTypeConverter; +import org.apache.gluten.util.PlanNodeIdGenerator; + +import io.github.zhztheplayer.velox4j.expression.TypedExpr; +import io.github.zhztheplayer.velox4j.plan.EmptyNode; +import io.github.zhztheplayer.velox4j.plan.ProjectNode; +import io.github.zhztheplayer.velox4j.plan.StatefulPlanNode; +import io.github.zhztheplayer.velox4j.plan.WatermarkAssignerNode; + +import org.apache.flink.api.common.serialization.SerializerConfigImpl; +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.streaming.api.watermark.Watermark; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.streaming.runtime.watermarkstatus.WatermarkStatus; +import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness; +import org.apache.flink.table.data.GenericRowData; +import org.apache.flink.table.data.RowData; +import org.apache.flink.table.planner.calcite.FlinkRexBuilder; +import org.apache.flink.table.planner.calcite.FlinkTypeFactory; +import org.apache.flink.table.planner.calcite.FlinkTypeSystem; +import org.apache.flink.table.runtime.typeutils.InternalTypeInfo; +import org.apache.flink.table.types.logical.BigIntType; +import org.apache.flink.table.types.logical.IntType; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.RowType; + +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.type.SqlTypeName; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +import java.util.List; +import java.util.Map; +import java.util.Queue; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Integration test for WatermarkAssigner idle detection. + * + *

Native {@code checkWatermarkStatus} is driven by the {@code next()} → {@code advance()} loop + * inside the {@code GlutenOneInputOperator}'s drain pipeline. Since {@code addInput()} resets the + * idle baseline on every record, the idle check must happen on an {@code advance()} call that + * processes no new input. We achieve this by calling {@code processWatermark()} (which triggers a + * drain cycle without new data) after waiting past the idle timeout. + */ +public class GlutenOneInputWatermarkAssignerIdleTest { + + private static FlinkTypeFactory typeFactory; + private static FlinkRexBuilder rexBuilder; + private static RowType inputFlinkRowType; + private static io.github.zhztheplayer.velox4j.type.RowType inputVeloxType; + private static TypeInformation typeInfo; + + @BeforeAll + static void setUpClass() { + Velox4jEnvironment.initializeOnce(); + + typeFactory = + new FlinkTypeFactory( + Thread.currentThread().getContextClassLoader(), FlinkTypeSystem.INSTANCE); + rexBuilder = new FlinkRexBuilder(typeFactory); + + inputFlinkRowType = + RowType.of(new LogicalType[] {new IntType(), new BigIntType()}, new String[] {"id", "ts"}); + inputVeloxType = + (io.github.zhztheplayer.velox4j.type.RowType) + LogicalTypeConverter.toVLType(inputFlinkRowType); + typeInfo = InternalTypeInfo.of(inputFlinkRowType); + } + + @Test + void testIdleDetectionWithRealTimePassage() throws Exception { + long idleTimeout = 100L; // 100 ms + long watermarkInterval = 50L; + + // ── Watermark expression: reference the ts field (index 1) ── + List fieldNames = Utils.getNamesFromRowType(inputFlinkRowType); + RexNode tsRef = rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.BIGINT), 1); + TypedExpr watermarkExpr = + RexNodeConverter.toTypedExpr(tsRef, new RexConversionContext(fieldNames)); + + ProjectNode watermarkProject = + new ProjectNode( + PlanNodeIdGenerator.newId(), + List.of(new EmptyNode(inputVeloxType)), + List.of("TIMESTAMP"), + List.of(watermarkExpr)); + + // ── WatermarkAssignerNode ── + WatermarkAssignerNode assignerNode = + new WatermarkAssignerNode( + PlanNodeIdGenerator.newId(), + null, + watermarkProject, + idleTimeout, + 1, // rowtimeFieldIndex (ts at index 1) + watermarkInterval); + + // ── GlutenOneInputOperator ── + GlutenOneInputOperator operator = + new GlutenOneInputOperator<>( + new StatefulPlanNode(assignerNode.getId(), assignerNode), + PlanNodeIdGenerator.newId(), + inputVeloxType, + Map.of(assignerNode.getId(), inputVeloxType), + RowData.class, + RowData.class, + "IdleDetectionTest"); + + // ── Test harness ── + OneInputStreamOperatorTestHarness harness = + new OneInputStreamOperatorTestHarness<>( + operator, typeInfo.createSerializer(new SerializerConfigImpl())); + harness.setup(typeInfo.createSerializer(new SerializerConfigImpl())); + harness.open(); + + try { + // ── Phase 1: feed one record ── + GenericRowData record1 = GenericRowData.of(1, 1000L); + harness.processElement(new StreamRecord<>(record1, 1000L)); + // output: StreamRecord(record1), Watermark(1000) + // After drain, checkWatermarkStatus(now1) scheduled timer at now1+100ms. + + // ── Phase 2: wait past idleTimeout, then trigger a drain WITHOUT new input ── + Thread.sleep(idleTimeout * 2); // 200 ms > 100 ms + harness.processWatermark(new Watermark(0)); + // Inside drain: advance → next → advanceWithFuture → blocked + // → checkWatermarkStatus(now2) → idle detected (now2 - lastRecordTime > 100ms) + // → push WatermarkStatus.IDLE to pendings_ + + // ── Phase 3: feed a second record — the drain first pops pending IDLE ── + GenericRowData record2 = GenericRowData.of(2, 2000L); + harness.processElement(new StreamRecord<>(record2, 2000L)); + // drain: advance → next → pendings_ non-empty (IDLE) → pop IDLE → emit IDLE + // → advance → next → process record2 → addInput → onRecord → idle was true + // → emit ACTIVE → push ACTIVE → advance → push record2 + watermark + // → pop ACTIVE → emit ACTIVE → pop record2 → collect → pop watermark → emit + + // ── Assertions ── + Queue output = harness.getOutput(); + assertThat(output) + .as("Output must contain WatermarkStatus.IDLE after idle timeout") + .anyMatch(e -> e instanceof WatermarkStatus && ((WatermarkStatus) e).isIdle()); + assertThat(output) + .as("Output must contain WatermarkStatus.ACTIVE after idle→active transition") + .anyMatch(e -> e instanceof WatermarkStatus && !((WatermarkStatus) e).isIdle()); + assertThat(output) + .as("All input records must be preserved") + .anyMatch( + e -> + e instanceof StreamRecord + && ((StreamRecord) e).getValue().getInt(0) == 1) + .anyMatch( + e -> + e instanceof StreamRecord + && ((StreamRecord) e).getValue().getInt(0) == 2); + } finally { + harness.close(); + } + } + + @Test + void testNoIdleWithContinuousRecords() throws Exception { + long idleTimeout = 100L; + long watermarkInterval = 50L; + + List fieldNames = Utils.getNamesFromRowType(inputFlinkRowType); + RexNode tsRef = rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.BIGINT), 1); + TypedExpr watermarkExpr = + RexNodeConverter.toTypedExpr(tsRef, new RexConversionContext(fieldNames)); + + ProjectNode watermarkProject = + new ProjectNode( + PlanNodeIdGenerator.newId(), + List.of(new EmptyNode(inputVeloxType)), + List.of("TIMESTAMP"), + List.of(watermarkExpr)); + + WatermarkAssignerNode assignerNode = + new WatermarkAssignerNode( + PlanNodeIdGenerator.newId(), null, watermarkProject, idleTimeout, 1, watermarkInterval); + + GlutenOneInputOperator operator = + new GlutenOneInputOperator<>( + new StatefulPlanNode(assignerNode.getId(), assignerNode), + PlanNodeIdGenerator.newId(), + inputVeloxType, + Map.of(assignerNode.getId(), inputVeloxType), + RowData.class, + RowData.class, + "NoIdleTest"); + + OneInputStreamOperatorTestHarness harness = + new OneInputStreamOperatorTestHarness<>( + operator, typeInfo.createSerializer(new SerializerConfigImpl())); + harness.setup(typeInfo.createSerializer(new SerializerConfigImpl())); + harness.open(); + + try { + GenericRowData record1 = GenericRowData.of(1, 1000L); + GenericRowData record2 = GenericRowData.of(2, 2000L); + + harness.processElement(new StreamRecord<>(record1, 1000L)); + // Feed second record immediately (no idle gap) → addInput resets baseline + harness.processElement(new StreamRecord<>(record2, 2000L)); + // Then drain without new data — idle timeout has NOT elapsed since record2 + harness.processWatermark(new Watermark(0)); + + Queue output = harness.getOutput(); + assertThat(output) + .as("No WatermarkStatus.IDLE should appear when records arrive continuously") + .noneMatch(e -> e instanceof WatermarkStatus && ((WatermarkStatus) e).isIdle()); + } finally { + harness.close(); + } + } +} diff --git a/gluten-flink/ut/src/test/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunctionWatermarkStatusTest.java b/gluten-flink/ut/src/test/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunctionWatermarkStatusTest.java index b9fe898fb6..c53d38c636 100644 --- a/gluten-flink/ut/src/test/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunctionWatermarkStatusTest.java +++ b/gluten-flink/ut/src/test/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunctionWatermarkStatusTest.java @@ -47,9 +47,7 @@ public class GlutenSourceFunctionWatermarkStatusTest { private static final String CONNECTOR_ID = "test-connector"; private static final RowType OUTPUT_TYPE = - new RowType( - Arrays.asList("id"), - Arrays.asList(new BigIntType())); + new RowType(Arrays.asList("id"), Arrays.asList(new BigIntType())); /** * A test spy that tracks {@link SourceContext} invocations so we can verify that idle/active @@ -149,8 +147,7 @@ public void testIdleThenActiveTransition() throws Exception { assertFalse(context.idleCalled); // SourceContext state is not visible from our spy — verify invocations sequence. - assertThat(context.invocations) - .containsExactly("markAsTemporarilyIdle"); + assertThat(context.invocations).containsExactly("markAsTemporarilyIdle"); } @Test @@ -167,10 +164,7 @@ public void testRepeatedIdleCallsAreIdempotent() throws Exception { invokeProcessWatermarkStatus(sourceFunction, context, idleStatus); assertThat(context.invocations) - .containsExactly( - "markAsTemporarilyIdle", - "markAsTemporarilyIdle", - "markAsTemporarilyIdle"); + .containsExactly("markAsTemporarilyIdle", "markAsTemporarilyIdle", "markAsTemporarilyIdle"); } @Test From 15c246dcca59f5284fbcdc34af692e6d602bd599 Mon Sep 17 00:00:00 2001 From: lgbo-ustc Date: Tue, 7 Jul 2026 04:21:28 +0000 Subject: [PATCH 4/9] fix: Update velox4j reference to feature/idle-source-handling branch --- .github/workflows/flink.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/flink.yml b/.github/workflows/flink.yml index e32cc1b7fc..d066e6859b 100644 --- a/.github/workflows/flink.yml +++ b/.github/workflows/flink.yml @@ -87,8 +87,8 @@ jobs: 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 3ccc4e61dec4d569dc3dd9408b78ee1d4facf3ba + git clone -b feature/idle-source-handling https://github.com/bigo-sg/velox4j.git + cd velox4j && git reset --hard 9d2ae596a6b31c2f7d81f05d3dcba1ff953c6fab git apply $GITHUB_WORKSPACE/gluten-flink/patches/fix-velox4j.patch $GITHUB_WORKSPACE/build/mvn clean install -DskipTests -Dgpg.skip -Dspotless.skip=true cd .. From ec01991b7a19dcf5a89a99ecbf353dbe83aea1ea Mon Sep 17 00:00:00 2001 From: lgbo-ustc Date: Tue, 7 Jul 2026 14:39:54 +0800 Subject: [PATCH 5/9] fix: Add shouldCallNoMoreSplits option to GlutenSourceFunction for unbounded test scenarios --- .../operators/GlutenSourceFunction.java | 16 ++++++- ...nSourceFunctionWatermarkStatusE2ETest.java | 43 +++++++++++++++++++ 2 files changed, 58 insertions(+), 1 deletion(-) create mode 100644 gluten-flink/ut/src/test/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunctionWatermarkStatusE2ETest.java 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 e7dfab6f09..e2d4b787e8 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 @@ -65,6 +65,13 @@ public class GlutenSourceFunction extends RichParallelSourceFunction 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; @@ -86,6 +93,11 @@ public GlutenSourceFunction( 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; } @@ -283,7 +295,9 @@ private void initSession() { VeloxConnectorConfig.getConfig(getRuntimeContext())); task = session.queryOps().execute(query); task.addSplit(id, activeSplit); - task.noMoreSplits(id); + if (shouldCallNoMoreSplits) { + task.noMoreSplits(id); + } metrics = new SourceOperatorMetrics(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(); + } +} From bf1aa13c9a64d0f330f383ba224089e019a8bab6 Mon Sep 17 00:00:00 2001 From: lgbo-ustc Date: Tue, 7 Jul 2026 16:31:07 +0800 Subject: [PATCH 6/9] feat: E2E test for idle WatermarkStatus detection with Kafka + MiniCluster - Add GlutenStreamSource.isShouldCallNoMoreSplits() delegating to source - Add GlutenSourceFunction.isShouldCallNoMoreSplits() getter - Extend OffloadedJobGraphGenerator to preserve shouldCallNoMoreSplits when creating a new GlutenSourceFunction during offloading - Rewrite GlutenSourceFunctionWatermarkStatusE2ETest as a real E2E test using embedded Kafka broker + Flink MiniCluster, verifying that WatermarkStatus.IDLE is emitted after idle timeout - Fix EmptyNode output type in WatermarkPushDownSpec project to match table scan schema (avoids FieldNotFound error during plan init) --- .../client/OffloadedJobGraphGenerator.java | 18 +- .../api/operators/GlutenStreamSource.java | 4 + .../operators/GlutenSourceFunction.java | 5 + ...nSourceFunctionWatermarkStatusE2ETest.java | 207 +++++++++++++++++- 4 files changed, 219 insertions(+), 15 deletions(-) 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 a9fa29557d..e24fb4c069 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 @@ -187,14 +187,16 @@ private OperatorChainSlice createOffloadedOperatorChainSlice( boolean supportsVectorOutput = supportsVectorOutput(sourceChainSlice, chainSliceGraph, jobVertex); Class outClass = supportsVectorOutput ? StatefulRecord.class : RowData.class; - GlutenStreamSource newSourceOp = - new GlutenStreamSource( - new GlutenSourceFunction<>( - planNode, - sourceOperator.getOutputTypes(), - sourceOperator.getId(), - ((GlutenStreamSource) sourceOperator).getConnectorSplit(), - outClass)); + GlutenSourceFunction newFn = + new GlutenSourceFunction<>( + planNode, + sourceOperator.getOutputTypes(), + sourceOperator.getId(), + ((GlutenStreamSource) sourceOperator).getConnectorSplit(), + outClass); + newFn.setShouldCallNoMoreSplits( + ((GlutenStreamSource) sourceOperator).isShouldCallNoMoreSplits()); + GlutenStreamSource newSourceOp = new GlutenStreamSource(newFn); offloadedOpConfig.setStreamOperator(newSourceOp); if (supportsVectorOutput) { setOffloadedOutputSerializer(offloadedOpConfig, sourceOperator); diff --git a/gluten-flink/runtime/src/main/java/org/apache/gluten/streaming/api/operators/GlutenStreamSource.java b/gluten-flink/runtime/src/main/java/org/apache/gluten/streaming/api/operators/GlutenStreamSource.java index acc6b2e60a..dae9f84410 100644 --- a/gluten-flink/runtime/src/main/java/org/apache/gluten/streaming/api/operators/GlutenStreamSource.java +++ b/gluten-flink/runtime/src/main/java/org/apache/gluten/streaming/api/operators/GlutenStreamSource.java @@ -79,6 +79,10 @@ public ConnectorSplit getConnectorSplit() { return sourceFunction.getConnectorSplit(); } + public boolean isShouldCallNoMoreSplits() { + return sourceFunction.isShouldCallNoMoreSplits(); + } + @SuppressWarnings("rawtypes") private SourceFunction.SourceContext getSourceContext() { return (SourceFunction.SourceContext) 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 e2d4b787e8..c8c4b2fda8 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 @@ -98,6 +98,11 @@ public void setShouldCallNoMoreSplits(boolean value) { this.shouldCallNoMoreSplits = value; } + /** Returns whether noMoreSplits() should be called after adding the initial split. */ + public boolean isShouldCallNoMoreSplits() { + return shouldCallNoMoreSplits; + } + public StatefulPlanNode getPlanNode() { return planNode; } 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 index 9ae81ed71e..d5303eda29 100644 --- 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 @@ -16,28 +16,221 @@ */ package org.apache.gluten.table.runtime.operators; +import org.apache.gluten.streaming.api.operators.GlutenStreamSource; import org.apache.gluten.table.runtime.stream.common.Velox4jEnvironment; +import io.github.zhztheplayer.velox4j.connector.KafkaConnectorSplit; +import io.github.zhztheplayer.velox4j.connector.KafkaTableHandle; +import io.github.zhztheplayer.velox4j.expression.FieldAccessTypedExpr; +import io.github.zhztheplayer.velox4j.plan.EmptyNode; +import io.github.zhztheplayer.velox4j.plan.ProjectNode; +import io.github.zhztheplayer.velox4j.plan.StatefulPlanNode; +import io.github.zhztheplayer.velox4j.plan.TableScanWithWatermarkNode; +import io.github.zhztheplayer.velox4j.plan.WatermarkPushDownSpec; +import io.github.zhztheplayer.velox4j.stateful.StatefulRecord; +import io.github.zhztheplayer.velox4j.type.BigIntType; +import io.github.zhztheplayer.velox4j.type.RowType; +import io.github.zhztheplayer.velox4j.type.VarCharType; + +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.streaming.api.datastream.DataStreamSource; +import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; +import org.apache.flink.streaming.api.operators.AbstractStreamOperator; +import org.apache.flink.streaming.api.operators.OneInputStreamOperator; +import org.apache.flink.streaming.api.watermark.Watermark; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.streaming.runtime.watermarkstatus.WatermarkStatus; + +import com.salesforce.kafka.test.junit5.SharedKafkaTestResource; +import com.salesforce.kafka.test.listeners.PlainListener; +import org.apache.kafka.clients.producer.ProducerRecord; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; -import java.lang.reflect.Method; +import java.nio.charset.StandardCharsets; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.CopyOnWriteArrayList; import static org.assertj.core.api.Assertions.assertThat; -/** Tests that GlutenSourceFunction correctly handles WatermarkStatus elements. */ +/** + * End-to-end test that verifies WatermarkStatus.IDLE is emitted through GlutenSourceFunction when + * the native Kafka source detects idleness. + * + *

The test produces a few records to an embedded Kafka broker, starts a Flink MiniCluster job + * with the Gluten native source pipeline, waits for the idle timeout to expire, and checks that + * WatermarkStatus.IDLE is captured by a downstream operator. + */ class GlutenSourceFunctionWatermarkStatusE2ETest { + private static final int KAFKA_PORT = 19093; + private static final long IDLE_TIMEOUT_MS = 5000; + private static final long WATERMARK_INTERVAL_MS = 500; + + @RegisterExtension + static final SharedKafkaTestResource KAFKA = + new SharedKafkaTestResource() + .withBrokerProperty("host.name", "127.0.0.1") + .withBrokers(1) + .registerListener(new PlainListener().onPorts(KAFKA_PORT)); + + private static final CopyOnWriteArrayList capturedStatuses = + new CopyOnWriteArrayList<>(); + @BeforeAll static void setupGluten() { Velox4jEnvironment.initializeOnce(); } + @BeforeEach + void clearCaptured() { + capturedStatuses.clear(); + } + + @AfterEach + void ensureJobCancelled() { + if (jobThread != null && jobThread.isAlive()) { + jobThread.interrupt(); + } + } + + private Thread jobThread; + @Test - void testGlutenSourceFunctionClassHasWatermarkStatusHandling() throws Exception { - // Verify GlutenSourceFunction has the setShouldCallNoMoreSplits method - Method setter = - GlutenSourceFunction.class.getMethod("setShouldCallNoMoreSplits", boolean.class); - assertThat(setter).isNotNull(); + void testIdleDetectionAfterStopWritingToKafka() throws Exception { + String topic = "idle_e2e_" + UUID.randomUUID().toString().replace("-", ""); + KAFKA.getKafkaTestUtils().createTopic(topic, 1, (short) 1); + KAFKA + .getKafkaTestUtils() + .produceRecords( + List.of( + jsonRecord(topic, "{\"id\":1000,\"name\":\"r0\"}"), + jsonRecord(topic, "{\"id\":2000,\"name\":\"r1\"}"), + jsonRecord(topic, "{\"id\":3000,\"name\":\"r2\"}"))); + + // Build pipeline: GlutenStreamSource -> WatermarkStatusCaptureOperator + StreamExecutionEnvironment env = StreamExecutionEnvironment.createLocalEnvironment(1); + env.getConfig().disableClosureCleaner(); + env.setParallelism(1); + + DataStreamSource source = addSourceToEnv(env, topic); + source + .transform("capture", TypeInformation.of(StatefulRecord.class), new StatusCaptureOp()) + .setParallelism(1); + + // Run job in background thread + jobThread = + new Thread( + () -> { + try { + env.execute("IdleDetectionE2ETest"); + } catch (Exception e) { + // ignore cancellation + } + }, + "job-runner"); + jobThread.start(); + + // Wait for initial records to be consumed + Thread.sleep(WATERMARK_INTERVAL_MS * 4); + + // Wait for idle timeout + Thread.sleep(IDLE_TIMEOUT_MS + 3000); + + // Cancel the job + jobThread.interrupt(); + + // Verify idle was detected + assertThat(capturedStatuses) + .as("Should have received WatermarkStatus.IDLE after idle timeout") + .contains(WatermarkStatus.IDLE); + } + + // -- helpers -- + + private static DataStreamSource addSourceToEnv( + StreamExecutionEnvironment env, String topic) { + RowType veloxRowType = + new RowType(List.of("id", "name"), List.of(new BigIntType(), new VarCharType())); + + ProjectNode watermarkProject = + new ProjectNode( + "watermark_project", + List.of(new EmptyNode(veloxRowType)), + List.of("watermark"), + List.of(FieldAccessTypedExpr.create(new BigIntType(), "id"))); + WatermarkPushDownSpec watermarkSpec = + new WatermarkPushDownSpec(watermarkProject, IDLE_TIMEOUT_MS, WATERMARK_INTERVAL_MS, 0); + + Map tableParams = new HashMap<>(); + tableParams.put("bootstrap.servers", "127.0.0.1:" + KAFKA_PORT); + tableParams.put("client.id", "test-client-e2e-" + UUID.randomUUID()); + tableParams.put("group.id", "test-group-e2e"); + tableParams.put("topic", topic); + tableParams.put("format", "json"); + tableParams.put("scan.startup.mode", "earliest-offsets"); + tableParams.put("enable.auto.commit", "false"); + + KafkaTableHandle tableHandle = + new KafkaTableHandle("connector-kafka", topic, veloxRowType, tableParams); + String planId = "plan_" + UUID.randomUUID().toString().replace("-", ""); + TableScanWithWatermarkNode scanNode = + new TableScanWithWatermarkNode(planId, veloxRowType, tableHandle, List.of(), watermarkSpec); + KafkaConnectorSplit connectorSplit = + new KafkaConnectorSplit( + "connector-kafka", + 0, + false, + "127.0.0.1:" + KAFKA_PORT, + "test-group-e2e", + "json", + false, + "earliest-offset", + List.of(new KafkaConnectorSplit.TopicPartitionOffset(topic, 0, -1L))); + + GlutenSourceFunction sourceFn = + new GlutenSourceFunction<>( + new StatefulPlanNode(scanNode.getId(), scanNode), + Map.of(scanNode.getId(), veloxRowType), + scanNode.getId(), + connectorSplit, + StatefulRecord.class); + sourceFn.setShouldCallNoMoreSplits(false); + + GlutenStreamSource sourceOp = new GlutenStreamSource(sourceFn, "KafkaSource"); + return new DataStreamSource( + env, TypeInformation.of(StatefulRecord.class), sourceOp, false, "KafkaSource"); + } + + private static ProducerRecord jsonRecord(String topic, String value) { + return new ProducerRecord<>(topic, value.getBytes(StandardCharsets.UTF_8)); + } + + // -- Capture operator -- + + private static class StatusCaptureOp extends AbstractStreamOperator + implements OneInputStreamOperator { + + @Override + public void processElement(StreamRecord element) throws Exception { + output.collect(element); + } + + @Override + public void processWatermark(Watermark mark) throws Exception { + output.emitWatermark(mark); + } + + @Override + public void processWatermarkStatus(WatermarkStatus status) throws Exception { + capturedStatuses.add(status); + super.processWatermarkStatus(status); + } } } From a2ab1ae6222009546383f8e44aed6405fe943225 Mon Sep 17 00:00:00 2001 From: lgbo-ustc Date: Tue, 7 Jul 2026 16:44:58 +0800 Subject: [PATCH 7/9] chore: remove obsolete GlutenSourceFunctionWatermarkStatusTest Covered by GlutenSourceFunctionWatermarkStatusE2ETest which tests the same behavior end-to-end with Kafka + MiniCluster. --- ...utenSourceFunctionWatermarkStatusTest.java | 196 ------------------ 1 file changed, 196 deletions(-) delete mode 100644 gluten-flink/ut/src/test/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunctionWatermarkStatusTest.java diff --git a/gluten-flink/ut/src/test/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunctionWatermarkStatusTest.java b/gluten-flink/ut/src/test/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunctionWatermarkStatusTest.java deleted file mode 100644 index c53d38c636..0000000000 --- a/gluten-flink/ut/src/test/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunctionWatermarkStatusTest.java +++ /dev/null @@ -1,196 +0,0 @@ -/* - * 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 io.github.zhztheplayer.velox4j.connector.ConnectorSplit; -import io.github.zhztheplayer.velox4j.connector.FromElementsConnectorSplit; -import io.github.zhztheplayer.velox4j.plan.EmptyNode; -import io.github.zhztheplayer.velox4j.plan.StatefulPlanNode; -import io.github.zhztheplayer.velox4j.stateful.StatefulWatermarkStatus; -import io.github.zhztheplayer.velox4j.type.BigIntType; -import io.github.zhztheplayer.velox4j.type.RowType; - -import org.apache.flink.streaming.api.functions.source.SourceFunction.SourceContext; -import org.apache.flink.streaming.api.watermark.Watermark; -import org.apache.flink.table.data.RowData; - -import org.junit.jupiter.api.Test; - -import java.lang.reflect.Method; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.List; -import java.util.Map; - -import static org.assertj.core.api.Assertions.assertThat; -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertTrue; - -/** Unit tests for {@link GlutenSourceFunction} watermark status (IDLE/ACTIVE) handling. */ -public class GlutenSourceFunctionWatermarkStatusTest { - - private static final String NODE_ID = "test-node"; - private static final String CONNECTOR_ID = "test-connector"; - - private static final RowType OUTPUT_TYPE = - new RowType(Arrays.asList("id"), Arrays.asList(new BigIntType())); - - /** - * A test spy that tracks {@link SourceContext} invocations so we can verify that idle/active - * transitions happen correctly. - */ - private static class TrackingSourceContext implements SourceContext { - final List invocations = new ArrayList<>(); - final List collected = new ArrayList<>(); - final List watermarks = new ArrayList<>(); - boolean idleCalled = false; - - @Override - public void collect(T element) { - invocations.add("collect"); - collected.add(element); - } - - @Override - public void collectWithTimestamp(T element, long timestamp) { - invocations.add("collectWithTimestamp"); - collected.add(element); - } - - @Override - public void emitWatermark(Watermark mark) { - invocations.add("emitWatermark"); - watermarks.add(mark); - } - - @Override - public void markAsTemporarilyIdle() { - invocations.add("markAsTemporarilyIdle"); - idleCalled = true; - } - - @Override - public Object getCheckpointLock() { - return this; - } - - @Override - public void close() {} - } - - // ─── createMinimalSourceFunction ──────────────────────────────────────── - - private GlutenSourceFunction createSourceFunction() { - ConnectorSplit split = new FromElementsConnectorSplit(CONNECTOR_ID, 0, false); - return new GlutenSourceFunction<>( - new StatefulPlanNode(NODE_ID, new EmptyNode(OUTPUT_TYPE)), - Map.of(NODE_ID, OUTPUT_TYPE), - NODE_ID, - split, - RowData.class); - } - - // ─── tests ────────────────────────────────────────────────────────────── - - @Test - public void testIdleWatermarkStatusCallsMarkAsTemporarilyIdle() throws Exception { - GlutenSourceFunction sourceFunction = createSourceFunction(); - TrackingSourceContext context = new TrackingSourceContext<>(); - StatefulWatermarkStatus idleStatus = new StatefulWatermarkStatus(NODE_ID, true); - - invokeProcessWatermarkStatus(sourceFunction, context, idleStatus); - - assertTrue(context.idleCalled, "IDLE status should trigger markAsTemporarilyIdle()"); - } - - @Test - public void testActiveWatermarkStatusDoesNotCallMarkAsTemporarilyIdle() throws Exception { - GlutenSourceFunction sourceFunction = createSourceFunction(); - TrackingSourceContext context = new TrackingSourceContext<>(); - StatefulWatermarkStatus activeStatus = new StatefulWatermarkStatus(NODE_ID, false); - - invokeProcessWatermarkStatus(sourceFunction, context, activeStatus); - - assertFalse(context.idleCalled, "ACTIVE status should NOT trigger markAsTemporarilyIdle()"); - } - - @Test - public void testIdleThenActiveTransition() throws Exception { - // Verifies that after receiving IDLE, subsequent ACTIVE does not call markAsTemporarilyIdle. - // The ACTIVE status is emitted by native when data resumes; Flink reactivates the source - // when the next record/watermark is collected. - GlutenSourceFunction sourceFunction = createSourceFunction(); - TrackingSourceContext context = new TrackingSourceContext<>(); - - StatefulWatermarkStatus idleStatus = new StatefulWatermarkStatus(NODE_ID, true); - StatefulWatermarkStatus activeStatus = new StatefulWatermarkStatus(NODE_ID, false); - - invokeProcessWatermarkStatus(sourceFunction, context, idleStatus); - assertTrue(context.idleCalled); - - context.idleCalled = false; - invokeProcessWatermarkStatus(sourceFunction, context, activeStatus); - assertFalse(context.idleCalled); - - // SourceContext state is not visible from our spy — verify invocations sequence. - assertThat(context.invocations).containsExactly("markAsTemporarilyIdle"); - } - - @Test - public void testRepeatedIdleCallsAreIdempotent() throws Exception { - // Multiple IDLE signals should each call markAsTemporarilyIdle — it is safe to - // call it multiple times (Flink's implementation is idempotent). - GlutenSourceFunction sourceFunction = createSourceFunction(); - TrackingSourceContext context = new TrackingSourceContext<>(); - - StatefulWatermarkStatus idleStatus = new StatefulWatermarkStatus(NODE_ID, true); - - invokeProcessWatermarkStatus(sourceFunction, context, idleStatus); - invokeProcessWatermarkStatus(sourceFunction, context, idleStatus); - invokeProcessWatermarkStatus(sourceFunction, context, idleStatus); - - assertThat(context.invocations) - .containsExactly("markAsTemporarilyIdle", "markAsTemporarilyIdle", "markAsTemporarilyIdle"); - } - - @Test - public void testActiveWatermarkStatusIsNoOp() throws Exception { - // ACTIVE status should produce no invocations on the SourceContext. - GlutenSourceFunction sourceFunction = createSourceFunction(); - TrackingSourceContext context = new TrackingSourceContext<>(); - - StatefulWatermarkStatus activeStatus = new StatefulWatermarkStatus(NODE_ID, false); - - invokeProcessWatermarkStatus(sourceFunction, context, activeStatus); - - assertThat(context.invocations).isEmpty(); - } - - // ─── reflection helper ────────────────────────────────────────────────── - - private void invokeProcessWatermarkStatus( - GlutenSourceFunction sourceFunction, - SourceContext context, - StatefulWatermarkStatus status) - throws Exception { - Method method = - GlutenSourceFunction.class.getDeclaredMethod( - "processWatermarkStatus", SourceContext.class, StatefulWatermarkStatus.class); - method.setAccessible(true); - method.invoke(sourceFunction, context, status); - } -} From c8d9460f7888b4b8afef7a253cbcddfd54f78886 Mon Sep 17 00:00:00 2001 From: lgbo-ustc Date: Tue, 7 Jul 2026 17:25:49 +0800 Subject: [PATCH 8/9] test: verify idle inputs are excluded from combined min-watermark Add testIdleInputExcludedFromMinWatermark to GlutenStreamTwoInputWatermarkStatusTest: when one input is marked IDLE, its watermark is excluded from min-watermark calculation so the other active input can advance freely. --- ...utenStreamTwoInputWatermarkStatusTest.java | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/gluten-flink/ut/src/test/java/org/apache/gluten/streaming/api/operators/GlutenStreamTwoInputWatermarkStatusTest.java b/gluten-flink/ut/src/test/java/org/apache/gluten/streaming/api/operators/GlutenStreamTwoInputWatermarkStatusTest.java index 8306c9d20d..a35b21716d 100644 --- a/gluten-flink/ut/src/test/java/org/apache/gluten/streaming/api/operators/GlutenStreamTwoInputWatermarkStatusTest.java +++ b/gluten-flink/ut/src/test/java/org/apache/gluten/streaming/api/operators/GlutenStreamTwoInputWatermarkStatusTest.java @@ -79,6 +79,32 @@ public void testWatermarkReactivatesIdleNativeJoinOperator() throws Exception { } } + @Test + public void testIdleInputExcludedFromMinWatermark() throws Exception { + // When one input is idle, its watermark is excluded from the combined min-watermark + // calculation. The other active input's watermark can advance freely. + GlutenTwoInputOperator operator = createGlutenJoinOperator(FlinkJoinType.INNER); + + try (TwoInputStreamOperatorTestHarness harness = + new TwoInputStreamOperatorTestHarness<>(operator)) { + harness.setup(); + harness.open(); + + harness.processWatermark1(new Watermark(100L)); + harness.processWatermark2(new Watermark(90L)); + assertThat(harness.getOutput()).containsExactly(new Watermark(90L)); + + harness.processWatermarkStatus1(WatermarkStatus.IDLE); + // Input 1 (watermark=100) is idle and excluded. Combined = input 2 (90). No change. + assertThat(harness.getOutput()).containsExactly(new Watermark(90L)); + + // Input 2 advances to 120. Since input 1 is idle and excluded, combined = 120. + // If input 1 were still active, combined would be min(100, 120) = 100. + harness.processWatermark2(new Watermark(120L)); + assertThat(harness.getOutput()).containsExactly(new Watermark(90L), new Watermark(120L)); + } + } + @Test public void testWatermarksUseNativeTwoInputMinimum() throws Exception { // While both inputs are active, native execution should combine indexed input watermarks by From 9cd1b66fcbb795ec2e9d781b111c097b555f220e Mon Sep 17 00:00:00 2001 From: lgbo-ustc Date: Wed, 8 Jul 2026 09:30:28 +0800 Subject: [PATCH 9/9] fix: upgrade surefire in gluten-flink-ut from 3.0.0-M5 to 3.3.0 --- gluten-flink/ut/pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gluten-flink/ut/pom.xml b/gluten-flink/ut/pom.xml index c97c3838e1..0f80c36292 100644 --- a/gluten-flink/ut/pom.xml +++ b/gluten-flink/ut/pom.xml @@ -288,7 +288,7 @@ org.apache.maven.plugins maven-surefire-plugin - 3.0.0-M5 + 3.3.0 false