Skip to content

Commit ddaea75

Browse files
sdrp713sdrp713
committed
Init Commit
Signed-off-by: Rahul Prabhu <raprabhu@nvidia.com>
1 parent 6dfa312 commit ddaea75

3 files changed

Lines changed: 172 additions & 0 deletions

File tree

‎delta-lake/delta-spark400db173/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuCheckOverflowInTableWrite.scala‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
/*
2+
* Copyright (c) 2026, NVIDIA CORPORATION.
3+
*
4+
* This file was derived from CheckOverflowInTableWrite in the
5+
* Delta Lake project at https://github.com/delta-io/delta.
6+
*
7+
* Copyright (2021) The Delta Lake Project Authors.
8+
*
9+
* Licensed under the Apache License, Version 2.0 (the "License");
10+
* you may not use this file except in compliance with the License.
11+
* You may obtain a copy of the License at
12+
*
13+
* http://www.apache.org/licenses/LICENSE-2.0
14+
*
15+
* Unless required by applicable law or agreed to in writing, software
16+
* distributed under the License is distributed on an "AS IS" BASIS,
17+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
18+
* See the License for the specific language governing permissions and
19+
* limitations under the License.
20+
*/
21+
22+
package com.databricks.sql.transaction.tahoe.rapids
23+
24+
import com.databricks.sql.transaction.tahoe.{CheckOverflowInTableWrite, DeltaErrors}
25+
import com.nvidia.spark.rapids._
26+
import com.nvidia.spark.rapids.shims.ShimUnaryExpression
27+
28+
import org.apache.spark.sql.catalyst.expressions.Expression
29+
import org.apache.spark.sql.types.DataType
30+
import org.apache.spark.sql.vectorized.ColumnarBatch
31+
32+
/** GPU version of Delta's CheckOverflowInTableWrite expression. */
33+
case class GpuCheckOverflowInTableWrite(child: GpuCast, columnName: String)
34+
extends ShimUnaryExpression with GpuExpression {
35+
36+
override def dataType: DataType = child.dataType
37+
38+
override def columnarEval(batch: ColumnarBatch): GpuColumnVector = {
39+
try {
40+
child.columnarEval(batch)
41+
} catch {
42+
case _: ArithmeticException =>
43+
throw DeltaErrors.castingCauseOverflowErrorInTableWrite(
44+
child.child.dataType,
45+
dataType,
46+
columnName)
47+
}
48+
}
49+
50+
override def sql: String = child.sql
51+
52+
override def toString: String = child.toString
53+
}
54+
55+
object GpuCheckOverflowInTableWrite {
56+
val exprRule: ExprRule[CheckOverflowInTableWrite] =
57+
GpuOverrides.expr[CheckOverflowInTableWrite](
58+
"Casting a numeric value as another numeric type in a Delta table write",
59+
ExprChecks.unaryProjectInputMatchesOutput(TypeSig.all, TypeSig.all),
60+
(check, conf, parent, rule) =>
61+
new UnaryExprMeta[CheckOverflowInTableWrite](check, conf, parent, rule) {
62+
override def convertToGpu(child: Expression): GpuExpression = child match {
63+
case cast: GpuCast => GpuCheckOverflowInTableWrite(cast, check.columnName)
64+
case _ =>
65+
throw new IllegalStateException("Expression child is not of type GpuCast")
66+
}
67+
})
68+
}

‎delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/DeltaSpark400DB173Provider.scala‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ import com.databricks.sql.transaction.tahoe.coordinatedcommits.{
3131
CoordinatedCommitsUtils
3232
}
3333
import com.databricks.sql.transaction.tahoe.rapids.{
34+
GpuCheckOverflowInTableWrite,
3435
GpuDeltaLog,
3536
GpuDeltaV1Write,
3637
GpuWriteIntoDelta,
@@ -83,6 +84,11 @@ object DeltaSpark400DB173Provider extends DatabricksDeltaProviderBase {
8384
(a, conf, p, r) => new GpuWriteIntoDeltaCommandMeta(a, conf, p, r))
8485
).map(r => (r.getClassFor.asSubclass(classOf[DataWritingCommand]), r)).toMap
8586
}
87+
88+
override def getExprs: Map[Class[_ <: Expression], ExprRule[_ <: Expression]] = {
89+
val rule = GpuCheckOverflowInTableWrite.exprRule
90+
super.getExprs + (rule.getClassFor.asSubclass(classOf[Expression]) -> rule)
91+
}
8692

