Merge remote-tracking branch 'apache/main' into ghj

Author: andygrove

dev/diffs/3.5.8.diff

@@ -76,7 +76,8 @@ index 27ae10b3d59..78e69902dfd 100644
      val extensionConfClassNames = sparkContext.getConf.get(StaticSQLConf.SPARK_SESSION_EXTENSIONS)
        .getOrElse(Seq.empty)
 +    val extensionClassNames = extensionConfClassNames ++ loadCometExtension(sparkContext)
-     extensionConfClassNames.foreach { extensionConfClassName =>
+-    extensionConfClassNames.foreach { extensionConfClassName =>
++    extensionClassNames.foreach { extensionConfClassName =>
        try {
          val extensionConfClass = Utils.classForName(extensionConfClassName)
 @@ -1396,4 +1405,12 @@ object SparkSession extends Logging {
@@ -3271,7 +3272,38 @@ index f3be79f9022..b4b1ea8dbc4 100644
 +    IgnoreCometNativeDataFusion("https://github.com/apache/datafusion-comet/issues/3312")) {
 
      withTempDir { tempDir =>
- 
+
+diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
+index 8b4ac474..3f79f208 100644
+--- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
++++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
+@@ -223,6 +223,8 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt
+     withSession(extensions) { session =>
+       session.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED, true)
+       session.conf.set(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key, "-1")
++      // https://github.com/apache/datafusion-comet/issues/1197
++      session.conf.set("spark.comet.enabled", false)
+       assert(session.sessionState.columnarRules.contains(
+         MyColumnarRule(PreRuleReplaceAddWithBrokenVersion(), MyPostRule())))
+       import session.sqlContext.implicits._
+@@ -281,6 +283,8 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt
+     }
+     withSession(extensions) { session =>
+       session.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED, enableAQE)
++      // https://github.com/apache/datafusion-comet/issues/1197
++      session.conf.set("spark.comet.enabled", false)
+       assert(session.sessionState.columnarRules.contains(
+         MyColumnarRule(PreRuleReplaceAddWithBrokenVersion(), MyPostRule())))
+       import session.sqlContext.implicits._
+@@ -319,6 +323,8 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper with Adapt
+     val session = SparkSession.builder()
+       .master("local[1]")
+       .config(COLUMN_BATCH_SIZE.key, 2)
++      // https://github.com/apache/datafusion-comet/issues/1197
++      .config("spark.comet.enabled", false)
+       .withExtensions { extensions =>
+         extensions.injectColumnar(session =>
+           MyColumnarRule(PreRuleReplaceAddWithBrokenVersion(), MyPostRule())) }
 diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala
 index 6160c3e5f6c..0956d7d9edc 100644
 --- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala

native/spark-expr/src/datetime_funcs/date_trunc.rs

@@ -16,7 +16,9 @@
 // under the License.
 
 use arrow::datatypes::DataType;
-use datafusion::common::{utils::take_function_args, DataFusionError, Result, ScalarValue::Utf8};
+use datafusion::common::{
+    utils::take_function_args, DataFusionError, Result, ScalarValue, ScalarValue::Utf8,
+};
 use datafusion::logical_expr::{
     ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
 };
@@ -76,9 +78,14 @@ impl ScalarUDFImpl for SparkDateTrunc {
                 let result = date_trunc_array_fmt_dyn(&date, &formats)?;
                 Ok(ColumnarValue::Array(result))
             }
+            (ColumnarValue::Scalar(date_scalar), ColumnarValue::Scalar(Utf8(Some(format)))) => {
+                let date_arr = date_scalar.to_array()?;
+                let result = date_trunc_dyn(&date_arr, format)?;
+                let scalar = ScalarValue::try_from_array(&result, 0)?;
+                Ok(ColumnarValue::Scalar(scalar))
+            }
             _ => Err(DataFusionError::Execution(
-                "Invalid input to function DateTrunc. Expected (PrimitiveArray<Date32>, Scalar) or \
-                    (PrimitiveArray<Date32>, StringArray)".to_string(),
+                "Invalid input to function DateTrunc. Expected (Date32, Utf8)".to_string(),
             )),
         }
     }

native/spark-expr/src/datetime_funcs/timestamp_trunc.rs

@@ -18,7 +18,7 @@
 use crate::utils::array_with_timezone;
 use arrow::datatypes::{DataType, Schema, TimeUnit::Microsecond};
 use arrow::record_batch::RecordBatch;
-use datafusion::common::{DataFusionError, ScalarValue::Utf8};
+use datafusion::common::{DataFusionError, ScalarValue, ScalarValue::Utf8};
 use datafusion::logical_expr::ColumnarValue;
 use datafusion::physical_expr::PhysicalExpr;
 use std::hash::Hash;
@@ -130,10 +130,20 @@ impl PhysicalExpr for TimestampTruncExpr {
                 let result = timestamp_trunc_array_fmt_dyn(&ts, &formats)?;
                 Ok(ColumnarValue::Array(result))
             }
+            (ColumnarValue::Scalar(ts_scalar), ColumnarValue::Scalar(Utf8(Some(format)))) => {
+                let ts_arr = ts_scalar.to_array()?;
+                let ts = array_with_timezone(
+                    ts_arr,
+                    tz.clone(),
+                    Some(&DataType::Timestamp(Microsecond, Some(tz.into()))),
+                )?;
+                let result = timestamp_trunc_dyn(&ts, format)?;
+                let scalar = ScalarValue::try_from_array(&result, 0)?;
+                Ok(ColumnarValue::Scalar(scalar))
+            }
             _ => Err(DataFusionError::Execution(
                 "Invalid input to function TimestampTrunc. \
-                    Expected (PrimitiveArray<TimestampMicrosecondType>, Scalar, String) or \
-                    (PrimitiveArray<TimestampMicrosecondType>, StringArray, String)"
+                    Expected (Timestamp, Utf8)"
                     .to_string(),
             )),
         }

spark/src/test/resources/sql-tests/expressions/datetime/trunc_date.sql

@@ -33,5 +33,5 @@ query
 SELECT trunc(d, 'quarter') FROM test_trunc_date
 
 -- literal arguments
-query ignore(https://github.com/apache/datafusion-comet/issues/3342)
+query
 SELECT trunc(date('2024-06-15'), 'year'), trunc(date('2024-06-15'), 'month'), trunc(date('2024-06-15'), 'quarter')

spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp.sql

@@ -37,5 +37,5 @@ query
 SELECT date_trunc('hour', ts) FROM test_trunc_ts
 
 -- literal arguments
-query ignore(https://github.com/apache/datafusion-comet/issues/3342)
+query
 SELECT date_trunc('year', timestamp('2024-06-15 10:30:45')), date_trunc('month', timestamp('2024-06-15 10:30:45')), date_trunc('day', timestamp('2024-06-15 10:30:45'))