fix aggregation wrapping now that we don't have an extra CometExecIterator.

Author: mbutrovich

native/shuffle/src/shuffle_writer.rs

@@ -29,7 +29,7 @@ use datafusion::physical_expr::{EquivalenceProperties, Partitioning};
 use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
 use datafusion::physical_plan::EmptyRecordBatchStream;
 use datafusion::{
-    arrow::{datatypes::SchemaRef, error::ArrowError},
+    arrow::datatypes::SchemaRef,
     error::Result,
     execution::context::TaskContext,
     physical_plan::{
@@ -38,7 +38,7 @@ use datafusion::{
         DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream,
     },
 };
-use futures::{StreamExt, TryFutureExt, TryStreamExt};
+use futures::{StreamExt, TryStreamExt};
 use std::{
     any::Any,
     fmt,
@@ -171,23 +171,23 @@ impl ExecutionPlan for ShuffleWriterExec {
         let input = self.input.execute(partition, Arc::clone(&context))?;
         let metrics = ShufflePartitionerMetrics::new(&self.metrics, 0);
 
+        // Propagate DataFusionError unchanged: the JNI bridge only downcasts a single
+        // `DataFusionError::External(SparkError)` layer, so any extra wrap here loses the
+        // typed exception (e.g. SparkArithmeticException on decimal overflow).
         Ok(Box::pin(RecordBatchStreamAdapter::new(
             self.schema(),
-            futures::stream::once(
-                external_shuffle(
-                    input,
-                    partition,
-                    self.output_data_file.clone(),
-                    self.output_index_file.clone(),
-                    self.partitioning.clone(),
-                    metrics,
-                    context,
-                    self.codec.clone(),
-                    self.tracing_enabled,
-                    self.write_buffer_size,
-                )
-                .map_err(|e| ArrowError::ExternalError(Box::new(e))),
-            )
+            futures::stream::once(external_shuffle(
+                input,
+                partition,
+                self.output_data_file.clone(),
+                self.output_index_file.clone(),
+                self.partitioning.clone(),
+                metrics,
+                context,
+                self.codec.clone(),
+                self.tracing_enabled,
+                self.write_buffer_size,
+            ))
             .try_flatten(),
         )))
     }

spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala

@@ -1817,6 +1817,10 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
                 // make sure that the error message throws overflow exception only
                 assert(sparkExc.getMessage.contains("ARITHMETIC_OVERFLOW"))
                 assert(cometExc.getMessage.contains("ARITHMETIC_OVERFLOW"))
+                assert(
+                  cometExc.isInstanceOf[ArithmeticException],
+                  "expected ArithmeticException, got " +
+                    s"${cometExc.getClass.getName}: ${cometExc.getMessage}")
               case _ => fail("Exception should be thrown for Long overflow in ANSI mode")
             }
           } else {
@@ -1831,6 +1835,10 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
               case (Some(sparkExc), Some(cometExc)) =>
                 assert(sparkExc.getMessage.contains("ARITHMETIC_OVERFLOW"))
                 assert(cometExc.getMessage.contains("ARITHMETIC_OVERFLOW"))
+                assert(
+                  cometExc.isInstanceOf[ArithmeticException],
+                  "expected ArithmeticException, got " +
+                    s"${cometExc.getClass.getName}: ${cometExc.getMessage}")
               case _ => fail("Exception should be thrown for Long underflow in ANSI mode")
             }
           } else {
@@ -1870,6 +1878,10 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
               case (Some(sparkExc), Some(cometExc)) =>
                 assert(sparkExc.getMessage.contains("ARITHMETIC_OVERFLOW"))
                 assert(cometExc.getMessage.contains("ARITHMETIC_OVERFLOW"))
+                assert(
+                  cometExc.isInstanceOf[ArithmeticException],
+                  "expected ArithmeticException, got " +
+                    s"${cometExc.getClass.getName}: ${cometExc.getMessage}")
               case _ =>
                 fail("Exception should be thrown for decimal overflow in ANSI mode")
             }
@@ -1893,6 +1905,10 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
               case (Some(sparkExc), Some(cometExc)) =>
                 assert(sparkExc.getMessage.contains("ARITHMETIC_OVERFLOW"))
                 assert(cometExc.getMessage.contains("ARITHMETIC_OVERFLOW"))
+                assert(
+                  cometExc.isInstanceOf[ArithmeticException],
+                  "expected ArithmeticException, got " +
+                    s"${cometExc.getClass.getName}: ${cometExc.getMessage}")
               case _ =>
                 fail("Exception should be thrown for Long overflow with GROUP BY in ANSI mode")
             }
@@ -1910,6 +1926,10 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
               case (Some(sparkExc), Some(cometExc)) =>
                 assert(sparkExc.getMessage.contains("ARITHMETIC_OVERFLOW"))
                 assert(cometExc.getMessage.contains("ARITHMETIC_OVERFLOW"))
+                assert(
+                  cometExc.isInstanceOf[ArithmeticException],
+                  "expected ArithmeticException, got " +
+                    s"${cometExc.getClass.getName}: ${cometExc.getMessage}")
               case _ =>
                 fail("Exception should be thrown for Long underflow with GROUP BY in ANSI mode")
             }
@@ -1951,6 +1971,10 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
               case (Some(sparkExc), Some(cometExc)) =>
                 assert(sparkExc.getMessage.contains("ARITHMETIC_OVERFLOW"))
                 assert(cometExc.getMessage.contains("ARITHMETIC_OVERFLOW"))
+                assert(
+                  cometExc.isInstanceOf[ArithmeticException],
+                  "expected ArithmeticException, got " +
+                    s"${cometExc.getClass.getName}: ${cometExc.getMessage}")
               case _ =>
                 fail("Exception should be thrown for decimal overflow with GROUP BY in ANSI mode")
             }