Fix test cases

szehon-ho · Apr 28, 2024 · 9776242 · 9776242
1 parent 41481b4
commit 9776242
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 9 deletions.
diff --git a/core/src/main/scala/org/apache/spark/Partitioner.scala b/core/src/main/scala/org/apache/spark/Partitioner.scala
@@ -151,7 +151,8 @@ private[spark] class KeyGroupedPartitioner(
   override def getPartition(key: Any): Int = {
     val keys = key.asInstanceOf[Seq[Any]]
     val normalizedKeys = ArraySeq.from(keys)
-    valueMap.getOrElseUpdate(normalizedKeys, Utils.nonNegativeMod(keys.hashCode, numPartitions))
+    valueMap.getOrElseUpdate(normalizedKeys,
+      Utils.nonNegativeMod(normalizedKeys.hashCode, numPartitions))
   }
 }
 

diff --git a/...ore/src/test/scala/org/apache/spark/sql/connector/WriteDistributionAndOrderingSuite.scala b/...ore/src/test/scala/org/apache/spark/sql/connector/WriteDistributionAndOrderingSuite.scala
@@ -20,10 +20,9 @@ package org.apache.spark.sql.connector
 
 import java.sql.Date
 import java.util.Collections
-
-import org.apache.spark.sql.{catalyst, AnalysisException, DataFrame, Row}
+import org.apache.spark.sql.{AnalysisException, DataFrame, Row, catalyst}
 import org.apache.spark.sql.catalyst.expressions.{ApplyFunctionExpression, Cast, Literal}
-import org.apache.spark.sql.catalyst.expressions.objects.Invoke
+import org.apache.spark.sql.catalyst.expressions.objects.StaticInvoke
 import org.apache.spark.sql.catalyst.plans.physical
 import org.apache.spark.sql.catalyst.plans.physical.{CoalescedBoundary, CoalescedHashPartitioning, HashPartitioning, RangePartitioning, UnknownPartitioning}
 import org.apache.spark.sql.connector.catalog.{Column, Identifier}
@@ -40,7 +39,7 @@ import org.apache.spark.sql.execution.streaming.sources.ContinuousMemoryStream
 import org.apache.spark.sql.functions.lit
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.streaming.{StreamingQueryException, Trigger}
-import org.apache.spark.sql.types.{DateType, IntegerType, LongType, ObjectType, StringType, TimestampType}
+import org.apache.spark.sql.types.{DataTypes, DateType, IntegerType, LongType, StringType, TimestampType}
 import org.apache.spark.sql.util.QueryExecutionListener
 import org.apache.spark.tags.SlowSQLTest
 
@@ -1128,13 +1127,16 @@ class WriteDistributionAndOrderingSuite extends DistributionAndOrderingSuiteBase
         Seq.empty
       ),
       catalyst.expressions.SortOrder(
-        Invoke(
-          Literal.create(YearsFunction, ObjectType(YearsFunction.getClass)),
+        StaticInvoke(
+          YearsFunction.getClass,
+          DataTypes.LongType,
           "invoke",
-          LongType,
           Seq(Cast(attr("day"), TimestampType, Some("America/Los_Angeles"))),
           Seq(TimestampType),
-          propagateNull = false),
+          propagateNull = false,
+          returnNullable = true,
+          isDeterministic = true,
+          Some(YearsFunction)),
         catalyst.expressions.Descending,
         catalyst.expressions.NullsFirst,
         Seq.empty