Skip to content

Commit 56a2f7e

Browse files
andygroveandygrove
authored
chore: Split expression serde hash map into separate categories (apache#2322)
1 parent b9ce50c commit 56a2f7e

File tree

1 file changed

+142
-119
lines changed

1 file changed

+142
-119
lines changed

‎spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala‎

Lines changed: 142 additions & 119 deletions
Original file line numberDiff line numberDiff line change
@@ -71,20 +71,9 @@ object QueryPlanSerde extends Logging with CometExprShim {
7171
private val opSerdeMap: Map[Class[_ <: SparkPlan], CometOperatorSerde[_]] =
7272
Map(classOf[ProjectExec] -> CometProject, classOf[SortExec] -> CometSort)
7373

74-
/**
75-
* Mapping of Spark expression class to Comet expression handler.
76-
*/
77-
private val exprSerdeMap: Map[Class[_ <: Expression], CometExpressionSerde[_]] = Map(
78-
classOf[AttributeReference] -> CometAttributeReference,
79-
classOf[Alias] -> CometAlias,
80-
classOf[Add] -> CometAdd,
81-
classOf[Subtract] -> CometSubtract,
82-
classOf[Multiply] -> CometMultiply,
83-
classOf[Divide] -> CometDivide,
84-
classOf[IntegralDivide] -> CometIntegralDivide,
85-
classOf[Remainder] -> CometRemainder,
86-
classOf[Round] -> CometRound,
74+
private val arrayExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] = Map(
8775
classOf[ArrayAppend] -> CometArrayAppend,
76+
// TODO ArrayCompact
8877
classOf[ArrayContains] -> CometArrayContains,
8978
classOf[ArrayDistinct] -> CometArrayDistinct,
9079
classOf[ArrayExcept] -> CometArrayExcept,
@@ -98,129 +87,179 @@ object QueryPlanSerde extends Logging with CometExprShim {
9887
classOf[ArraysOverlap] -> CometArraysOverlap,
9988
classOf[ArrayUnion] -> CometArrayUnion,
10089
classOf[CreateArray] -> CometCreateArray,
101-
classOf[GetArrayItem] -> CometGetArrayItem,
10290
classOf[ElementAt] -> CometElementAt,
103-
classOf[Ascii] -> CometScalarFunction("ascii"),
104-
classOf[ConcatWs] -> CometScalarFunction("concat_ws"),
105-
classOf[Chr] -> CometScalarFunction("char"),
106-
classOf[InitCap] -> CometInitCap,
107-
classOf[BitwiseCount] -> CometBitwiseCount,
108-
classOf[BitwiseGet] -> CometBitwiseGet,
109-
classOf[BitwiseNot] -> CometBitwiseNot,
110-
classOf[BitwiseAnd] -> CometBitwiseAnd,
111-
classOf[BitwiseOr] -> CometBitwiseOr,
112-
classOf[BitwiseXor] -> CometBitwiseXor,
113-
classOf[BitLength] -> CometScalarFunction("bit_length"),
114-
classOf[FromUnixTime] -> CometFromUnixTime,
115-
classOf[Length] -> CometScalarFunction("length"),
116-
classOf[Acos] -> CometScalarFunction("acos"),
117-
classOf[Cos] -> CometScalarFunction("cos"),
118-
classOf[Asin] -> CometScalarFunction("asin"),
119-
classOf[Sin] -> CometScalarFunction("sin"),
120-
classOf[Atan] -> CometScalarFunction("atan"),
121-
classOf[Tan] -> CometScalarFunction("tan"),
122-
classOf[Exp] -> CometScalarFunction("exp"),
123-
classOf[Expm1] -> CometScalarFunction("expm1"),
124-
classOf[Sqrt] -> CometScalarFunction("sqrt"),
125-
classOf[Signum] -> CometScalarFunction("signum"),
126-
classOf[Md5] -> CometScalarFunction("md5"),
127-
classOf[ShiftLeft] -> CometShiftLeft,
128-
classOf[ShiftRight] -> CometShiftRight,
129-
classOf[StringInstr] -> CometScalarFunction("instr"),
130-
classOf[StringRepeat] -> CometStringRepeat,
131-
classOf[StringReplace] -> CometScalarFunction("replace"),
132-
classOf[StringTranslate] -> CometScalarFunction("translate"),
133-
classOf[StringTrim] -> CometScalarFunction("trim"),
134-
classOf[StringTrimLeft] -> CometScalarFunction("ltrim"),
135-
classOf[StringTrimRight] -> CometScalarFunction("rtrim"),
136-
classOf[StringTrimBoth] -> CometScalarFunction("btrim"),
137-
classOf[Upper] -> CometUpper,
138-
classOf[Lower] -> CometLower,
139-
classOf[Murmur3Hash] -> CometMurmur3Hash,
140-
classOf[XxHash64] -> CometXxHash64,
141-
classOf[Sha2] -> CometSha2,
142-
classOf[MapKeys] -> CometMapKeys,
143-
classOf[MapEntries] -> CometMapEntries,
144-
classOf[MapValues] -> CometMapValues,
145-
classOf[MapFromArrays] -> CometMapFromArrays,
146-
classOf[GetMapValue] -> CometMapExtract,
91+
classOf[Flatten] -> CometFlatten,
92+
classOf[GetArrayItem] -> CometGetArrayItem)
93+
94+
private val conditionalExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] =
95+
Map(classOf[CaseWhen] -> CometCaseWhen, classOf[If] -> CometIf)
96+
97+
private val predicateExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] = Map(
98+
classOf[And] -> CometAnd,
14799
classOf[EqualTo] -> CometEqualTo,
148100
classOf[EqualNullSafe] -> CometEqualNullSafe,
149-
classOf[Not] -> CometNot,
150-
classOf[And] -> CometAnd,
151-
classOf[Or] -> CometOr,
152101
classOf[GreaterThan] -> CometGreaterThan,
153102
classOf[GreaterThanOrEqual] -> CometGreaterThanOrEqual,
154103
classOf[LessThan] -> CometLessThan,
155104
classOf[LessThanOrEqual] -> CometLessThanOrEqual,
156-
classOf[IsNull] -> CometIsNull,
157-
classOf[IsNotNull] -> CometIsNotNull,
158-
classOf[IsNaN] -> CometIsNaN,
159105
classOf[In] -> CometIn,
106+
classOf[IsNotNull] -> CometIsNotNull,
107+
classOf[IsNull] -> CometIsNull,
160108
classOf[InSet] -> CometInSet,
109+
classOf[Not] -> CometNot,
110+
classOf[Or] -> CometOr)
111+
112+
private val mathExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] = Map(
113+
classOf[Acos] -> CometScalarFunction("acos"),
114+
classOf[Add] -> CometAdd,
115+
classOf[Asin] -> CometScalarFunction("asin"),
116+
classOf[Atan] -> CometScalarFunction("atan"),
117+
classOf[Atan2] -> CometAtan2,
118+
classOf[Ceil] -> CometCeil,
119+
classOf[Cos] -> CometScalarFunction("cos"),
120+
classOf[Divide] -> CometDivide,
121+
classOf[Exp] -> CometScalarFunction("exp"),
122+
classOf[Expm1] -> CometScalarFunction("expm1"),
123+
classOf[Floor] -> CometFloor,
124+
classOf[Hex] -> CometHex,
125+
classOf[IntegralDivide] -> CometIntegralDivide,
126+
classOf[IsNaN] -> CometIsNaN,
127+
classOf[Log] -> CometLog,
128+
classOf[Log2] -> CometLog2,
129+
classOf[Log10] -> CometLog10,
130+
classOf[Multiply] -> CometMultiply,
131+
classOf[Pow] -> CometScalarFunction("pow"),
161132
classOf[Rand] -> CometRand,
162133
classOf[Randn] -> CometRandn,
163-
classOf[SparkPartitionID] -> CometSparkPartitionId,
164-
classOf[MonotonicallyIncreasingID] -> CometMonotonicallyIncreasingId,
165-
classOf[StringSpace] -> CometScalarFunction("string_space"),
166-
classOf[StartsWith] -> CometScalarFunction("starts_with"),
167-
classOf[EndsWith] -> CometScalarFunction("ends_with"),
134+
classOf[Remainder] -> CometRemainder,
135+
classOf[Round] -> CometRound,
136+
classOf[Signum] -> CometScalarFunction("signum"),
137+
classOf[Sin] -> CometScalarFunction("sin"),
138+
classOf[Sqrt] -> CometScalarFunction("sqrt"),
139+
classOf[Subtract] -> CometSubtract,
140+
classOf[Tan] -> CometScalarFunction("tan"),
141+
// TODO UnaryMinus
142+
classOf[Unhex] -> CometUnhex)
143+
144+
private val mapExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] = Map(
145+
classOf[GetMapValue] -> CometMapExtract,
146+
classOf[MapKeys] -> CometMapKeys,
147+
classOf[MapEntries] -> CometMapEntries,
148+
classOf[MapValues] -> CometMapValues,
149+
classOf[MapFromArrays] -> CometMapFromArrays)
150+
151+
private val structExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] = Map(
152+
classOf[CreateNamedStruct] -> CometCreateNamedStruct,
153+
classOf[GetArrayStructFields] -> CometGetArrayStructFields,
154+
classOf[GetStructField] -> CometGetStructField,
155+
classOf[StructsToJson] -> CometStructsToJson)
156+
157+
private val hashExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] = Map(
158+
classOf[Md5] -> CometScalarFunction("md5"),
159+
classOf[Murmur3Hash] -> CometMurmur3Hash,
160+
classOf[Sha2] -> CometSha2,
161+
classOf[XxHash64] -> CometXxHash64)
162+
163+
private val stringExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] = Map(
164+
classOf[Ascii] -> CometScalarFunction("ascii"),
165+
classOf[BitLength] -> CometScalarFunction("bit_length"),
166+
classOf[Chr] -> CometScalarFunction("char"),
167+
classOf[ConcatWs] -> CometScalarFunction("concat_ws"),
168168
classOf[Contains] -> CometScalarFunction("contains"),
169-
classOf[Substring] -> CometSubstring,
169+
classOf[EndsWith] -> CometScalarFunction("ends_with"),
170+
classOf[InitCap] -> CometInitCap,
171+
classOf[Length] -> CometScalarFunction("length"),
170172
classOf[Like] -> CometLike,
171-
classOf[RLike] -> CometRLike,
173+
classOf[Lower] -> CometLower,
172174
classOf[OctetLength] -> CometScalarFunction("octet_length"),
173175
classOf[Reverse] -> CometScalarFunction("reverse"),
176+
classOf[RLike] -> CometRLike,
177+
classOf[StartsWith] -> CometScalarFunction("starts_with"),
178+
classOf[StringInstr] -> CometScalarFunction("instr"),
179+
classOf[StringRepeat] -> CometStringRepeat,
180+
classOf[StringReplace] -> CometScalarFunction("replace"),
174181
classOf[StringRPad] -> CometStringRPad,
175-
classOf[Year] -> CometYear,
182+
classOf[StringSpace] -> CometScalarFunction("string_space"),
183+
classOf[StringTranslate] -> CometScalarFunction("translate"),
184+
classOf[StringTrim] -> CometScalarFunction("trim"),
185+
classOf[StringTrimBoth] -> CometScalarFunction("btrim"),
186+
classOf[StringTrimLeft] -> CometScalarFunction("ltrim"),
187+
classOf[StringTrimRight] -> CometScalarFunction("rtrim"),
188+
classOf[Substring] -> CometSubstring,
189+
classOf[Upper] -> CometUpper)
190+
191+
private val bitwiseExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] = Map(
192+
classOf[BitwiseAnd] -> CometBitwiseAnd,
193+
classOf[BitwiseCount] -> CometBitwiseCount,
194+
classOf[BitwiseGet] -> CometBitwiseGet,
195+
classOf[BitwiseOr] -> CometBitwiseOr,
196+
classOf[BitwiseNot] -> CometBitwiseNot,
197+
classOf[BitwiseXor] -> CometBitwiseXor,
198+
classOf[ShiftLeft] -> CometShiftLeft,
199+
classOf[ShiftRight] -> CometShiftRight)
200+
201+
private val temporalExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] = Map(
202+
classOf[DateAdd] -> CometDateAdd,
203+
classOf[DateSub] -> CometDateSub,
204+
classOf[FromUnixTime] -> CometFromUnixTime,
176205
classOf[Hour] -> CometHour,
177206
classOf[Minute] -> CometMinute,
178207
classOf[Second] -> CometSecond,
179-
classOf[DateAdd] -> CometDateAdd,
180-
classOf[DateSub] -> CometDateSub,
181208
classOf[TruncDate] -> CometTruncDate,
182209
classOf[TruncTimestamp] -> CometTruncTimestamp,
183-
classOf[Cast] -> CometCast,
184-
classOf[CreateNamedStruct] -> CometCreateNamedStruct,
185-
classOf[GetStructField] -> CometGetStructField,
186-
classOf[GetArrayStructFields] -> CometGetArrayStructFields,
187-
classOf[StructsToJson] -> CometStructsToJson,
188-
classOf[Flatten] -> CometFlatten,
189-
classOf[Atan2] -> CometAtan2,
190-
classOf[Ceil] -> CometCeil,
191-
classOf[Floor] -> CometFloor,
192-
classOf[Log] -> CometLog,
193-
classOf[Log10] -> CometLog10,
194-
classOf[Log2] -> CometLog2,
195-
classOf[Hex] -> CometHex,
196-
classOf[Unhex] -> CometUnhex,
197-
classOf[Pow] -> CometScalarFunction[Pow]("pow"),
198-
classOf[If] -> CometIf,
199-
classOf[CaseWhen] -> CometCaseWhen,
210+
classOf[Year] -> CometYear)
211+
212+
private val conversionExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] = Map(
213+
classOf[Cast] -> CometCast)
214+
215+
private val miscExpressions: Map[Class[_ <: Expression], CometExpressionSerde[_]] = Map(
216+
// TODO Literal
217+
// TODO SortOrder (?)
218+
// TODO PromotePrecision
219+
// TODO CheckOverflow
220+
// TODO KnownFloatingPointNormalized
221+
// TODO ScalarSubquery
222+
// TODO UnscaledValue
223+
// TODO MakeDecimal
224+
// TODO BloomFilterMightContain
225+
// TODO RegExpReplace
226+
classOf[Alias] -> CometAlias,
227+
classOf[AttributeReference] -> CometAttributeReference,
228+
classOf[SparkPartitionID] -> CometSparkPartitionId,
229+
classOf[MonotonicallyIncreasingID] -> CometMonotonicallyIncreasingId,
200230
classOf[Coalesce] -> CometCoalesce)
201231

