maropu commented on a change in pull request #32488: URL: https://github.com/apache/spark/pull/32488#discussion_r637672796
########## File path: sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/UnwrapCastInBinaryComparison.scala ########## @@ -121,6 +129,49 @@ object UnwrapCastInBinaryComparison extends Rule[LogicalPlan] { if canImplicitlyCast(fromExp, toType, literalType) => simplifyNumericComparison(be, fromExp, toType, value) + // As the analyzer makes sure that the list of In is already of the same data type, then the + // rule can simply check the first literal in `in.list` can implicitly cast to `toType` or not, + // and this rule doesn't convert in when `in.list` is empty. + case in @ In(Cast(fromExp, toType: NumericType, _), list @ Seq(firstLit, _*)) + if canImplicitlyCast(fromExp, toType, firstLit.dataType) && in.inSetConvertible => + val (newValueList, exp) = + list.map(lit => unwrapCast(EqualTo(in.value, lit))) + .partition { + case EqualTo(_, _: Literal) => true + case And(IsNull(_), Literal(null, BooleanType)) => false Review comment: `case And(IsNull(_), Literal(null, BooleanType)) => false` => `case _ => false`? ########## File path: sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/UnwrapCastInBinaryComparison.scala ########## @@ -121,6 +129,49 @@ object UnwrapCastInBinaryComparison extends Rule[LogicalPlan] { if canImplicitlyCast(fromExp, toType, literalType) => simplifyNumericComparison(be, fromExp, toType, value) + // As the analyzer makes sure that the list of In is already of the same data type, then the + // rule can simply check the first literal in `in.list` can implicitly cast to `toType` or not, + // and this rule doesn't convert in when `in.list` is empty. + case in @ In(Cast(fromExp, toType: NumericType, _), list @ Seq(firstLit, _*)) + if canImplicitlyCast(fromExp, toType, firstLit.dataType) && in.inSetConvertible => + val (newValueList, exp) = + list.map(lit => unwrapCast(EqualTo(in.value, lit))) + .partition { + case EqualTo(_, _: Literal) => true + case And(IsNull(_), Literal(null, BooleanType)) => false + } + + val (nonNullValueList, nullValueList) = newValueList.partition { + case EqualTo(_, NonNullLiteral(_, _: NumericType)) => true + case EqualTo(_, Literal(null, _)) => false + } + // make sure the new return list have the same dataType. + val newList = { + if (nonNullValueList.nonEmpty) { + // cast the null value to the dataType of nonNullValueList + // when the nonNullValueList is nonEmpty. + nullValueList.map { + case EqualTo(_, lit) => + Cast(lit, nonNullValueList.head.asInstanceOf[EqualTo].left.dataType) + } ++ nonNullValueList.map {case EqualTo(_, lit) => lit} + } else { + // the new value list only contains null value, + // cast the null value to fromExp.dataType. + nullValueList.map { + case EqualTo(_, lit) => + Cast(lit, fromExp.dataType) + } + } + } + + val unwrapIn = In(fromExp, newList) + // since `exp` are all the same, + // convert to a single value `And(IsNull(_), Literal(null, BooleanType))`. + exp.headOption match { + case Some(falseIfNotNull) => Or(falseIfNotNull, unwrapIn) Review comment: We still need to unwrap casts in this case? IIUC we unwrap casts so that the later optimizer rules can easily push down predicates into data sources. But, predicates having `Or` makes it hard to push them down? ########## File path: sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/UnwrapCastInBinaryComparison.scala ########## @@ -121,6 +129,49 @@ object UnwrapCastInBinaryComparison extends Rule[LogicalPlan] { if canImplicitlyCast(fromExp, toType, literalType) => simplifyNumericComparison(be, fromExp, toType, value) + // As the analyzer makes sure that the list of In is already of the same data type, then the + // rule can simply check the first literal in `in.list` can implicitly cast to `toType` or not, + // and this rule doesn't convert in when `in.list` is empty. + case in @ In(Cast(fromExp, toType: NumericType, _), list @ Seq(firstLit, _*)) + if canImplicitlyCast(fromExp, toType, firstLit.dataType) && in.inSetConvertible => + val (newValueList, exp) = + list.map(lit => unwrapCast(EqualTo(in.value, lit))) + .partition { + case EqualTo(_, _: Literal) => true + case And(IsNull(_), Literal(null, BooleanType)) => false + } + + val (nonNullValueList, nullValueList) = newValueList.partition { + case EqualTo(_, NonNullLiteral(_, _: NumericType)) => true + case EqualTo(_, Literal(null, _)) => false Review comment: ditto: `case EqualTo(_, Literal(null, _)) => false` => `case _ => false`? ########## File path: sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/UnwrapCastInBinaryComparison.scala ########## @@ -21,15 +21,15 @@ import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.expressions.Literal.FalseLiteral import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.catalyst.trees.TreePattern.BINARY_COMPARISON +import org.apache.spark.sql.catalyst.trees.TreePattern.{BINARY_COMPARISON, IN} import org.apache.spark.sql.types._ /** - * Unwrap casts in binary comparison operations with patterns like following: + * Unwrap casts in binary comparison or `In` operations with patterns like following: * - * `BinaryComparison(Cast(fromExp, toType), Literal(value, toType))` - * or - * `BinaryComparison(Literal(value, toType), Cast(fromExp, toType))` + * - `BinaryComparison(Cast(fromExp, toType), Literal(value, toType))` + * - `BinaryComparison(Literal(value, toType), Cast(fromExp, toType))` + * - `In(Cast(fromExp, toType), Seq((v1, toType), (v2, toType), ...)` Review comment: `v1` and `v2` should be `Literal`? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For additional commands, e-mail: reviews-h...@spark.apache.org