CometCast.scala

/*
 * 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.comet.expressions

import org.apache.spark.sql.catalyst.expressions.{Attribute, Cast, Expression, Literal}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.{ArrayType, DataType, DataTypes, DecimalType, NullType, StructType, TimestampNTZType, TimestampType}

import org.apache.comet.CometConf
import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus, withFallbackReason}
import org.apache.comet.serde.{CometExpressionSerde, Compatible, ExprOuterClass, Incompatible, SupportLevel, Unsupported}
import org.apache.comet.serde.ExprOuterClass.Expr
import org.apache.comet.serde.QueryPlanSerde.{evalModeToProto, exprToProtoInternal, serializeDataType}
import org.apache.comet.shims.CometExprShim

object CometCast extends CometExpressionSerde[Cast] with CometExprShim {

  def supportedTypes: Seq[DataType] =
    Seq(
      DataTypes.BooleanType,
      DataTypes.ByteType,
      DataTypes.ShortType,
      DataTypes.IntegerType,
      DataTypes.LongType,
      DataTypes.FloatType,
      DataTypes.DoubleType,
      DataTypes.createDecimalType(10, 2),
      DataTypes.StringType,
      DataTypes.BinaryType,
      DataTypes.DateType,
      DataTypes.TimestampType,
      DataTypes.TimestampNTZType)

  override def getIncompatibleReasons(): Seq[String] = Seq(
    "Some cast operations between specific type pairs may produce different results than Spark." +
      " Refer to the compatibility guide for the full matrix of supported cast operations.")

  override def getUnsupportedReasons(): Seq[String] = Seq(
    "Not all cast type combinations are supported. Unsupported casts fall back to Spark.")

  override def getSupportLevel(cast: Cast): SupportLevel = {
    if (cast.child.isInstanceOf[Literal]) {
      // casting from literal is compatible because we delegate to Spark
      // further data type checks will be performed by CometLiteral
      Compatible()
    } else {
      isSupported(cast.child.dataType, cast.dataType, cast.timeZoneId, evalMode(cast))
    }
  }

  override def convert(
      cast: Cast,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    val cometEvalMode = evalMode(cast)
    cast.child match {
      case _: Literal =>
        exprToProtoInternal(Literal.create(cast.eval(), cast.dataType), inputs, binding)
      case _ =>
        if (isAlwaysCastToNull(cast.child.dataType, cast.dataType, cometEvalMode)) {
          exprToProtoInternal(Literal.create(null, cast.dataType), inputs, binding)
        } else {
          val childExpr = exprToProtoInternal(cast.child, inputs, binding)
          if (childExpr.isDefined) {
            castToProto(cast, cast.timeZoneId, cast.dataType, childExpr.get, cometEvalMode)
          } else {
            withFallbackReason(cast, cast.child)
            None
          }
        }
    }
  }

//  Some casts like date -> int/byte / long are always null. Terminate early in planning
  private def isAlwaysCastToNull(
      fromType: DataType,
      toType: DataType,
      evalMode: CometEvalMode.Value): Boolean = {
    (fromType, toType) match {
      case (
            DataTypes.DateType,
            DataTypes.BooleanType | DataTypes.ByteType | DataTypes.ShortType |
            DataTypes.IntegerType | DataTypes.LongType | DataTypes.FloatType |
            DataTypes.DoubleType | _: DecimalType) if evalMode == CometEvalMode.LEGACY =>
        true
      case _ => false
    }
  }

  /**
   * Wrap an already serialized expression in a cast.
   */
  def castToProto(
      expr: Expression,
      timeZoneId: Option[String],
      dt: DataType,
      childExpr: Expr,
      evalMode: CometEvalMode.Value): Option[Expr] = {
    serializeDataType(dt) match {
      case Some(dataType) =>
        val castBuilder = ExprOuterClass.Cast.newBuilder()
        castBuilder.setChild(childExpr)
        castBuilder.setDatatype(dataType)
        castBuilder.setEvalMode(evalModeToProto(evalMode))
        castBuilder.setAllowIncompat(
          SQLConf.get
            .getConfString(CometConf.getExprAllowIncompatConfigKey(classOf[Cast]), "false")
            .toBoolean)
        castBuilder.setTimezone(timeZoneId.getOrElse("UTC"))
        castBuilder.setIsSpark4Plus(isSpark40Plus)
        Some(
          ExprOuterClass.Expr
            .newBuilder()
            .setCast(castBuilder)
            .build())
      case _ =>
        withFallbackReason(expr, s"Unsupported datatype in castToProto: $dt")
        None
    }
  }

  def isSupported(
      fromType: DataType,
      toType: DataType,
      timeZoneId: Option[String],
      evalMode: CometEvalMode.Value): SupportLevel = {

    if (fromType == toType) {
      return Compatible()
    }

    (fromType, toType) match {
      case (dt: ArrayType, _: ArrayType) if dt.elementType == NullType => Compatible()
      case (ArrayType(DataTypes.DateType, _), ArrayType(toElementType, _))
          if toElementType != DataTypes.IntegerType && toElementType != DataTypes.StringType =>
        unsupported(fromType, toType)
      case (dt: ArrayType, DataTypes.StringType) =>
        isSupported(dt.elementType, DataTypes.StringType, timeZoneId, evalMode)
      case (dt: ArrayType, dt1: ArrayType) =>
        isSupported(dt.elementType, dt1.elementType, timeZoneId, evalMode)
      case (dt: DataType, _) if dt.typeName == "timestamp_ntz" =>
        toType match {
          case DataTypes.StringType => Compatible()
          case DataTypes.DateType => Compatible()
          case DataTypes.TimestampType => Compatible()
          case _ => unsupported(fromType, toType)
        }
      case (_: DecimalType, _: DecimalType) =>
        Compatible()
      case (DataTypes.StringType, _) =>
        canCastFromString(toType, timeZoneId, evalMode)
      case (_, DataTypes.StringType) =>
        canCastToString(fromType, timeZoneId, evalMode)
      case (DataTypes.TimestampType, _) =>
        canCastFromTimestamp(toType)
      case (_: DecimalType, _) =>
        canCastFromDecimal(toType)
      case (DataTypes.BooleanType, _) =>
        canCastFromBoolean(toType, evalMode)
      case (DataTypes.ByteType, _) =>
        canCastFromByte(toType, evalMode)
      case (DataTypes.ShortType, _) =>
        canCastFromShort(toType, evalMode)
      case (DataTypes.IntegerType, _) =>
        canCastFromInt(toType, evalMode)
      case (DataTypes.LongType, _) =>
        canCastFromLong(toType, evalMode)
      case (DataTypes.FloatType, _) =>
        canCastFromFloat(toType)
      case (DataTypes.DoubleType, _) =>
        canCastFromDouble(toType)
      case (from_struct: StructType, to_struct: StructType) =>
        from_struct.fields.zip(to_struct.fields).foreach { case (a, b) =>
          isSupported(a.dataType, b.dataType, timeZoneId, evalMode) match {
            case Compatible(_) =>
            // all good
            case other =>
              return other
          }
        }
        Compatible()
      case (DataTypes.DateType, toType) => canCastFromDate(toType, evalMode)
      case _ => unsupported(fromType, toType)
    }
  }

  private def canCastFromString(
      toType: DataType,
      timeZoneId: Option[String],
      evalMode: CometEvalMode.Value): SupportLevel = {
    toType match {
      case DataTypes.BooleanType =>
        Compatible()
      case DataTypes.ByteType | DataTypes.ShortType | DataTypes.IntegerType |
          DataTypes.LongType =>
        Compatible()
      case DataTypes.BinaryType =>
        Compatible()
      case DataTypes.FloatType | DataTypes.DoubleType =>
        Compatible()
      case _: DecimalType =>
        Compatible()
      case DataTypes.DateType =>
        // https://github.com/apache/datafusion-comet/issues/327
        Compatible(Some("Only supports years between 262143 BC and 262142 AD"))
      case DataTypes.TimestampType =>
        Compatible()
      case _: TimestampNTZType =>
        Compatible()
      case _ =>
        unsupported(DataTypes.StringType, toType)
    }
  }

  private def canCastToString(
      fromType: DataType,
      timeZoneId: Option[String],
      evalMode: CometEvalMode.Value): SupportLevel = {
    fromType match {
      case DataTypes.BooleanType => Compatible()
      case DataTypes.ByteType | DataTypes.ShortType | DataTypes.IntegerType |
          DataTypes.LongType =>
        Compatible()
      case DataTypes.DateType => Compatible()
      case DataTypes.TimestampType => Compatible()
      case DataTypes.FloatType | DataTypes.DoubleType =>
        Compatible(Some("There can be differences in precision"))
      case d: DecimalType if d.scale < 0 =>
        // Negative-scale decimals require spark.sql.legacy.allowNegativeScaleOfDecimal=true.
        // When that config is enabled, Spark formats them using Java BigDecimal.toString()
        // which produces scientific notation (e.g. "1.23E+4"). Comet matches this behavior.
        // When the config is disabled, negative-scale decimals cannot be created in Spark,
        // so we mark this as incompatible to avoid native execution on unexpected inputs.
        val allowNegativeScale = SQLConf.get
          .getConfString("spark.sql.legacy.allowNegativeScaleOfDecimal", "false")
          .toBoolean
        if (allowNegativeScale) Compatible() else Incompatible()
      case _: DecimalType =>
        // Compatible across all eval modes: LEGACY uses cast_decimal128_to_utf8 which
        // replicates Java BigDecimal.toString() (scientific notation when adj_exp < -6);
        // TRY and ANSI fall through to DataFusion's plain-notation cast, which matches Spark.
        Compatible()
      case DataTypes.BinaryType =>
        Compatible()
      case StructType(fields) =>
        for (field <- fields) {
          isSupported(field.dataType, DataTypes.StringType, timeZoneId, evalMode) match {
            case s: Incompatible =>
              return s
            case u: Unsupported =>
              return u
            case _ =>
          }
        }
        Compatible()
      case _ => unsupported(fromType, DataTypes.StringType)
    }
  }

  private def canCastFromTimestamp(toType: DataType): SupportLevel = {
    toType match {
      case DataTypes.BooleanType | DataTypes.ByteType | DataTypes.ShortType |
          DataTypes.IntegerType =>
        // https://github.com/apache/datafusion-comet/issues/352
        // this seems like an edge case that isn't important for us to support
        unsupported(DataTypes.TimestampType, toType)
      case DataTypes.LongType =>
        // https://github.com/apache/datafusion-comet/issues/352
        Compatible()
      case DataTypes.StringType => Compatible()
      case DataTypes.DateType => Compatible()
      case _: TimestampNTZType => Compatible()
      case _ => unsupported(DataTypes.TimestampType, toType)
    }
  }

  private def canCastFromBoolean(toType: DataType, evalMode: CometEvalMode.Value): SupportLevel =
    toType match {
      case DataTypes.ByteType | DataTypes.ShortType | DataTypes.IntegerType | DataTypes.LongType |
          DataTypes.FloatType | DataTypes.DoubleType | _: DecimalType =>
        Compatible()
      case _: TimestampType if evalMode == CometEvalMode.LEGACY =>
        Compatible()
      case _ => unsupported(DataTypes.BooleanType, toType)
    }

  private def canCastFromByte(toType: DataType, evalMode: CometEvalMode.Value): SupportLevel =
    toType match {
      case DataTypes.BooleanType =>
        Compatible()
      case DataTypes.ShortType | DataTypes.IntegerType | DataTypes.LongType =>
        Compatible()
      case DataTypes.FloatType | DataTypes.DoubleType | _: DecimalType =>
        Compatible()
      case DataTypes.BinaryType if (evalMode == CometEvalMode.LEGACY) =>
        Compatible()
      case DataTypes.TimestampType =>
        Compatible()
      case _ =>
        unsupported(DataTypes.ByteType, toType)
    }

  private def canCastFromShort(toType: DataType, evalMode: CometEvalMode.Value): SupportLevel =
    toType match {
      case DataTypes.BooleanType =>
        Compatible()
      case DataTypes.ByteType | DataTypes.IntegerType | DataTypes.LongType =>
        Compatible()
      case DataTypes.FloatType | DataTypes.DoubleType | _: DecimalType =>
        Compatible()
      case DataTypes.BinaryType if (evalMode == CometEvalMode.LEGACY) =>
        Compatible()
      case DataTypes.TimestampType =>
        Compatible()
      case _ =>
        unsupported(DataTypes.ShortType, toType)
    }

  private def canCastFromInt(toType: DataType, evalMode: CometEvalMode.Value): SupportLevel =
    toType match {
      case DataTypes.BooleanType =>
        Compatible()
      case DataTypes.ByteType | DataTypes.ShortType | DataTypes.LongType =>
        Compatible()
      case DataTypes.FloatType | DataTypes.DoubleType =>
        Compatible()
      case _: DecimalType =>
        Compatible()
      case DataTypes.BinaryType if (evalMode == CometEvalMode.LEGACY) => Compatible()
      case DataTypes.TimestampType =>
        Compatible()
      case _ =>
        unsupported(DataTypes.IntegerType, toType)
    }

  private def canCastFromLong(toType: DataType, evalMode: CometEvalMode.Value): SupportLevel =
    toType match {
      case DataTypes.BooleanType =>
        Compatible()
      case DataTypes.ByteType | DataTypes.ShortType | DataTypes.IntegerType =>
        Compatible()
      case DataTypes.FloatType | DataTypes.DoubleType =>
        Compatible()
      case _: DecimalType =>
        Compatible()
      case DataTypes.BinaryType if (evalMode == CometEvalMode.LEGACY) => Compatible()
      case DataTypes.TimestampType =>
        Compatible()
      case _ =>
        unsupported(DataTypes.LongType, toType)
    }

  private def canCastFromFloat(toType: DataType): SupportLevel = toType match {
    case DataTypes.BooleanType | DataTypes.DoubleType | DataTypes.ByteType | DataTypes.ShortType |
        DataTypes.IntegerType | DataTypes.LongType | DataTypes.TimestampType =>
      Compatible()
    case _: DecimalType =>
      // https://github.com/apache/datafusion-comet/issues/1371
      Incompatible(Some("There can be rounding differences"))
    case _ =>
      unsupported(DataTypes.FloatType, toType)
  }

  private def canCastFromDouble(toType: DataType): SupportLevel = toType match {
    case DataTypes.BooleanType | DataTypes.FloatType | DataTypes.ByteType | DataTypes.ShortType |
        DataTypes.IntegerType | DataTypes.LongType | DataTypes.TimestampType =>
      Compatible()
    case _: DecimalType =>
      // https://github.com/apache/datafusion-comet/issues/1371
      Incompatible(Some("There can be rounding differences"))
    case _ => unsupported(DataTypes.DoubleType, toType)
  }

  private def canCastFromDecimal(toType: DataType): SupportLevel = toType match {
    case DataTypes.FloatType | DataTypes.DoubleType | DataTypes.ByteType | DataTypes.ShortType |
        DataTypes.IntegerType | DataTypes.LongType | DataTypes.BooleanType |
        DataTypes.TimestampType =>
      Compatible()
    case _ => Unsupported(Some(s"Cast from DecimalType to $toType is not supported"))
  }

  private def canCastFromDate(toType: DataType, evalMode: CometEvalMode.Value): SupportLevel =
    toType match {
      case DataTypes.TimestampType =>
        Compatible()
      case _: TimestampNTZType =>
        Compatible()
      case DataTypes.BooleanType | DataTypes.ByteType | DataTypes.ShortType |
          DataTypes.IntegerType | DataTypes.LongType | DataTypes.FloatType |
          DataTypes.DoubleType | _: DecimalType if evalMode == CometEvalMode.LEGACY =>
        Compatible()
      case _ => Unsupported(Some(s"Cast from DateType to $toType is not supported"))
    }

  private def unsupported(fromType: DataType, toType: DataType): Unsupported = {
    Unsupported(Some(s"Cast from $fromType to $toType is not supported"))
  }
}