232+
/**
233+
* Mapping of Spark expression class to Comet expression handler.
234+
*/
235+
private val exprSerdeMap: Map[Class[_ <: Expression], CometExpressionSerde[_]] =
236+
mathExpressions ++ hashExpressions ++ stringExpressions ++
237+
conditionalExpressions ++ mapExpressions ++ predicateExpressions ++
238+
structExpressions ++ bitwiseExpressions ++ miscExpressions ++ arrayExpressions ++
239+
temporalExpressions ++ conversionExpressions
240+
202241
/**
203242
* Mapping of Spark aggregate expression class to Comet expression handler.
204243
*/
205244
private val aggrSerdeMap: Map[Class[_], CometAggregateExpressionSerde[_]] = Map(
206-
classOf[Sum] -> CometSum,
207245
classOf[Average] -> CometAverage,
208-
classOf[Count] -> CometCount,
209-
classOf[Min] -> CometMin,
210-
classOf[Max] -> CometMax,
211-
classOf[First] -> CometFirst,
212-
classOf[Last] -> CometLast,
213246
classOf[BitAndAgg] -> CometBitAndAgg,
214247
classOf[BitOrAgg] -> CometBitOrAgg,
215248
classOf[BitXorAgg] -> CometBitXOrAgg,
216-
classOf[CovSample] -> CometCovSample,
249+
classOf[BloomFilterAggregate] -> CometBloomFilterAggregate,
250+
classOf[Corr] -> CometCorr,
251+
classOf[Count] -> CometCount,
217252
classOf[CovPopulation] -> CometCovPopulation,
218-
classOf[VarianceSamp] -> CometVarianceSamp,
219-
classOf[VariancePop] -> CometVariancePop,
220-
classOf[StddevSamp] -> CometStddevSamp,
253+
classOf[CovSample] -> CometCovSample,
254+
classOf[First] -> CometFirst,
255+
classOf[Last] -> CometLast,
256+
classOf[Max] -> CometMax,
257+
classOf[Min] -> CometMin,
221258
classOf[StddevPop] -> CometStddevPop,
222-
classOf[Corr] -> CometCorr,
223-
classOf[BloomFilterAggregate] -> CometBloomFilterAggregate)
259+
classOf[StddevSamp] -> CometStddevSamp,
260+
classOf[Sum] -> CometSum,
261+
classOf[VariancePop] -> CometVariancePop,
262+
classOf[VarianceSamp] -> CometVarianceSamp)
224263

