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]