Preserve tags on FileSourceScanExec (#10461)

Signed-off-by: Jason Lowe <[email protected]>
NVIDIA · Feb 22, 2024 · 3f6196d · 3f6196d
1 parent 3d16605
commit 3f6196d
Show file tree

Hide file tree

Showing 7 changed files with 28 additions and 14 deletions.
diff --git a/sql-plugin/src/main/spark311/scala/com/nvidia/spark/rapids/shims/Spark31XShims.scala b/sql-plugin/src/main/spark311/scala/com/nvidia/spark/rapids/shims/Spark31XShims.scala
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2021-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -339,7 +339,9 @@ abstract class Spark31XShims extends Spark31Xuntil33XShims with Logging {
           override def tagPlanForGpu(): Unit = GpuFileSourceScanExec.tagSupport(this)
 
           override def convertToCpu(): SparkPlan = {
-            wrapped.copy(partitionFilters = partitionFilters)
+            val cpu = wrapped.copy(partitionFilters = partitionFilters)
+            cpu.copyTagsFrom(wrapped)
+            cpu
           }
 
           override def convertToGpu(): GpuExec = {

diff --git a/sql-plugin/src/main/spark320/scala/com/nvidia/spark/rapids/shims/BatchScanExecMeta.scala b/sql-plugin/src/main/spark320/scala/com/nvidia/spark/rapids/shims/BatchScanExecMeta.scala
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -85,7 +85,9 @@ class BatchScanExecMeta(p: BatchScanExec,
   }
 
   override def convertToCpu(): SparkPlan = {
-    wrapped.copy(runtimeFilters = runtimeFilters)
+    val cpu = wrapped.copy(runtimeFilters = runtimeFilters)
+    cpu.copyTagsFrom(wrapped)
+    cpu
   }
 
   override def convertToGpu(): GpuExec =

diff --git a/...plugin/src/main/spark320/scala/com/nvidia/spark/rapids/shims/FileSourceScanExecMeta.scala b/...plugin/src/main/spark320/scala/com/nvidia/spark/rapids/shims/FileSourceScanExecMeta.scala
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -82,7 +82,9 @@ class FileSourceScanExecMeta(plan: FileSourceScanExec,
   override def tagPlanForGpu(): Unit = ScanExecShims.tagGpuFileSourceScanExecSupport(this)
 
   override def convertToCpu(): SparkPlan = {
-    wrapped.copy(partitionFilters = partitionFilters)
+    val cpu = wrapped.copy(partitionFilters = partitionFilters)
+    cpu.copyTagsFrom(wrapped)
+    cpu
   }
 
   override def convertToGpu(): GpuExec = {

diff --git a/...ugin/src/main/spark321db/scala/com/nvidia/spark/rapids/shims/FileSourceScanExecMeta.scala b/...ugin/src/main/spark321db/scala/com/nvidia/spark/rapids/shims/FileSourceScanExecMeta.scala
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -126,7 +126,9 @@ class FileSourceScanExecMeta(plan: FileSourceScanExec,
   }
 
   override def convertToCpu(): SparkPlan = {
-    wrapped.copy(partitionFilters = partitionFilters, dataFilters = dataFilters)
+    val cpu = wrapped.copy(partitionFilters = partitionFilters, dataFilters = dataFilters)
+    cpu.copyTagsFrom(wrapped)
+    cpu
   }
 
   override def convertToGpu(): GpuExec = {

diff --git a/sql-plugin/src/main/spark330db/scala/com/nvidia/spark/rapids/shims/BatchScanExecMeta.scala b/sql-plugin/src/main/spark330db/scala/com/nvidia/spark/rapids/shims/BatchScanExecMeta.scala
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023, NVIDIA CORPORATION.
+ * Copyright (c) 2023-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -76,7 +76,9 @@ class BatchScanExecMeta(p: BatchScanExec,
   }
 
   override def convertToCpu(): SparkPlan = {
-    wrapped.copy(runtimeFilters = runtimeFilters)
+    val cpu = wrapped.copy(runtimeFilters = runtimeFilters)
+    cpu.copyTagsFrom(wrapped)
+    cpu
   }
 
   override def convertToGpu(): GpuExec =

diff --git a/sql-plugin/src/main/spark340/scala/com/nvidia/spark/rapids/shims/BatchScanExecMeta.scala b/sql-plugin/src/main/spark340/scala/com/nvidia/spark/rapids/shims/BatchScanExecMeta.scala
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023, NVIDIA CORPORATION.
+ * Copyright (c) 2023-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -75,7 +75,9 @@ class BatchScanExecMeta(p: BatchScanExec,
   }
 
   override def convertToCpu(): SparkPlan = {
-    wrapped.copy(runtimeFilters = runtimeFilters)
+    val cpu = wrapped.copy(runtimeFilters = runtimeFilters)
+    cpu.copyTagsFrom(wrapped)
+    cpu
   }
 
   override def convertToGpu(): GpuExec = {

diff --git a/sql-plugin/src/main/spark350/scala/com/nvidia/spark/rapids/shims/BatchScanExecMeta.scala b/sql-plugin/src/main/spark350/scala/com/nvidia/spark/rapids/shims/BatchScanExecMeta.scala
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023, NVIDIA CORPORATION.
+ * Copyright (c) 2023-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -73,7 +73,9 @@ class BatchScanExecMeta(p: BatchScanExec,
   }
 
   override def convertToCpu(): SparkPlan = {
-    wrapped.copy(runtimeFilters = runtimeFilters)
+    val cpu = wrapped.copy(runtimeFilters = runtimeFilters)
+    cpu.copyTagsFrom(wrapped)
+    cpu
   }
 
   override def convertToGpu(): GpuExec = {