|
20 | 20 | package org.apache.comet.shims |
21 | 21 |
|
22 | 22 | 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} |
24 | 24 | import org.apache.spark.sql.internal.SQLConf |
25 | 25 | 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} |
27 | 27 |
|
28 | 28 | import org.apache.comet.CometSparkSessionExtensions.withInfo |
29 | 29 | import org.apache.comet.expressions.{CometCast, CometEvalMode} |
30 | 30 | import org.apache.comet.serde.{CommonStringExprs, Compatible, ExprOuterClass, Incompatible} |
31 | 31 | 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} |
33 | 33 |
|
34 | 34 | /** |
35 | 35 | * `CometExprShim` acts as a shim for parsing expressions from different Spark versions. |
@@ -103,6 +103,17 @@ trait CometExprShim extends CommonStringExprs { |
103 | 103 | None |
104 | 104 | } |
105 | 105 |
|
| 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 | + |
106 | 117 | case wb: WidthBucket => |
107 | 118 | withInfo( |
108 | 119 | wb, |
|
0 commit comments