/*
 * 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.execution.datasources.parquet

import org.apache.hudi.common.util.{Option => HOption}
import org.apache.hudi.internal.schema.InternalSchema
import org.apache.hudi.io.storage.HoodieSparkParquetReader.ENABLE_LOGICAL_TIMESTAMP_REPAIR

import org.apache.hadoop.conf.Configuration
import org.apache.hadoop.mapred.FileSplit
import org.apache.hadoop.mapreduce.{JobID, TaskAttemptID, TaskID, TaskType}
import org.apache.hadoop.mapreduce.task.TaskAttemptContextImpl
import org.apache.parquet.filter2.compat.FilterCompat
import org.apache.parquet.filter2.predicate.{FilterApi, FilterPredicate}
import org.apache.parquet.hadoop.{ParquetInputFormat, ParquetRecordReader}
import org.apache.parquet.hadoop.metadata.{FileMetaData, ParquetMetadata}
import org.apache.parquet.schema.{MessageType, SchemaRepair}
import org.apache.spark.TaskContext
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.JoinedRow
import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes
import org.apache.spark.sql.catalyst.util.{DateTimeUtils, RebaseDateTime}
import org.apache.spark.sql.execution.datasources.{DataSourceUtils, FileFormat, PartitionedFile, RecordReaderIterator, SparkColumnarFileReader}
import org.apache.spark.sql.execution.datasources.parquet.Spark41ParquetReader.repairFooterSchema
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.sources._
import org.apache.spark.sql.types.StructType
import org.apache.spark.util.Utils

import java.io.Closeable
import java.time.ZoneId
import java.util.concurrent.atomic.AtomicBoolean

class Spark41ParquetReader(enableVectorizedReader: Boolean,
                           datetimeRebaseModeInRead: String,
                           int96RebaseModeInRead: String,
                           enableParquetFilterPushDown: Boolean,
                           pushDownDate: Boolean,
                           pushDownTimestamp: Boolean,
                           pushDownDecimal: Boolean,
                           pushDownInFilterThreshold: Int,
                           pushDownStringPredicate: Boolean,
                           isCaseSensitive: Boolean,
                           timestampConversion: Boolean,
                           enableOffHeapColumnVector: Boolean,
                           capacity: Int,
                           returningBatch: Boolean,
                           enableRecordFilter: Boolean,
                           enableLogicalTimestampRepair: Boolean,
                           timeZoneId: Option[String]) extends SparkParquetReaderBase(
  enableVectorizedReader = enableVectorizedReader,
  enableParquetFilterPushDown = enableParquetFilterPushDown,
  pushDownDate = pushDownDate,
  pushDownTimestamp = pushDownTimestamp,
  pushDownDecimal = pushDownDecimal,
  pushDownInFilterThreshold = pushDownInFilterThreshold,
  isCaseSensitive = isCaseSensitive,
  timestampConversion = timestampConversion,
  enableOffHeapColumnVector = enableOffHeapColumnVector,
  capacity = capacity,
  returningBatch = returningBatch,
  enableRecordFilter = enableRecordFilter,
  enableLogicalTimestampRepair = enableLogicalTimestampRepair,
  timeZoneId = timeZoneId) {

  /**
   * Read an individual parquet file
   * Code from ParquetFileFormat#buildReaderWithPartitionValues from Spark v3.5.1 adapted here
   *
   * @param file               parquet file to read
   * @param requiredSchema     desired output schema of the data
   * @param partitionSchema    schema of the partition columns. Partition values will be appended to the end of every row
   * @param internalSchemaOpt  option of internal schema for schema.on.read
   * @param filters            filters for data skipping. Not guaranteed to be used; the spark plan will also apply the filters.
   * @param sharedConf         the hadoop conf
   * @return iterator of rows read from the file output type says [[InternalRow]] but could be [[ColumnarBatch]]
   */
  override protected def doRead(file: PartitionedFile,
                                requiredSchema: StructType,
                                partitionSchema: StructType,
                                internalSchemaOpt: HOption[InternalSchema],
                                filters: scala.Seq[Filter],
                                sharedConf: Configuration,
                                tableSchemaOpt: HOption[MessageType]): Iterator[InternalRow] = {
    assert(file.partitionValues.numFields == partitionSchema.size)

    val filePath = file.toPath
    val split = new FileSplit(filePath, file.start, file.length, Array.empty[String])

    val schemaEvolutionUtils = new ParquetSchemaEvolutionUtils(sharedConf, filePath, requiredSchema,
      partitionSchema, internalSchemaOpt)

    // When there are vectorized reads, we can avoid
    // 1. opening the file twice by transfering the SeekableInputStream
    // 2. reading the footer twice by reading all row groups in advance and filter row groups
    //    according to filters that require push down
    val originalFooter =
      ParquetFooterReader.openFileAndReadFooter(sharedConf, file, enableVectorizedReader)
    val openedFooter = if (enableLogicalTimestampRepair) {
      repairFooterSchema(originalFooter, tableSchemaOpt)
    } else {
      originalFooter
    }
    assert(openedFooter.inputStreamOpt.isPresent == enableVectorizedReader)

    // Before transferring the ownership of inputStream to the vectorizedReader,
    // we must take responsibility to close the inputStream if something goes wrong
    // to avoid resource leak.
    val shouldCloseInputStream = new AtomicBoolean(openedFooter.inputStreamOpt.isPresent)
    try {
      val footerFileMetaData = openedFooter.footer.getFileMetaData
      val datetimeRebaseSpec = DataSourceUtils.datetimeRebaseSpec(
        footerFileMetaData.getKeyValueMetaData.get,
        datetimeRebaseModeInRead)
      val int96RebaseSpec = DataSourceUtils.int96RebaseSpec(
        footerFileMetaData.getKeyValueMetaData.get, int96RebaseModeInRead)

      // Try to push down filters when filter push-down is enabled.
      val pushed = if (enableParquetFilterPushDown) {
        val parquetSchema = footerFileMetaData.getSchema
        val parquetFilters = new ParquetFilters(
          parquetSchema,
          pushDownDate,
          pushDownTimestamp,
          pushDownDecimal,
          pushDownStringPredicate,
          pushDownInFilterThreshold,
          isCaseSensitive,
          datetimeRebaseSpec)
        filters.map(schemaEvolutionUtils.rebuildFilterFromParquet)
          // Collects all converted Parquet filter predicates. Notice that not all predicates
          // can be converted (`ParquetFilters.createFilter` returns an `Option`). That's why
          // a `flatMap` is used here.
          .flatMap(parquetFilters.createFilter)
          .reduceOption(FilterApi.and)
      } else {
        None
      }

      // PARQUET_INT96_TIMESTAMP_CONVERSION says to apply timezone conversions to int96
      // timestamps *only* if the file was created by something other than "parquet-mr",
      // so check the actual writer here for this file.  We have to do this per-file,
      // as each file in the table may have different writers.
      // Define isCreatedByParquetMr as function to avoid unnecessary parquet footer reads.
      def isCreatedByParquetMr: Boolean =
        footerFileMetaData.getCreatedBy.startsWith("parquet-mr")

      val convertTz =
        if (timestampConversion && !isCreatedByParquetMr) {
          Some(DateTimeUtils.getZoneId(sharedConf.get(SQLConf.SESSION_LOCAL_TIMEZONE.key)))
        } else {
          None
        }

      val attemptId = new TaskAttemptID(new TaskID(new JobID(), TaskType.MAP, 0), 0)
      val hadoopAttemptContext =
        new TaskAttemptContextImpl(schemaEvolutionUtils.getHadoopConfClone(footerFileMetaData, enableVectorizedReader), attemptId)

      // Try to push down filters when filter push-down is enabled.
      // Notice: This push-down is RowGroups level, not individual records.
      pushed.foreach {
        ParquetInputFormat.setFilterPredicate(hadoopAttemptContext.getConfiguration, _)
      }
      if (enableVectorizedReader) {
        buildVectorizedIterator(
          hadoopAttemptContext, split, file.partitionValues, partitionSchema, convertTz,
          datetimeRebaseSpec, int96RebaseSpec, enableOffHeapColumnVector, returningBatch,
          capacity, openedFooter, shouldCloseInputStream, schemaEvolutionUtils)
      } else {
        buildRowBasedIterator(
          hadoopAttemptContext, split, file.partitionValues, partitionSchema, convertTz,
          datetimeRebaseSpec, int96RebaseSpec, requiredSchema, pushed, enableRecordFilter,
          enableLogicalTimestampRepair, tableSchemaOpt, schemaEvolutionUtils)
      }
    } finally {
      if (shouldCloseInputStream.get) {
        openedFooter.inputStreamOpt.ifPresent(Utils.closeQuietly)
      }
    }
  }

  // scalastyle:off parameter.number
  private def buildVectorizedIterator(
                                       hadoopAttemptContext: TaskAttemptContextImpl,
                                       split: FileSplit,
                                       partitionValues: InternalRow,
                                       partitionSchema: StructType,
                                       convertTz: Option[ZoneId],
                                       datetimeRebaseSpec: RebaseDateTime.RebaseSpec,
                                       int96RebaseSpec: RebaseDateTime.RebaseSpec,
                                       enableOffHeapColumnVector: Boolean,
                                       returningBatch: Boolean,
                                       batchSize: Int,
                                       openedFooter: OpenedParquetFooter,
                                       shouldCloseInputStream: AtomicBoolean,
                                       schemaEvolutionUtils: ParquetSchemaEvolutionUtils): Iterator[InternalRow] = {
    // scalastyle:on parameter.number
    assert(openedFooter.inputStreamOpt.isPresent)
    val vectorizedReader = schemaEvolutionUtils.buildVectorizedReader(
      convertTz.orNull,
      datetimeRebaseSpec.mode.toString,
      datetimeRebaseSpec.timeZone,
      int96RebaseSpec.mode.toString,
      int96RebaseSpec.timeZone,
      enableOffHeapColumnVector && TaskContext.get() != null,
      batchSize)
    // SPARK-37089: We cannot register a task completion listener to close this iterator here
    // because downstream exec nodes have already registered their listeners. Since listeners
    // are executed in reverse order of registration, a listener registered here would close the
    // iterator while downstream exec nodes are still running. When off-heap column vectors are
    // enabled, this can cause a use-after-free bug leading to a segfault.
    //
    // Instead, we use FileScanRDD's task completion listener to close this iterator.
    val iter = new RecordReaderIterator(vectorizedReader)
    try {
      vectorizedReader.initialize(
        split, hadoopAttemptContext, Some(openedFooter.inputFile),
        Some(openedFooter.inputStream), Some(openedFooter.footer))
      // The caller don't need to take care of the close of inputStream after calling
      // `initialize` because the ownership of inputStream has been transferred to the
      // vectorizedReader
      shouldCloseInputStream.set(false)
      vectorizedReader.initBatch(partitionSchema, partitionValues)
      if (returningBatch) {
        vectorizedReader.enableReturningBatches()
      }

      // UnsafeRowParquetRecordReader appends the columns internally to avoid another copy.
      iter.asInstanceOf[Iterator[InternalRow]]
    } catch {
      case e: Throwable =>
        // SPARK-23457: In case there is an exception in initialization, close the iterator to
        // avoid leaking resources.
        iter.close()
        throw e
    }
  }

  // scalastyle:off parameter.number
  private def buildRowBasedIterator(
                                     hadoopAttemptContext: TaskAttemptContextImpl,
                                     split: FileSplit,
                                     partitionValues: InternalRow,
                                     partitionSchema: StructType,
                                     convertTz: Option[ZoneId],
                                     datetimeRebaseSpec: RebaseDateTime.RebaseSpec,
                                     int96RebaseSpec: RebaseDateTime.RebaseSpec,
                                     requiredSchema: StructType,
                                     pushed: Option[FilterPredicate],
                                     enableRecordFilter: Boolean,
                                     enableLogicalTimestampRepair: Boolean,
                                     tableSchemaOpt: HOption[MessageType],
                                     schemaEvolutionUtils: ParquetSchemaEvolutionUtils): Iterator[InternalRow] with Closeable = {
    // scalastyle:on parameter.number
    // ParquetRecordReader returns InternalRow
    val readSupport = new HoodieParquetReadSupport(
      convertTz,
      enableVectorizedReader = false,
      enableLogicalTimestampRepair,
      datetimeRebaseSpec,
      int96RebaseSpec,
      tableSchemaOpt)
    val reader = if (pushed.isDefined && enableRecordFilter) {
      val parquetFilter = FilterCompat.get(pushed.get, null)
      new ParquetRecordReader[InternalRow](readSupport, parquetFilter)
    } else {
      new ParquetRecordReader[InternalRow](readSupport)
    }
    val readerWithRowIndexes = ParquetRowIndexUtil.addRowIndexToRecordReaderIfNeeded(reader,
      requiredSchema)
    val iter = new RecordReaderIterator[InternalRow](readerWithRowIndexes)
    try {
      readerWithRowIndexes.initialize(split, hadoopAttemptContext)

      val fullSchema = toAttributes(requiredSchema) ++ toAttributes(partitionSchema)
      val unsafeProjection = schemaEvolutionUtils.generateUnsafeProjection(fullSchema, timeZoneId)

      if (partitionSchema.length == 0) {
        // There is no partition columns
        iter.map(unsafeProjection)
      } else {
        val joinedRow = new JoinedRow()
        iter.map(d => unsafeProjection(joinedRow(d, partitionValues)))
      }
    } catch {
      case e: Throwable =>
        // SPARK-23457: In case there is an exception in initialization, close the iterator to
        // avoid leaking resources.
        iter.close()
        throw e
    }
  }
}

