/*
 * 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.auron.plan

import java.util.UUID

import scala.jdk.CollectionConverters._

import org.apache.spark.Partition
import org.apache.spark.TaskContext
import org.apache.spark.sql.auron.{EmptyNativeRDD, NativeRDD}
import org.apache.spark.sql.execution.FileSourceScanExec
import org.apache.spark.sql.execution.datasources.FilePartition

import org.apache.auron.{protobuf => pb}
import org.apache.auron.metric.SparkMetricNode

abstract class NativeParquetScanBase(basedFileScan: FileSourceScanExec)
    extends NativeFileSourceScanBase(basedFileScan) {

  override def doExecuteNative(): NativeRDD = {
    val partitions = inputFileScanRDD.filePartitions.toArray
    if (partitions.length > 0) {
      val nativeMetrics = SparkMetricNode(
        metrics,
        Nil,
        Some({
          case ("bytes_scanned", v) =>
            val inputMetric = TaskContext.get.taskMetrics().inputMetrics
            inputMetric.incBytesRead(v)
          case ("output_rows", v) =>
            val inputMetric = TaskContext.get.taskMetrics().inputMetrics
            inputMetric.incRecordsRead(v)
          case _ =>
        }))
      val nativePruningPredicateFilters = this.nativePruningPredicateFilters
      val nativeFileSchema = this.nativeFileSchema
      val nativeFileGroups = this.nativeFileGroups
      val nativePartitionSchema = this.nativePartitionSchema

      val projection = schema.map(field => basedFileScan.relation.schema.fieldIndex(field.name))
      val broadcastedHadoopConf = this.broadcastedHadoopConf
      val numPartitions = partitions.length

      new NativeRDD(
        sparkContext,
        nativeMetrics,
        partitions.asInstanceOf[Array[Partition]],
        None,
        Nil,
        rddShuffleReadFull = true,
        (partition, _) => {
          val resourceId = s"NativeParquetScanExec:${UUID.randomUUID().toString}"
          putJniBridgeResource(resourceId, broadcastedHadoopConf)

          val nativeFileGroup = nativeFileGroups(partition.asInstanceOf[FilePartition])
          val nativeParquetScanConf = pb.FileScanExecConf
            .newBuilder()
            .setNumPartitions(numPartitions)
            .setPartitionIndex(partition.index)
            .setStatistics(pb.Statistics.getDefaultInstance)
            .setSchema(nativeFileSchema)
            .setFileGroup(nativeFileGroup)
            .addAllProjection(projection.map(Integer.valueOf).asJava)
            .setPartitionSchema(nativePartitionSchema)
            .build()

          val nativeParquetScanExecBuilder = pb.ParquetScanExecNode
            .newBuilder()
            .setBaseConf(nativeParquetScanConf)
            .setFsResourceId(resourceId)
            .addAllPruningPredicates(nativePruningPredicateFilters.asJava)

          pb.PhysicalPlanNode
            .newBuilder()
            .setParquetScan(nativeParquetScanExecBuilder.build())
            .build()
        },
        friendlyName = "NativeRDD.ParquetScan")
    } else {
      new EmptyNativeRDD(sparkContext)
    }
  }

  override val nodeName: String =
    s"NativeParquetScan ${basedFileScan.tableIdentifier.map(_.unquotedString).getOrElse("")}"
}
