structs.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.serde

import scala.jdk.CollectionConverters._
import scala.util.Try

import org.apache.spark.sql.catalyst.expressions.{Attribute, CreateNamedStruct, GetArrayStructFields, GetStructField, JsonToStructs, StructsToCsv, StructsToJson}
import org.apache.spark.sql.types._

import org.apache.comet.CometSparkSessionExtensions.withInfo
import org.apache.comet.DataTypeSupport
import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, serializeDataType}

object CometCreateNamedStruct extends CometExpressionSerde[CreateNamedStruct] {
  override def convert(
      expr: CreateNamedStruct,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    if (expr.names.length != expr.names.distinct.length) {
      withInfo(expr, "CreateNamedStruct with duplicate field names are not supported")
      return None
    }

    val valExprs = expr.valExprs.map(exprToProtoInternal(_, inputs, binding))

    if (valExprs.forall(_.isDefined)) {
      val structBuilder = ExprOuterClass.CreateNamedStruct.newBuilder()
      structBuilder.addAllValues(valExprs.map(_.get).asJava)
      structBuilder.addAllNames(expr.names.map(_.toString).asJava)

      Some(
        ExprOuterClass.Expr
          .newBuilder()
          .setCreateNamedStruct(structBuilder)
          .build())
    } else {
      withInfo(expr, "unsupported arguments for CreateNamedStruct", expr.valExprs: _*)
      None
    }

  }
}

object CometGetStructField extends CometExpressionSerde[GetStructField] {
  override def convert(
      expr: GetStructField,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    exprToProtoInternal(expr.child, inputs, binding).map { childExpr =>
      val getStructFieldBuilder = ExprOuterClass.GetStructField
        .newBuilder()
        .setChild(childExpr)
        .setOrdinal(expr.ordinal)

      ExprOuterClass.Expr
        .newBuilder()
        .setGetStructField(getStructFieldBuilder)
        .build()
    }
  }
}

object CometGetArrayStructFields extends CometExpressionSerde[GetArrayStructFields] {
  override def convert(
      expr: GetArrayStructFields,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    val childExpr = exprToProtoInternal(expr.child, inputs, binding)

    if (childExpr.isDefined) {
      val arrayStructFieldsBuilder = ExprOuterClass.GetArrayStructFields
        .newBuilder()
        .setChild(childExpr.get)
        .setOrdinal(expr.ordinal)

      Some(
        ExprOuterClass.Expr
          .newBuilder()
          .setGetArrayStructFields(arrayStructFieldsBuilder)
          .build())
    } else {
      withInfo(expr, "unsupported arguments for GetArrayStructFields", expr.child)
      None
    }
  }
}

object CometStructsToJson extends CometExpressionSerde[StructsToJson] {

  override def getSupportLevel(expr: StructsToJson): SupportLevel =
    Incompatible(
      Some(
        "Does not support Infinity/-Infinity for numeric types" +
          " (https://github.com/apache/datafusion-comet/issues/3016)"))

  override def convert(
      expr: StructsToJson,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    if (expr.options.nonEmpty) {
      withInfo(expr, "StructsToJson with options is not supported")
      None
    } else {
      val isSupported = expr.child.dataType match {
        case s: StructType =>
          s.fields.forall(f => isSupportedType(f.dataType))
        case _: MapType | _: ArrayType =>
          // Spark supports map and array in StructsToJson but this is not yet
          // implemented in Comet
          false
        case _ =>
          false
      }

      if (isSupported) {
        exprToProtoInternal(expr.child, inputs, binding) match {
          case Some(p) =>
            val toJson = ExprOuterClass.ToJson
              .newBuilder()
              .setChild(p)
              .setTimezone(expr.timeZoneId.getOrElse("UTC"))
              .setIgnoreNullFields(true)
              .build()
            Some(
              ExprOuterClass.Expr
                .newBuilder()
                .setToJson(toJson)
                .build())
          case _ =>
            withInfo(expr, expr.child)
            None
        }
      } else {
        withInfo(expr, "Unsupported data type", expr.child)
        None
      }
    }
  }

  def isSupportedType(dt: DataType): Boolean = {
    dt match {
      case StructType(fields) =>
        fields.forall(f => isSupportedType(f.dataType))
      case DataTypes.BooleanType | DataTypes.ByteType | DataTypes.ShortType |
          DataTypes.IntegerType | DataTypes.LongType | DataTypes.FloatType |
          DataTypes.DoubleType | DataTypes.StringType =>
        true
      case DataTypes.DateType | DataTypes.TimestampType =>
        // TODO implement these types with tests for formatting options and timezone
        false
      case _: MapType | _: ArrayType =>
        // Spark supports map and array in StructsToJson but this is not yet
        // implemented in Comet
        false
      case _ => false
    }
  }
}

object CometJsonToStructs extends CometExpressionSerde[JsonToStructs] {

