From c5dfedeb12f819fae518e57b7d4e102bb80c6681 Mon Sep 17 00:00:00 2001 From: Minni Mittal Date: Wed, 1 Apr 2026 07:01:45 +0000 Subject: [PATCH] [VL] Fix streaming UTs by copying tags in EnsureLocalSortRequirements When EnsureLocalSortRequirements adds a local SortExec to satisfy required child ordering, the new SortExec node was missing tags from the original child. This caused streaming test failures because StateStoreWriter and other stateful operators rely on tags (e.g., isStatefulOperatorStreamingRestore) to propagate streaming execution metadata. Fix: Call newChild.copyTagsFrom(originalChild) to preserve tags when wrapping a child plan with SortExec. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../extension/columnar/EnsureLocalSortRequirements.scala | 1 + .../org/apache/gluten/utils/velox/VeloxTestSettings.scala | 6 +++--- .../org/apache/gluten/utils/velox/VeloxTestSettings.scala | 6 +++--- 3 files changed, 7 insertions(+), 6 deletions(-) diff --git a/gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/EnsureLocalSortRequirements.scala b/gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/EnsureLocalSortRequirements.scala index e17a8e746035..cbedac924b0f 100644 --- a/gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/EnsureLocalSortRequirements.scala +++ b/gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/EnsureLocalSortRequirements.scala @@ -38,6 +38,7 @@ object EnsureLocalSortRequirements extends Rule[SparkPlan] { requiredOrdering: Seq[SortOrder]): SparkPlan = { // FIXME: HeuristicTransform is costly. Re-applying it may cause performance issues. val newChild = SortExec(requiredOrdering, global = false, child = originalChild) + newChild.copyTagsFrom(originalChild) transform.apply(newChild) } diff --git a/gluten-ut/spark40/src/test/scala/org/apache/gluten/utils/velox/VeloxTestSettings.scala b/gluten-ut/spark40/src/test/scala/org/apache/gluten/utils/velox/VeloxTestSettings.scala index ac547c57cd75..63b1d9da5152 100644 --- a/gluten-ut/spark40/src/test/scala/org/apache/gluten/utils/velox/VeloxTestSettings.scala +++ b/gluten-ut/spark40/src/test/scala/org/apache/gluten/utils/velox/VeloxTestSettings.scala @@ -1200,15 +1200,15 @@ class VeloxTestSettings extends BackendTestSettings { enableSuite[GlutenFileStreamStressSuite] // TODO: 4.x enableSuite[GlutenFlatMapGroupsInPandasWithStateDistributionSuite] // failures with GlutenPlugin enableSuite[GlutenFlatMapGroupsInPandasWithStateSuite] - // TODO: 4.x enableSuite[GlutenFlatMapGroupsWithStateDistributionSuite] - // TODO: 4.x enableSuite[GlutenFlatMapGroupsWithStateSuite] + enableSuite[GlutenFlatMapGroupsWithStateDistributionSuite] + enableSuite[GlutenFlatMapGroupsWithStateSuite] enableSuite[GlutenFlatMapGroupsWithStateWithInitialStateSuite] enableSuite[GlutenGroupStateSuite] enableSuite[GlutenLongOffsetSuite] enableSuite[GlutenMemorySourceStressSuite] // TODO: 4.x enableSuite[GlutenMultiStatefulOperatorsSuite] // 2 failures enableSuite[GlutenReportSinkMetricsSuite] - // TODO: 4.x enableSuite[GlutenRocksDBStateStoreFlatMapGroupsWithStateSuite] + enableSuite[GlutenRocksDBStateStoreFlatMapGroupsWithStateSuite] // TODO: 4.x enableSuite[GlutenRocksDBStateStoreStreamingAggregationSuite] // TODO: 4.x enableSuite[GlutenRocksDBStateStoreStreamingDeduplicationSuite] // TODO: 4.x enableSuite[GlutenStreamSuite] diff --git a/gluten-ut/spark41/src/test/scala/org/apache/gluten/utils/velox/VeloxTestSettings.scala b/gluten-ut/spark41/src/test/scala/org/apache/gluten/utils/velox/VeloxTestSettings.scala index 0f83f0d653cd..23b12ea2a840 100644 --- a/gluten-ut/spark41/src/test/scala/org/apache/gluten/utils/velox/VeloxTestSettings.scala +++ b/gluten-ut/spark41/src/test/scala/org/apache/gluten/utils/velox/VeloxTestSettings.scala @@ -1214,15 +1214,15 @@ class VeloxTestSettings extends BackendTestSettings { enableSuite[GlutenFileStreamStressSuite] // TODO: 4.x enableSuite[GlutenFlatMapGroupsInPandasWithStateDistributionSuite] // failures with GlutenPlugin enableSuite[GlutenFlatMapGroupsInPandasWithStateSuite] - // TODO: 4.x enableSuite[GlutenFlatMapGroupsWithStateDistributionSuite] - // TODO: 4.x enableSuite[GlutenFlatMapGroupsWithStateSuite] + enableSuite[GlutenFlatMapGroupsWithStateDistributionSuite] + enableSuite[GlutenFlatMapGroupsWithStateSuite] enableSuite[GlutenFlatMapGroupsWithStateWithInitialStateSuite] enableSuite[GlutenGroupStateSuite] enableSuite[GlutenLongOffsetSuite] enableSuite[GlutenMemorySourceStressSuite] // TODO: 4.x enableSuite[GlutenMultiStatefulOperatorsSuite] // 2 failures enableSuite[GlutenReportSinkMetricsSuite] - // TODO: 4.x enableSuite[GlutenRocksDBStateStoreFlatMapGroupsWithStateSuite] + enableSuite[GlutenRocksDBStateStoreFlatMapGroupsWithStateSuite] // TODO: 4.x enableSuite[GlutenRocksDBStateStoreStreamingAggregationSuite] // TODO: 4.x enableSuite[GlutenRocksDBStateStoreStreamingDeduplicationSuite] // TODO: 4.x enableSuite[GlutenStreamSuite]