Skip to content

Commit e050c4f

Browse files
committed
feat: add MapSort support for Spark 4.0 via shim
Adds MapSort serialization to Spark 4.0 version shim
1 parent 07d1013 commit e050c4f

1 file changed

Lines changed: 14 additions & 3 deletions

File tree

spark/src/main/spark-4.0/org/apache/comet/shims/CometExprShim.scala

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,16 +20,16 @@
2020
package org.apache.comet.shims
2121

2222
import org.apache.spark.sql.catalyst.expressions._
23-
import org.apache.spark.sql.catalyst.expressions.objects.StaticInvoke
23+
import org.apache.spark.sql.catalyst.expressions.objects.{RowOrdering, StaticInvoke}
2424
import org.apache.spark.sql.internal.SQLConf
2525
import org.apache.spark.sql.internal.types.StringTypeWithCollation
26-
import org.apache.spark.sql.types.{BinaryType, BooleanType, DataTypes, StringType}
26+
import org.apache.spark.sql.types.{BinaryType, BooleanType, DataTypes, MapType, StringType}
2727

2828
import org.apache.comet.CometSparkSessionExtensions.withInfo
2929
import org.apache.comet.expressions.{CometCast, CometEvalMode}
3030
import org.apache.comet.serde.{CommonStringExprs, Compatible, ExprOuterClass, Incompatible}
3131
import org.apache.comet.serde.ExprOuterClass.{BinaryOutputStyle, Expr}
32-
import org.apache.comet.serde.QueryPlanSerde.exprToProtoInternal
32+
import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, optExprWithInfo, scalarFunctionExprToProto}
3333

3434
/**
3535
* `CometExprShim` acts as a shim for parsing expressions from different Spark versions.
@@ -103,6 +103,17 @@ trait CometExprShim extends CommonStringExprs {
103103
None
104104
}
105105

106+
case expr: MapSort =>
107+
val keyType = expr.base.dataType.asInstanceOf[MapType].keyType
108+
if (!RowOrdering.isOrderable(keyType)) {
109+
withInfo(expr, s"map_sort requires orderable key type, got: $keyType")
110+
return None
111+
}
112+
113+
val childExpr = exprToProtoInternal(expr.base, inputs, binding)
114+
val mapSortScalarExpr = scalarFunctionExprToProto("map_sort", childExpr)
115+
optExprWithInfo(mapSortScalarExpr, expr, expr.children: _*)
116+
106117
case wb: WidthBucket =>
107118
withInfo(
108119
wb,

0 commit comments

Comments
 (0)