  override def getSupportLevel(expr: JsonToStructs): SupportLevel = {
    // this feature is partially implemented and not comprehensively tested yet
    Incompatible()
  }

  override def convert(
      expr: JsonToStructs,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {

    if (expr.schema == null) {
      withInfo(expr, "from_json requires explicit schema")
      return None
    }

    def isSupportedType(dt: DataType): Boolean = {
      dt match {
        case StructType(fields) =>
          fields.nonEmpty && fields.forall(f => isSupportedType(f.dataType))
        case DataTypes.IntegerType | DataTypes.LongType | DataTypes.FloatType |
            DataTypes.DoubleType | DataTypes.BooleanType | DataTypes.StringType =>
          true
        case _ => false
      }
    }

    val schemaType = expr.schema
    if (!isSupportedType(schemaType)) {
      withInfo(expr, "from_json: Unsupported schema type")
      return None
    }

    val options = expr.options
    if (options.nonEmpty) {
      val mode = options.getOrElse("mode", "PERMISSIVE")
      if (mode != "PERMISSIVE") {
        withInfo(expr, s"from_json: Only PERMISSIVE mode supported, got: $mode")
        return None
      }
      val knownOptions = Set("mode")
      val unknownOpts = options.keySet -- knownOptions
      if (unknownOpts.nonEmpty) {
        withInfo(expr, s"from_json: Ignoring unsupported options: ${unknownOpts.mkString(", ")}")
      }
    }

    // Convert child expression and schema to protobuf
    for {
      childProto <- exprToProtoInternal(expr.child, inputs, binding)
      schemaProto <- serializeDataType(schemaType)
    } yield {
      val fromJson = ExprOuterClass.FromJson
        .newBuilder()
        .setChild(childProto)
        .setSchema(schemaProto)
        .setTimezone(expr.timeZoneId.getOrElse("UTC"))
        .build()
      ExprOuterClass.Expr.newBuilder().setFromJson(fromJson).build()
    }
  }
}

object CometStructsToCsv extends CometExpressionSerde[StructsToCsv] {

  private val incompatibleDataTypes = Seq(DateType, TimestampType, TimestampNTZType, BinaryType)

  override def getSupportLevel(expr: StructsToCsv): SupportLevel = {
    val dataTypes = expr.inputSchema.fields.map(_.dataType)
    val containsComplexType = dataTypes.exists(DataTypeSupport.isComplexType)
    if (containsComplexType) {
      return Unsupported(
        s"The schema ${expr.inputSchema} is not supported because it includes a complex type")
    }
    val containsIncompatibleDataTypes = dataTypes.exists(incompatibleDataTypes.contains)
    if (containsIncompatibleDataTypes) {
      return Incompatible(
        Some(
          s"The schema ${expr.inputSchema} is not supported because " +
            s"it includes a incompatible data types: $incompatibleDataTypes"))
    }
    // https://github.com/apache/datafusion-comet/issues/3232
    Incompatible()
  }

  override def convert(
      expr: StructsToCsv,
      inputs: Seq[Attribute],
      binding: Boolean): Option[ExprOuterClass.Expr] = {
    for {
      childProto <- exprToProtoInternal(expr.child, inputs, binding)
    } yield {
      val optionsProto = options2Proto(expr.options, expr.timeZoneId)
      val toCsv = ExprOuterClass.ToCsv
        .newBuilder()
        .setChild(childProto)
        .setOptions(optionsProto)
        .build()
      ExprOuterClass.Expr.newBuilder().setToCsv(toCsv).build()
    }
  }

  private def options2Proto(
      options: Map[String, String],
      timeZoneId: Option[String]): ExprOuterClass.CsvWriteOptions = {
    ExprOuterClass.CsvWriteOptions
      .newBuilder()
      .setDelimiter(options.getOrElse("delimiter", ","))
      .setQuote(options.getOrElse("quote", "\""))
      .setEscape(options.getOrElse("escape", "\\"))
      .setNullValue(options.getOrElse("nullValue", ""))
      .setTimezone(timeZoneId.getOrElse("UTC"))
      .setIgnoreLeadingWhiteSpace(options
        .get("ignoreLeadingWhiteSpace")
        .flatMap(ignoreLeadingWhiteSpace => Try(ignoreLeadingWhiteSpace.toBoolean).toOption)
        .getOrElse(true))
      .setIgnoreTrailingWhiteSpace(options
        .get("ignoreTrailingWhiteSpace")
        .flatMap(ignoreTrailingWhiteSpace => Try(ignoreTrailingWhiteSpace.toBoolean).toOption)
        .getOrElse(true))
      .setQuoteAll(options
        .get("quoteAll")
        .flatMap(quoteAll => Try(quoteAll.toBoolean).toOption)
        .getOrElse(false))
      .build()
  }
}