8793
override def getRunnableCommandRules: Map[Class[_ <: RunnableCommand],
8894
RunnableCommandRule[_ <: RunnableCommand]] = {

‎integration_tests/src/main/python/delta_lake_merge_test.py‎

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,9 @@
3131
# Disable AQE temporarily until https://github.com/NVIDIA/spark-rapids/issues/14319 is resolved.
3232
delta_merge_enabled_conf = copy_and_update(delta_merge_enabled_conf, {"spark.sql.adaptive.enabled": "false"})
3333

34+
delta_merge_no_cpu_bridge_conf = copy_and_update(
35+
delta_merge_enabled_conf, {"spark.rapids.sql.expression.cpuBridge.enabled": "false"})
36+
3437
fallback_test_params = [{"spark.rapids.sql.format.delta.write.enabled": "false"},
3538
{"spark.rapids.sql.format.parquet.enabled": "false"},
3639
{"spark.rapids.sql.format.parquet.write.enabled": "false"},
@@ -43,6 +46,101 @@
4346
# See https://github.com/NVIDIA/spark-rapids/issues/13021#issuecomment-3166724473 for details.
4447
fallback_test_params.append(delta_writes_enabled_conf)
4548

49+
50+
def _assert_gpu_merge_processor(do_merge, data_path, conf, expect_write=True):
51+
assert expect_write
52+
cpu_result = with_cpu_session(lambda spark: do_merge(spark, data_path + "/CPU"), conf=conf)
53+
54+
callback = spark_jvm().org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback
55+
callback.startCapture()
56+
try:
57+
gpu_result = with_gpu_session(
58+
lambda spark: do_merge(spark, data_path + "/GPU"), conf=conf)
59+
captured_plans = callback.getResultsWithTimeout(10000)
60+
finally:
61+
callback.endCapture()
62+
63+
assert_equal(cpu_result, gpu_result)
64+
# The CPU expression bridge is disabled for this test, so finding the GPU merge processor
65+
# proves that CheckOverflowInTableWrite and the other merge expressions were replaced on GPU.
66+
class_name = "GpuRapidsProcessDeltaMergeJoinExec"
67+
assert any(callback.contains(plan, class_name) for plan in captured_plans), \
68+
f"{class_name} was not found in the captured MERGE plans"
69+
70+
71+
@allow_non_gpu(*delta_meta_allow)
72+
@delta_lake
73+
@ignore_order
74+
@pytest.mark.skipif(not is_databricks173_or_later(),
75+
reason="CheckOverflowInTableWrite is required on DBR 17.3+")
76+
def test_delta_merge_check_overflow_in_table_write(
77+
spark_tmp_path, spark_tmp_table_factory):
78+
def source_df(spark):
79+
return spark.sql("""
80+
SELECT timestamp '2024-01-01 00:00:00' AS timestampF,
81+
CAST(2 AS INT) AS byteF
82+
""")
83+
84+
def target_df(spark):
85+
return spark.sql("""
86+
SELECT timestamp '2024-01-01 00:00:00' AS timestampF,
87+
CAST(1 AS TINYINT) AS byteF
88+
UNION ALL
89+
SELECT timestamp '2024-01-02 00:00:00' AS timestampF,
90+
CAST(2 AS TINYINT) AS byteF
91+
""")
92+
93+
merge_sql = "MERGE INTO {dest_table} AS target USING {src_table} AS source " \
94+
"ON target.timestampF = source.timestampF " \
95+
"WHEN MATCHED THEN UPDATE SET byteF = source.byteF"
96+
assert_delta_sql_merge_collect(
97+
spark_tmp_path,
98+
spark_tmp_table_factory,
99+
use_cdf=False,
100+
enable_deletion_vectors=False,
101+
src_table_func=source_df,
102+
dest_table_func=target_df,
103+
merge_sql=merge_sql,
104+
compare_logs=False,
105+
assert_func=_assert_gpu_merge_processor,
106+
conf=delta_merge_no_cpu_bridge_conf)
107+
108+
109+
@allow_non_gpu(*delta_meta_allow)
110+
@delta_lake
111+
@pytest.mark.skipif(not is_databricks173_or_later(),
112+
reason="CheckOverflowInTableWrite is required on DBR 17.3+")
113+
def test_delta_merge_check_overflow_in_table_write_error(
114+
spark_tmp_path, spark_tmp_table_factory):
115+
updates_view = spark_tmp_table_factory.get()
116+
117+
def do_merge(spark):
118+
gpu_enabled = str(spark.conf.get("spark.rapids.sql.enabled", "false")).lower() == "true"
119+
target_path = spark_tmp_path + ("/GPU" if gpu_enabled else "/CPU")
120+
spark.sql("""
121+
SELECT timestamp '2024-01-01 00:00:00' AS timestampF,
122+
CAST(1 AS TINYINT) AS byteF
123+
""").write.format("delta") \
124+
.option("delta.enableDeletionVectors", "false") \
125+
.mode("overwrite") \
126+
.save(target_path)
127+
spark.sql("""
128+
SELECT timestamp '2024-01-01 00:00:00' AS timestampF,
129+
CAST(128 AS INT) AS byteF
130+
""").createOrReplaceTempView(updates_view)
131+
return spark.sql(f"""
132+
MERGE INTO delta.`{target_path}` AS target
133+
USING {updates_view} AS source
134+
ON target.timestampF = source.timestampF
135+
WHEN MATCHED THEN UPDATE SET byteF = source.byteF
136+
""").collect()
137+
138+
assert_gpu_and_cpu_error(
139+
do_merge,
140+
conf=delta_merge_no_cpu_bridge_conf,
141+
error_message="DELTA_CAST_OVERFLOW_IN_TABLE_WRITE")
142+
143+
46144
@allow_non_gpu(delta_write_fallback_allow, *delta_meta_allow)
47145
@delta_lake
48146
@ignore_order

0 commit comments

Comments
 (0)