From 3f6196d40fc6fe50d5e8853d7e84f54e9e09ecc6 Mon Sep 17 00:00:00 2001 From: Jason Lowe Date: Thu, 22 Feb 2024 11:47:31 -0600 Subject: [PATCH] Preserve tags on FileSourceScanExec (#10461) Signed-off-by: Jason Lowe --- .../scala/com/nvidia/spark/rapids/shims/Spark31XShims.scala | 6 ++++-- .../com/nvidia/spark/rapids/shims/BatchScanExecMeta.scala | 6 ++++-- .../nvidia/spark/rapids/shims/FileSourceScanExecMeta.scala | 6 ++++-- .../nvidia/spark/rapids/shims/FileSourceScanExecMeta.scala | 6 ++++-- .../com/nvidia/spark/rapids/shims/BatchScanExecMeta.scala | 6 ++++-- .../com/nvidia/spark/rapids/shims/BatchScanExecMeta.scala | 6 ++++-- .../com/nvidia/spark/rapids/shims/BatchScanExecMeta.scala | 6 ++++-- 7 files changed, 28 insertions(+), 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 index c97c63c2e3c..eab56a8348a 100644 --- 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 index 10781b0142e..af8d7ef7d95 100644 --- 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/sql-plugin/src/main/spark320/scala/com/nvidia/spark/rapids/shims/FileSourceScanExecMeta.scala b/sql-plugin/src/main/spark320/scala/com/nvidia/spark/rapids/shims/FileSourceScanExecMeta.scala index ff6c408f7d8..a318bef159d 100644 --- a/sql-plugin/src/main/spark320/scala/com/nvidia/spark/rapids/shims/FileSourceScanExecMeta.scala +++ b/sql-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/sql-plugin/src/main/spark321db/scala/com/nvidia/spark/rapids/shims/FileSourceScanExecMeta.scala b/sql-plugin/src/main/spark321db/scala/com/nvidia/spark/rapids/shims/FileSourceScanExecMeta.scala index bcfcc59d2a7..8f61c594715 100644 --- a/sql-plugin/src/main/spark321db/scala/com/nvidia/spark/rapids/shims/FileSourceScanExecMeta.scala +++ b/sql-plugin/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 index 34a45ca1920..d3e9fe903a3 100644 --- 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 index 7e6fd7d67d4..0dabfd1e46d 100644 --- 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 index e08765dda97..b9999dee246 100644 --- 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 = {