object Spark41ParquetReader extends SparkParquetReaderBuilder {
  /**
   * Get parquet file reader
   *
   * @param vectorized true if vectorized reading is not prohibited due to schema, reading mode, etc
   * @param sqlConf    the [[SQLConf]] used for the read
   * @param options    passed as a param to the file format
   * @param hadoopConf some configs will be set for the hadoopConf
   * @return parquet file reader
   */
  def build(vectorized: Boolean,
            sqlConf: SQLConf,
            options: Map[String, String],
            hadoopConf: Configuration): SparkColumnarFileReader = {
    //set hadoopconf
    hadoopConf.set(ParquetInputFormat.READ_SUPPORT_CLASS, classOf[ParquetReadSupport].getName)
    hadoopConf.set(SQLConf.SESSION_LOCAL_TIMEZONE.key, sqlConf.sessionLocalTimeZone)
    hadoopConf.setBoolean(SQLConf.NESTED_SCHEMA_PRUNING_ENABLED.key, sqlConf.nestedSchemaPruningEnabled)
    hadoopConf.setBoolean(SQLConf.CASE_SENSITIVE.key, sqlConf.caseSensitiveAnalysis)
    hadoopConf.setBoolean(SQLConf.PARQUET_BINARY_AS_STRING.key, sqlConf.isParquetBinaryAsString)
    hadoopConf.setBoolean(SQLConf.PARQUET_INT96_AS_TIMESTAMP.key, sqlConf.isParquetINT96AsTimestamp)
    // Using string value of this conf to preserve compatibility across spark versions. See [HUDI-5868]
    hadoopConf.setBoolean(
      SQLConf.LEGACY_PARQUET_NANOS_AS_LONG.key,
      sqlConf.getConfString(
        SQLConf.LEGACY_PARQUET_NANOS_AS_LONG.key,
        SQLConf.LEGACY_PARQUET_NANOS_AS_LONG.defaultValueString).toBoolean
    )
    hadoopConf.setBoolean(SQLConf.PARQUET_INFER_TIMESTAMP_NTZ_ENABLED.key, sqlConf.parquetInferTimestampNTZEnabled)

    val enableLogicalTimestampRepair = hadoopConf.getBoolean(ENABLE_LOGICAL_TIMESTAMP_REPAIR, true)
    val returningBatch = sqlConf.parquetVectorizedReaderEnabled &&
      options.getOrElse(FileFormat.OPTION_RETURNING_BATCH,
          throw new IllegalArgumentException(
            "OPTION_RETURNING_BATCH should always be set for ParquetFileFormat. " +
              "To workaround this issue, set spark.sql.parquet.enableVectorizedReader=false."))
        .equals("true")

    val parquetOptions = new ParquetOptions(options, sqlConf)
    new Spark41ParquetReader(
      enableVectorizedReader = vectorized,
      datetimeRebaseModeInRead = parquetOptions.datetimeRebaseModeInRead,
      int96RebaseModeInRead = parquetOptions.int96RebaseModeInRead,
      enableParquetFilterPushDown = sqlConf.parquetFilterPushDown,
      pushDownDate = sqlConf.parquetFilterPushDownDate,
      pushDownTimestamp = sqlConf.parquetFilterPushDownTimestamp,
      pushDownDecimal = sqlConf.parquetFilterPushDownDecimal,
      pushDownInFilterThreshold = sqlConf.parquetFilterPushDownInFilterThreshold,
      pushDownStringPredicate = sqlConf.parquetFilterPushDownStringPredicate,
      isCaseSensitive = sqlConf.caseSensitiveAnalysis,
      timestampConversion = sqlConf.isParquetINT96TimestampConversion,
      enableOffHeapColumnVector = sqlConf.offHeapColumnVectorEnabled,
      capacity = sqlConf.parquetVectorizedReaderBatchSize,
      returningBatch = returningBatch,
      enableRecordFilter = sqlConf.parquetRecordFilterEnabled,
      enableLogicalTimestampRepair = enableLogicalTimestampRepair,
      timeZoneId = Some(sqlConf.sessionLocalTimeZone))
  }

  // Helper to repair the schema if needed
  def repairFooterSchema(original: OpenedParquetFooter,
                         tableSchemaOpt: HOption[MessageType]): OpenedParquetFooter = {
    val originalParquetMetadata = original.footer();
    val repairedSchema = SchemaRepair.repairLogicalTypes(originalParquetMetadata.getFileMetaData.getSchema, tableSchemaOpt)
    val oldMeta = originalParquetMetadata.getFileMetaData
    new OpenedParquetFooter(new ParquetMetadata(
      new FileMetaData(
        repairedSchema,
        oldMeta.getKeyValueMetaData,
        oldMeta.getCreatedBy,
        oldMeta.getEncryptionType,
        oldMeta.getFileDecryptor
      ),
      originalParquetMetadata.getBlocks
    ), original.inputFile(), original.inputStreamOpt())
  }
}
