[SPARK-29837][SQL] PostgreSQL dialect: cast to boolean

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/PostgreSQLDialect.scala

@@ -19,27 +19,26 @@ package org.apache.spark.sql.catalyst.analysis
 
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.catalyst.expressions.Cast
-import org.apache.spark.sql.catalyst.expressions.postgreSQL.PostgreCastStringToBoolean
+import org.apache.spark.sql.catalyst.expressions.postgreSQL.PostgreCastToBoolean
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.{BooleanType, StringType}
 
 object PostgreSQLDialect {
   val postgreSQLDialectRules: List[Rule[LogicalPlan]] =
-    CastStringToBoolean ::
+    CastToBoolean ::
       Nil
 
-  object CastStringToBoolean extends Rule[LogicalPlan] with Logging {
+  object CastToBoolean extends Rule[LogicalPlan] with Logging {
     override def apply(plan: LogicalPlan): LogicalPlan = {
       // The SQL configuration `spark.sql.dialect` can be changed in runtime.
       // To make sure the configuration is effective, we have to check it during rule execution.
       val conf = SQLConf.get
       if (conf.usePostgreSQLDialect) {
         plan.transformExpressions {
-          case Cast(child, dataType, _)
-            if dataType == BooleanType && child.dataType == StringType =>
-            PostgreCastStringToBoolean(child)
+          case Cast(child, dataType, timeZoneId) if dataType == BooleanType =>
+            PostgreCastToBoolean(child, timeZoneId)
         }
       } else {
         plan

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Cast.scala

@@ -273,7 +273,7 @@ abstract class CastBase extends UnaryExpression with TimeZoneAwareExpression wit
   private[this] def needsTimeZone: Boolean = Cast.needsTimeZone(child.dataType, dataType)
 
   // [[func]] assumes the input is no longer null because eval already does the null check.
-  @inline private[this] def buildCast[T](a: Any, func: T => Any): Any = func(a.asInstanceOf[T])
+  @inline protected[this] def buildCast[T](a: Any, func: T => Any): Any = func(a.asInstanceOf[T])
 
   private lazy val dateFormatter = DateFormatter(zoneId)
   private lazy val timestampFormatter = TimestampFormatter.getFractionFormatter(zoneId)
@@ -376,7 +376,7 @@ abstract class CastBase extends UnaryExpression with TimeZoneAwareExpression wit
   }
 
   // UDFToBoolean
-  private[this] def castToBoolean(from: DataType): Any => Any = from match {
+  protected[this] def castToBoolean(from: DataType): Any => Any = from match {
     case StringType =>
       buildCast[UTF8String](_, s => {
         if (StringUtils.isTrueString(s)) {
@@ -781,7 +781,7 @@ abstract class CastBase extends UnaryExpression with TimeZoneAwareExpression wit
     }
   }
 
-  override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+  override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
     val eval = child.genCode(ctx)
     val nullSafeCast = nullSafeCastFunction(child.dataType, dataType, ctx)
 
@@ -791,7 +791,7 @@ abstract class CastBase extends UnaryExpression with TimeZoneAwareExpression wit
 
   // The function arguments are: `input`, `result` and `resultIsNull`. We don't need `inputIsNull`
   // in parameter list, because the returned code will be put in null safe evaluation region.
-  private[this] type CastFunction = (ExprValue, ExprValue, ExprValue) => Block
+  protected[this] type CastFunction = (ExprValue, ExprValue, ExprValue) => Block
 
   private[this] def nullSafeCastFunction(
       from: DataType,
@@ -1233,7 +1233,7 @@ abstract class CastBase extends UnaryExpression with TimeZoneAwareExpression wit
   private[this] def timestampToDoubleCode(ts: ExprValue): Block =
     code"$ts / (double)$MICROS_PER_SECOND"
 
-  private[this] def castToBooleanCode(from: DataType): CastFunction = from match {
+  protected[this] def castToBooleanCode(from: DataType): CastFunction = from match {
     case StringType =>
       val stringUtils = inline"${StringUtils.getClass.getName.stripSuffix("$")}"
       (c, evPrim, evNull) =>

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/postgreSQL/PostgreCastToBoolean.scala

@@ -0,0 +1,81 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.catalyst.expressions.postgreSQL
+
+import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
+import org.apache.spark.sql.catalyst.expressions.{CastBase, Expression, TimeZoneAwareExpression}
+import org.apache.spark.sql.catalyst.expressions.codegen.Block._
+import org.apache.spark.sql.catalyst.util.postgreSQL.StringUtils
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types._
+import org.apache.spark.unsafe.types.UTF8String
+
+case class PostgreCastToBoolean(child: Expression, timeZoneId: Option[String])
+  extends CastBase {
+
+  override protected def ansiEnabled = SQLConf.get.ansiEnabled
+
+  override def withTimeZone(timeZoneId: String): TimeZoneAwareExpression =
+    copy(timeZoneId = Option(timeZoneId))
+
+  override def checkInputDataTypes(): TypeCheckResult = child.dataType match {
+    case StringType | IntegerType => TypeCheckResult.TypeCheckSuccess
+    case dt => throw new UnsupportedOperationException(s"cannot cast type $dt to boolean")
+  }
+
+  override def castToBoolean(from: DataType): Any => Any = from match {
+    case StringType =>
+      buildCast[UTF8String](_, str => {
+        val s = str.trim().toLowerCase()
+        if (StringUtils.isTrueString(s)) {
+          true
+        } else if (StringUtils.isFalseString(s)) {
+          false
+        } else {
+          throw new IllegalArgumentException(s"invalid input syntax for type boolean: $s")
+        }
+      })
+    case IntegerType =>
+      super.castToBoolean(from)
+  }
+
+  override def castToBooleanCode(from: DataType): CastFunction = from match {
+    case StringType =>
+      val stringUtils = inline"${StringUtils.getClass.getName.stripSuffix("$")}"
+      (c, evPrim, evNull) =>
+        code"""
+          if ($stringUtils.isTrueString($c.trim().toLowerCase())) {
+            $evPrim = true;
+          } else if ($stringUtils.isFalseString($c.trim().toLowerCase())) {
+            $evPrim = false;
+          } else {
+            throw new IllegalArgumentException("invalid input syntax for type boolean: $c");
+          }
+        """
+
+    case IntegerType =>
+      super.castToBooleanCode(from)
+  }
+
+  override def dataType: DataType = BooleanType
+
+  override def nullable: Boolean = child.nullable
+
+  override def toString: String = s"PostgreCastToBoolean($child as ${dataType.simpleString})"
+
+  override def sql: String = s"CAST(${child.sql} AS ${dataType.sql})"
+}

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/postgreSQL/CastSuite.scala

@@ -16,44 +16,57 @@
  */
 package org.apache.spark.sql.catalyst.expressions.postgreSQL
 
+import java.sql.{Date, Timestamp}
+
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql.catalyst.expressions.{ExpressionEvalHelper, Literal}
 
 class CastSuite extends SparkFunSuite with ExpressionEvalHelper {
-  private def checkPostgreCastStringToBoolean(v: Any, expected: Any): Unit = {
-    checkEvaluation(PostgreCastStringToBoolean(Literal(v)), expected)
+  private def checkPostgreCastToBoolean(v: Any, expected: Any): Unit = {
+    checkEvaluation(PostgreCastToBoolean(Literal(v), None), expected)
   }
 
   test("cast string to boolean") {
-    checkPostgreCastStringToBoolean("true", true)
-    checkPostgreCastStringToBoolean("tru", true)
-    checkPostgreCastStringToBoolean("tr", true)
-    checkPostgreCastStringToBoolean("t", true)
-    checkPostgreCastStringToBoolean("tRUe", true)
-    checkPostgreCastStringToBoolean("    tRue   ", true)
-    checkPostgreCastStringToBoolean("    tRu   ", true)
-    checkPostgreCastStringToBoolean("yes", true)
-    checkPostgreCastStringToBoolean("ye", true)
-    checkPostgreCastStringToBoolean("y", true)
-    checkPostgreCastStringToBoolean("1", true)
-    checkPostgreCastStringToBoolean("on", true)
+    checkPostgreCastToBoolean("true", true)
+    checkPostgreCastToBoolean("tru", true)
+    checkPostgreCastToBoolean("tr", true)
+    checkPostgreCastToBoolean("t", true)
+    checkPostgreCastToBoolean("tRUe", true)
+    checkPostgreCastToBoolean("    tRue   ", true)
+    checkPostgreCastToBoolean("    tRu   ", true)
+    checkPostgreCastToBoolean("yes", true)
+    checkPostgreCastToBoolean("ye", true)
+    checkPostgreCastToBoolean("y", true)
+    checkPostgreCastToBoolean("1", true)
+    checkPostgreCastToBoolean("on", true)
+
+    checkPostgreCastToBoolean("false", false)
+    checkPostgreCastToBoolean("fals", false)
+    checkPostgreCastToBoolean("fal", false)
+    checkPostgreCastToBoolean("fa", false)
+    checkPostgreCastToBoolean("f", false)
+    checkPostgreCastToBoolean("    fAlse    ", false)
+    checkPostgreCastToBoolean("    fAls    ", false)
+    checkPostgreCastToBoolean("    FAlsE    ", false)
+    checkPostgreCastToBoolean("no", false)
+    checkPostgreCastToBoolean("n", false)
+    checkPostgreCastToBoolean("0", false)
+    checkPostgreCastToBoolean("off", false)
+    checkPostgreCastToBoolean("of", false)
 
-    checkPostgreCastStringToBoolean("false", false)
-    checkPostgreCastStringToBoolean("fals", false)
-    checkPostgreCastStringToBoolean("fal", false)
-    checkPostgreCastStringToBoolean("fa", false)
-    checkPostgreCastStringToBoolean("f", false)
-    checkPostgreCastStringToBoolean("    fAlse    ", false)
-    checkPostgreCastStringToBoolean("    fAls    ", false)
-    checkPostgreCastStringToBoolean("    FAlsE    ", false)
-    checkPostgreCastStringToBoolean("no", false)
-    checkPostgreCastStringToBoolean("n", false)
-    checkPostgreCastStringToBoolean("0", false)
-    checkPostgreCastStringToBoolean("off", false)
-    checkPostgreCastStringToBoolean("of", false)
+    intercept[Exception](checkPostgreCastToBoolean("o", null))
+    intercept[Exception](checkPostgreCastToBoolean("abc", null))
+    intercept[Exception](checkPostgreCastToBoolean("", null))
+  }
 
-    checkPostgreCastStringToBoolean("o", null)
-    checkPostgreCastStringToBoolean("abc", null)
-    checkPostgreCastStringToBoolean("", null)
+  test("unsupported data types to cast to boolean") {
+    intercept[Exception](checkPostgreCastToBoolean(new Timestamp(1), null))
+    intercept[Exception](checkPostgreCastToBoolean(new Date(1), null))
+    intercept[Exception](checkPostgreCastToBoolean(1.toLong, null))
+    intercept[Exception](checkPostgreCastToBoolean(1.toShort, null))
+    intercept[Exception](checkPostgreCastToBoolean(1.toByte, null))
+    intercept[Exception](checkPostgreCastToBoolean(BigDecimal(1.0), null))
+    intercept[Exception](checkPostgreCastToBoolean(1.toDouble, null))
+    intercept[Exception](checkPostgreCastToBoolean(1.toFloat, null))
   }
 }