225264
def supportedDataType(dt: DataType, allowComplex: Boolean = false): Boolean = dt match {
226265
case _: ByteType | _: ShortType | _: IntegerType | _: LongType | _: FloatType |
@@ -873,22 +912,6 @@ object QueryPlanSerde extends Logging with CometExprShim {
873912
None
874913
}
875914

876-
// abs implementation is not correct
877-
// https://github.com/apache/datafusion-comet/issues/666
878-
// case Abs(child, failOnErr) =>
879-
// val childExpr = exprToProtoInternal(child, inputs)
880-
// if (childExpr.isDefined) {
881-
// val evalModeStr =
882-
// if (failOnErr) ExprOuterClass.EvalMode.ANSI else ExprOuterClass.EvalMode.LEGACY
883-
// val absBuilder = ExprOuterClass.Abs.newBuilder()
884-
// absBuilder.setChild(childExpr.get)
885-
// absBuilder.setEvalMode(evalModeStr)
886-
// Some(Expr.newBuilder().setAbs(absBuilder).build())
887-
// } else {
888-
// withInfo(expr, child)
889-
// None
890-
// }
891-
892915
case RegExpReplace(subject, pattern, replacement, startPosition) =>
893916
if (!RegExp.isSupportedPattern(pattern.toString) &&
894917
!CometConf.COMET_REGEXP_ALLOW_INCOMPATIBLE.get()) {

0 commit comments

Comments
 (0)