RyanBerti commented on code in PR #41203:
URL: https://github.com/apache/spark/pull/41203#discussion_r1198077075
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ExpectsInputTypes.scala:
##########
@@ -74,3 +74,44 @@ object ExpectsInputTypes extends QueryErrorsBase {
trait ImplicitCastInputTypes extends ExpectsInputTypes {
// No other methods
}
+
+/**
+ * An extension of the ExpectsInputTypes trait that also checks each
+ * input expression's foldable attribute against inputIsFoldable;
+ * inputInputFoldable is a Seq of type Option[Boolean], specifying
+ * a value of None bypasses the check for the given column
+ */
+trait ExpectsInputTypesAndFoldable extends ExpectsInputTypes {
+ def inputIsFoldable: Seq[Option[Boolean]]
+
+ override def checkInputDataTypes(): TypeCheckResult = {
+ super.checkInputDataTypes() match {
+ case TypeCheckResult.TypeCheckSuccess =>
+ ExpectsInputTypesAndFoldable.checkInputIsFoldable(children,
inputIsFoldable)
+ case failure => failure
+ }
+ }
+}
+
+object ExpectsInputTypesAndFoldable extends QueryErrorsBase {
+
+ def checkInputIsFoldable(
+ inputs: Seq[Expression],
+ inputIsFoldable: Seq[Option[Boolean]]): TypeCheckResult = {
+ val mismatch = inputs.zip(inputIsFoldable).zipWithIndex.filter {
+ case ((_, Some(_)), _) => true
+ case _ => false
+ }.collectFirst {
+ case ((input, expected), idx) if input.foldable != expected.get =>
+ DataTypeMismatch(
+ errorSubClass = "UNEXPECTED_INPUT_FOLDABLE_VALUE",
+ messageParameters = Map(
+ "paramIndex" -> (idx + 1).toString,
+ "inputSql" -> toSQLExpr(input),
+ "inputFoldable" -> input.foldable.toString,
+ "requiredFoldable" -> expected.get.toString))
+ }
Review Comment:
Thanks for the suggestion - that's a bit cleaner!
--
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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
