[GLUTEN-6151] Reset local property after finishing write operator (ap…

…ache#6163) quick bug fix. Need to reset local property on fallback
ccat3z · Jun 21, 2024 · 1e06169 · 1e06169
1 parent ef7b2c5
commit 1e06169
Showing 1 changed file with 21 additions and 12 deletions.
diff --git a/...src/main/scala/org/apache/spark/sql/execution/datasources/GlutenWriterColumnarRules.scala b/...src/main/scala/org/apache/spark/sql/execution/datasources/GlutenWriterColumnarRules.scala
@@ -162,19 +162,28 @@ object GlutenWriterColumnarRules {
           if write.getClass.getName == NOOP_WRITE &&
             BackendsApiManager.getSettings.enableNativeWriteFiles() =>
         injectFakeRowAdaptor(rc, rc.child)
-      case rc @ DataWritingCommandExec(cmd, child)
-          if BackendsApiManager.getSettings.supportNativeWrite(child.output.toStructType.fields) =>
-        val format = getNativeFormat(cmd)
-        session.sparkContext.setLocalProperty(
-          "staticPartitionWriteOnly",
-          BackendsApiManager.getSettings.staticPartitionWriteOnly().toString)
-        // FIXME: We should only use context property if having no other approaches.
-        //  Should see if there is another way to pass these options.
-        session.sparkContext.setLocalProperty("isNativeAppliable", format.isDefined.toString)
-        session.sparkContext.setLocalProperty("nativeFormat", format.getOrElse(""))
-        if (format.isDefined) {
-          injectFakeRowAdaptor(rc, child)
+      case rc @ DataWritingCommandExec(cmd, child) =>
+        if (BackendsApiManager.getSettings.supportNativeWrite(child.output.toStructType.fields)) {
+          val format = getNativeFormat(cmd)
+          session.sparkContext.setLocalProperty(
+            "staticPartitionWriteOnly",
+            BackendsApiManager.getSettings.staticPartitionWriteOnly().toString)
+          // FIXME: We should only use context property if having no other approaches.
+          //  Should see if there is another way to pass these options.
+          session.sparkContext.setLocalProperty("isNativeAppliable", format.isDefined.toString)
+          session.sparkContext.setLocalProperty("nativeFormat", format.getOrElse(""))
+          if (format.isDefined) {
+            injectFakeRowAdaptor(rc, child)
+          } else {
+            rc.withNewChildren(rc.children.map(apply))
+          }
         } else {
+          session.sparkContext.setLocalProperty(
+            "staticPartitionWriteOnly",
+            BackendsApiManager.getSettings.staticPartitionWriteOnly().toString)
+          session.sparkContext.setLocalProperty("isNativeAppliable", "false")
+          session.sparkContext.setLocalProperty("nativeFormat", "")
+
           rc.withNewChildren(rc.children.map(apply))
         }
       case plan: SparkPlan => plan.withNewChildren(plan.children.map(apply))