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

import java.io.File
import java.math.{BigDecimal, BigInteger}
import java.time.{ZoneId, ZoneOffset}

import scala.reflect.ClassTag
import scala.reflect.runtime.universe.TypeTag
import scala.util.control.Breaks.breakable

import org.scalactic.source.Position
import org.scalatest.Tag

import org.apache.hadoop.fs.Path
import org.apache.parquet.example.data.simple.SimpleGroup
import org.apache.parquet.schema.MessageTypeParser
import org.apache.spark.SparkException
import org.apache.spark.sql.{CometTestBase, DataFrame, Row}
import org.apache.spark.sql.catalyst.util.DateTimeUtils
import org.apache.spark.sql.comet.{CometNativeScanExec, CometScanExec}
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.execution.datasources.parquet.ParquetUtils
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types._

import com.google.common.primitives.UnsignedLong

import org.apache.comet.CometConf

abstract class ParquetReadSuite extends CometTestBase {
  import testImplicits._

  testStandardAndLegacyModes("decimals") {
    Seq(true, false).foreach { useDecimal128 =>
      Seq(16, 1024).foreach { batchSize =>
        withSQLConf(
          CometConf.COMET_EXEC_ENABLED.key -> false.toString,
          CometConf.COMET_USE_DECIMAL_128.key -> useDecimal128.toString,
          CometConf.COMET_BATCH_SIZE.key -> batchSize.toString) {
          var combinations = Seq((5, 2), (1, 0), (18, 10), (18, 17), (19, 0), (38, 37))
          // If ANSI mode is on, the combination (1, 1) will cause a runtime error. Otherwise, the
          // decimal RDD contains all null values and should be able to read back from Parquet.

          if (!SQLConf.get.ansiEnabled) {
            combinations = combinations ++ Seq((1, 1))
          }
          for ((precision, scale) <- combinations; useDictionary <- Seq(false, true)) {
            withTempPath { dir =>
              val data = makeDecimalRDD(1000, DecimalType(precision, scale), useDictionary)
              data.write.parquet(dir.getCanonicalPath)
              readParquetFile(dir.getCanonicalPath) { df =>
                {
                  checkAnswer(df, data.collect().toSeq)
                }
              }
            }
          }
        }
      }
    }
  }

  test("simple count") {
    withParquetTable((0 until 10).map(i => (i, i.toString)), "tbl") {
      assert(sql("SELECT * FROM tbl WHERE _1 % 2 == 0").count() == 5)
    }
  }

  test("basic data types") {
    Seq(7, 1024).foreach { batchSize =>
      withSQLConf(CometConf.COMET_BATCH_SIZE.key -> batchSize.toString) {
        val data = (-100 to 100).map { i =>
          (
            i % 2 == 0,
            i,
            i.toByte,
            i.toShort,
            i.toLong,
            i.toFloat,
            i.toDouble,
            DateTimeUtils.toJavaDate(i))
        }
        if (!hasUnsignedSmallIntSafetyCheck(conf)) {
          checkParquetScan(data)
        }
        checkParquetFile(data)
      }
    }
  }

  test("basic data types with dictionary") {
    Seq(7, 1024).foreach { batchSize =>
      withSQLConf(CometConf.COMET_BATCH_SIZE.key -> batchSize.toString) {
        val data = (-100 to 100).map(_ % 4).map { i =>
          (
            i % 2 == 0,
            i,
            i.toByte,
            i.toShort,
            i.toLong,
            i.toFloat,
            i.toDouble,
            DateTimeUtils.toJavaDate(i))
        }
        if (!hasUnsignedSmallIntSafetyCheck(conf)) {
          checkParquetScan(data)
        }
        checkParquetFile(data)
      }
    }
  }

  test("basic filters") {
    val data = (-100 to 100).map { i =>
      (
        i % 2 == 0,
        i,
        i.toByte,
        i.toShort,
        i.toLong,
        i.toFloat,
        i.toDouble,
        DateTimeUtils.toJavaDate(i))
    }
    val filter = (row: Row) => row.getBoolean(0)
    if (!hasUnsignedSmallIntSafetyCheck(conf)) {
      checkParquetScan(data, filter)
    }
    checkParquetFile(data, filter)
  }

  test("raw binary test") {
    val data = (1 to 4).map(i => Tuple1(Array.fill(3)(i.toByte)))
    withParquetDataFrame(data) { df =>
      assertResult(data.map(_._1.mkString(",")).sorted) {
        df.collect().map(_.getAs[Array[Byte]](0).mkString(",")).sorted
      }
    }
  }

  test("string") {
    val data = (1 to 4).map(i => Tuple1(i.toString))
    // Property spark.sql.parquet.binaryAsString shouldn't affect Parquet files written by Spark SQL
    // as we store Spark SQL schema in the extra metadata.
    withSQLConf(SQLConf.PARQUET_BINARY_AS_STRING.key -> "false")(checkParquetFile(data))
    withSQLConf(SQLConf.PARQUET_BINARY_AS_STRING.key -> "true")(checkParquetFile(data))
  }

  test("string with dictionary") {
    Seq((100, 5), (1000, 10)).foreach { case (total, divisor) =>
      val data = (1 to total).map(i => Tuple1((i % divisor).toString))
      // Property spark.sql.parquet.binaryAsString shouldn't affect Parquet files written by Spark SQL
      // as we store Spark SQL schema in the extra metadata.
      withSQLConf(SQLConf.PARQUET_BINARY_AS_STRING.key -> "false")(checkParquetFile(data))
      withSQLConf(SQLConf.PARQUET_BINARY_AS_STRING.key -> "true")(checkParquetFile(data))
    }
  }

  test("long string + reserve additional space for value buffer") {
    withSQLConf(CometConf.COMET_BATCH_SIZE.key -> 16.toString) {
      val data = (1 to 100).map(i => (i, i.toString * 10))
      checkParquetFile(data)
    }
  }

  test("timestamp") {
    Seq(true, false).foreach { dictionaryEnabled =>
      withTempDir { dir =>
        val path = new Path(dir.toURI.toString, "part-r-0.parquet")
        val expected = makeRawTimeParquetFile(path, dictionaryEnabled = dictionaryEnabled, 10000)
        readParquetFile(path.toString) { df =>
          checkAnswer(
            df.select($"_0", $"_1", $"_2", $"_3", $"_4", $"_5"),
            expected.map {
              case None =>
                Row(null, null, null, null, null, null)
              case Some(i) =>
                val ts = new java.sql.Timestamp(i)
                val ldt = ts.toLocalDateTime
                  .atZone(ZoneId.systemDefault())
                  .withZoneSameInstant(ZoneOffset.UTC)
                  .toLocalDateTime
                Row(ts, ts, ts, ldt, ts, ldt)
            })
        }
      }
    }
  }

  test("timestamp as int96") {
    import testImplicits._

    val N = 100
    val ts = "2020-01-01 01:02:03.123456"
    Seq(false, true).foreach { dictionaryEnabled =>
      Seq(false, true).foreach { conversionEnabled =>
        withSQLConf(
          SQLConf.PARQUET_OUTPUT_TIMESTAMP_TYPE.key -> "INT96",
          SQLConf.PARQUET_INT96_TIMESTAMP_CONVERSION.key -> conversionEnabled.toString) {
          withTempPath { path =>
            Seq
              .tabulate(N)(_ => ts)
              .toDF("ts1")
              .select($"ts1".cast("timestamp").as("ts"))
              .repartition(1)
              .write
              .option("parquet.enable.dictionary", dictionaryEnabled)
              .parquet(path.getCanonicalPath)

            checkAnswer(
              spark.read.parquet(path.getCanonicalPath).select($"ts".cast("string")),
              Seq.tabulate(N)(_ => Row(ts)))
          }
        }
      }
    }
  }

  test("batch paging on basic types") {
    Seq(1, 2, 4, 9).foreach { batchSize =>
      withSQLConf(CometConf.COMET_BATCH_SIZE.key -> batchSize.toString) {
        val data = (1 to 10).map(i => (i, i.toByte, i.toShort, i.toFloat, i.toDouble, i.toString))
        checkParquetFile(data)
      }
    }
  }

  test("nulls") {
    val allNulls = (
      null.asInstanceOf[java.lang.Boolean],
      null.asInstanceOf[Integer],
      null.asInstanceOf[java.lang.Long],
      null.asInstanceOf[java.lang.Float],
      null.asInstanceOf[java.lang.Double],
      null.asInstanceOf[java.lang.String])

    withParquetDataFrame(allNulls :: Nil) { df =>
      val rows = df.collect()
      assert(rows.length === 1)
      assert(rows.head === Row(Seq.fill(6)(null): _*))
      assert(df.where("_1 is null").count() == 1)
    }
  }

  test("mixed nulls and non-nulls") {
    val rand = new scala.util.Random(42)
    val data = (0 to 100).map { i =>
      val row: (Boolean, Integer, java.lang.Long, java.lang.Float, java.lang.Double, String) = {
        if (rand.nextBoolean()) {
          (i % 2 == 0, i, i.toLong, i.toFloat, i.toDouble, i.toString)
        } else {
          (
            null.asInstanceOf[java.lang.Boolean],
            null.asInstanceOf[Integer],
            null.asInstanceOf[java.lang.Long],
            null.asInstanceOf[java.lang.Float],
            null.asInstanceOf[java.lang.Double],
            null.asInstanceOf[String])
        }
      }
      row
    }
    checkParquetFile(data)
  }

  test("vector reloading with all non-null values") {
    def makeRawParquetFile(
        path: Path,
        dictionaryEnabled: Boolean,
        n: Int,
        numNonNulls: Int): Seq[Option[Int]] = {
      val schemaStr =
        """
          |message root {
          | optional int32 _1;
          |}
        """.stripMargin

      val schema = MessageTypeParser.parseMessageType(schemaStr)
      val writer = createParquetWriter(schema, path, dictionaryEnabled = dictionaryEnabled)

      val expected = (0 until n).map { i =>
        if (i >= numNonNulls) {
          None
        } else {
          Some(i)
        }
      }
      expected.foreach { opt =>
        val record = new SimpleGroup(schema)
        opt match {
          case Some(i) =>
            record.add(0, i)
          case _ =>
        }
        writer.write(record)
      }

      writer.close()
      expected
    }

    Seq(2, 99, 1024).foreach { numNonNulls =>
      withTempDir { dir =>
        val path = new Path(dir.toURI.toString, "part-r-0.parquet")
        val expected = makeRawParquetFile(path, dictionaryEnabled = false, 1024, numNonNulls)
        withSQLConf(CometConf.COMET_BATCH_SIZE.key -> "2") {
          readParquetFile(path.toString) { df =>
            checkAnswer(
              df,
              expected.map {
                case None =>
                  Row(null)
                case Some(i) =>
                  Row(i)
              })
          }
        }
      }
    }
  }

  test("test lazy materialization skipping") {
    def makeRawParquetFile(
        path: Path,
        dictionaryEnabled: Boolean,
        pageSize: Int,
        pageRowCountLimit: Int,
        expected: Seq[Row]): Unit = {
      val schemaStr =
        """
          |message root {
          |  optional int32   _1;
          |  optional binary  _2(UTF8);
          |}
        """.stripMargin

      val schema = MessageTypeParser.parseMessageType(schemaStr)
      val writer = createParquetWriter(
        schema,
        path,
        dictionaryEnabled = dictionaryEnabled,
        pageSize = pageSize,
        dictionaryPageSize = pageSize,
        pageRowCountLimit = pageRowCountLimit)

      expected.foreach { row =>
        val record = new SimpleGroup(schema)
        record.add(0, row.getInt(0))
        record.add(1, row.getString(1))
        writer.write(record)
      }

      writer.close()
    }

    val skip = Row(0, "a") // row to skip by lazy materialization
    val read = Row(1, "b") // row not to skip
    // The initial page row count is always 100 in ParquetWriter, even with pageRowCountLimit config
    // Thus, use this header to fill in the first 100
    val header = Seq.fill(100)(skip)

    val expected = Seq( // spotless:off
      read, read, read, read, // read all rows in the page
      skip, skip, skip, skip, // skip all rows in the page
      skip, skip, skip, skip, // consecutively skip all rows in the page
      read, skip, skip, read, // skip middle rows in the page
      skip, read, read, skip, // read middle rows in the page
      skip, read, skip, read, // skip and read in turns
      read, skip, read, skip // skip and read in turns
    ) // spotless:on

    withTempDir { dir =>
      val path = new Path(dir.toURI.toString, "part-r-0.parquet")
      withSQLConf(
        CometConf.COMET_BATCH_SIZE.key -> "4",
        CometConf.COMET_EXEC_ENABLED.key -> "false") {
        makeRawParquetFile(path, dictionaryEnabled = false, 1024, 4, header ++ expected)
        readParquetFile(path.toString) { df =>
          checkAnswer(df.filter("_1 != 0"), expected.filter(_.getInt(0) != 0))
        }
      }
    }
  }

  test("test multiple pages with mixed PLAIN_DICTIONARY and PLAIN encoding") {
    // TODO: consider merging this with the same method above
    def makeRawParquetFile(path: Path, n: Int): Seq[Option[Int]] = {
      val dictionaryPageSize = 1024
      val pageRowCount = 500
      val schemaStr =
        """
          |message root {
          |  optional boolean _1;
          |  optional int32   _2(INT_8);
          |  optional int32   _3(INT_16);
          |  optional int32   _4;
          |  optional int64   _5;
          |  optional float   _6;
          |  optional double  _7;
          |  optional binary  _8(UTF8);
          |}
        """.stripMargin

      val schema = MessageTypeParser.parseMessageType(schemaStr)
      val writer = createParquetWriter(
        schema,
        path,
        dictionaryEnabled = true,
        dictionaryPageSize = dictionaryPageSize,
        pageRowCountLimit = pageRowCount)

      val rand = new scala.util.Random(42)
      val expected = (0 until n).map { i =>
        // use a single value for the first page, to make sure dictionary encoding kicks in
        val value = if (i < pageRowCount) i % 8 else i
        if (rand.nextBoolean()) None
        else Some(value)
      }

      expected.foreach { opt =>
        val record = new SimpleGroup(schema)
        opt match {
          case Some(i) =>
            record.add(0, i % 2 == 0)
            record.add(1, i.toByte)
            record.add(2, i.toShort)
            record.add(3, i)
            record.add(4, i.toLong)
            record.add(5, i.toFloat)
            record.add(6, i.toDouble)
            record.add(7, i.toString * 100)
          case _ =>
        }
        writer.write(record)
      }

      writer.close()
      expected
    }

    Seq(16, 128).foreach { batchSize =>
      withSQLConf(CometConf.COMET_BATCH_SIZE.key -> batchSize.toString) {
        withTempDir { dir =>
          val path = new Path(dir.toURI.toString, "part-r-0.parquet")
          val expected = makeRawParquetFile(path, 10000)
          readParquetFile(path.toString) { df =>
            checkAnswer(
              df,
              expected.map {
                case None =>
                  Row(null, null, null, null, null, null, null, null)
                case Some(i) =>
                  Row(
                    i % 2 == 0,
                    i.toByte,
                    i.toShort,
                    i,
                    i.toLong,
                    i.toFloat,
                    i.toDouble,
                    i.toString * 100)
              })
          }
        }
      }
    }
  }

  test("skip vector re-loading") {
    Seq(false, true).foreach { enableDictionary =>
      withSQLConf(
        CometConf.COMET_BATCH_SIZE.key -> 7.toString,
        CometConf.COMET_EXEC_ENABLED.key -> "true") {
        // Make sure this works with Comet native execution too
        val data = (1 to 100)
          .map(_ % 5) // trigger dictionary encoding
          .map(i => (i, i.toByte, i.toShort, i.toFloat, i.toDouble, i.toString))
        withParquetTable(data, "tbl", withDictionary = enableDictionary) {
          val df = sql("SELECT count(*) FROM tbl WHERE _1 >= 0")
          checkAnswer(df, Row(100) :: Nil)
        }
      }
    }
  }

  test("partition columns - multiple batch") {
    withSQLConf(
      CometConf.COMET_BATCH_SIZE.key -> 7.toString,
      CometConf.COMET_EXEC_ENABLED.key -> "false",
      CometConf.COMET_ENABLED.key -> "true") {
      Seq("a", null).foreach { partValue =>
        withTempPath { dir =>
          (1 to 100)
            .map(v => (partValue.asInstanceOf[String], v))
            .toDF("pcol", "col")
            .repartition(1)
            .write
            .format("parquet")
            .partitionBy("pcol")
            .save(dir.getCanonicalPath)
          val df = spark.read.format("parquet").load(dir.getCanonicalPath)
          assert(df.filter("col > 90").count() == 10)
        }
      }
    }
  }

  test("fix: string partition column with incorrect offset buffer") {
    def makeRawParquetFile(
        path: Path,
        dictionaryEnabled: Boolean,
        n: Int,
        pageSize: Int): Seq[Option[Int]] = {
      val schemaStr =
        """
          |message root {
          |  optional binary                  _1(UTF8);
          |}
    """.stripMargin

      val schema = MessageTypeParser.parseMessageType(schemaStr)
      val writer = createParquetWriter(
        schema,
        path,
        dictionaryEnabled = dictionaryEnabled,
        pageSize = pageSize,
        dictionaryPageSize = pageSize,
        rowGroupSize = 1024 * 128)

      val rand = new scala.util.Random(42)
      val expected = (0 until n).map { i =>
        if (rand.nextBoolean()) {
          None
        } else {
          Some(i)
        }
      }
      expected.foreach { opt =>
        val record = new SimpleGroup(schema)
        opt match {
          case Some(i) =>
            record.add(0, i.toString * 48)
          case _ =>
        }
        writer.write(record)
      }

      writer.close()
      expected
    }

    withTable("tbl") {
      withTempDir { dir =>
        val path = new Path(dir.toURI.toString, "part-r-0.parquet")
        makeRawParquetFile(path, false, 10000, 128)

        sql("CREATE TABLE tbl (value STRING, p STRING) USING PARQUET PARTITIONED BY (p) ")
        sql(s"ALTER TABLE tbl ADD PARTITION (p='a') LOCATION '$dir'")
        assert(sql("SELECT DISTINCT p FROM tbl").count() == 1)
      }
    }

  }

  test("missing columns") {
    withTempPath { dir =>
      Seq("a", "b").toDF("col1").write.parquet(dir.getCanonicalPath)

      // Create a schema where `col2` doesn't exist in the file schema
      var schema =
        StructType(Seq(StructField("col1", StringType), StructField("col2", IntegerType)))
      var df = spark.read.schema(schema).parquet(dir.getCanonicalPath)
      checkAnswer(df, Row("a", null) :: Row("b", null) :: Nil)

      // Should be the same when the missing column is at the beginning of the schema

      schema = StructType(Seq(StructField("col0", BooleanType), StructField("col1", StringType)))
      df = spark.read.schema(schema).parquet(dir.getCanonicalPath)
      checkAnswer(df, Row(null, "a") :: Row(null, "b") :: Nil)
    }
  }

  test("nested struct with new child field") {
    withTempPath { dir =>
      val path = dir.getCanonicalPath

      val df1 = sql("SELECT 1 c1, named_struct('c3', 2, 'c4', named_struct('c5', 3, 'c6', 4)) c2")
      val df2 = sql(
        "SELECT 1 c1, named_struct('c3', 2, 'c4', named_struct('c5', 3, 'c6', 4, 'c7', 5)) c2")
      val dir1 = s"$path${File.separator}part=one"
      val dir2 = s"$path${File.separator}part=two"

      df1.write.parquet(dir1)
      df2.write.parquet(dir2)

      var df = spark.read
        .schema(df2.schema)
        .load(path)
        .select($"c1", $"c2.c3")

      // read a leaf field
      val expected1 =
        Seq(Row(1, 2), Row(1, 2))
      checkAnswer(df, expected1)

      // read a struct field
      df = spark.read
        .schema(df2.schema)
        .load(path)
        .select($"c1", $"c2")
      val expected2 =
        Seq(Row(1, Row(2, Row(3, 4, null))), Row(1, Row(2, Row(3, 4, 5))))
      checkAnswer(df, expected2)

      // read a missing field
      df = spark.read
        .schema(df2.schema)
        .load(path)
        .select($"c1", $"c2.c4.c7")
      val expected3 =
        Seq(Row(1, null), Row(1, 5))
      checkAnswer(df, expected3)
    }
  }

  test("reading required fields in structs") {
    Seq(true, false).foreach { dictionaryEnabled =>
      def makeRawParquetFile(path: Path): Unit = {
        val schemaStr =
          """message schema {
            |  optional int32 a;
            |  required group b {
            |    optional int32 b1;
            |    required int32 b2;
            |  }
            |  optional group c {
            |    required int32 c1;
            |    optional int32 c2;
            |  }
            |}
        """.stripMargin
        val schema = MessageTypeParser.parseMessageType(schemaStr)

        val writer = createParquetWriter(schema, path, dictionaryEnabled)

        // All optional fields are specified
        var record = new SimpleGroup(schema)
        record.add("a", 1)
        var b = record.addGroup("b")
        b.add("b1", 1)
        b.add("b2", 1)
        var c = record.addGroup("c")
        c.add("c1", 1)
        c.add("c2", 1)
        writer.write(record)

        // No optional values
        record = new SimpleGroup(schema)
        record.add("a", 2)
        b = record.addGroup("b")
        b.add("b1", 2)
        b.add("b2", 2)
        writer.write(record)

        writer.close()
      }

      withTempDir { dir =>
        val path = new Path(dir.toURI.toString, "part-r-0.parquet")
        makeRawParquetFile(path)

        withParquetTable(spark.read.format("parquet").load(path.toString), "complex_types") {
          // required struct field with optional field
          var df = sql("select a, b from complex_types")
          checkSparkAnswer(df)

          // Missing optional struct field
          df = sql("select a, c from complex_types")
          checkSparkAnswer(df)

          // Missing optional struct field with nested required field
          // TODO: This produces incorrect results in both native_datafusion and native_iceberg_compat
          //          df = sql("select a, c.c1 from complex_types")
          //          checkSparkAnswer(df)

          // required nested field in a struct
          df = sql("select a, b.b2 from complex_types")
          checkSparkAnswer(df)

        }
      }
    }
  }

  // TODO: This test fails for both native_datafusion and native_iceberg_compat
  ignore(" Missing optional struct field with nested required field") {
    Seq(true, false).foreach { dictionaryEnabled =>
      def makeRawParquetFile(path: Path): Unit = {
        val schemaStr =
          """message schema {
            |  optional int32 a;
            |  required group b {
            |    optional int32 b1;
            |    required int32 b2;
            |  }
            |  optional group c {
            |    required int32 c1;
            |    optional int32 c2;
            |  }
            |}
        """.stripMargin
        val schema = MessageTypeParser.parseMessageType(schemaStr)

        val writer = createParquetWriter(schema, path, dictionaryEnabled)

        // All optional fields are specified
        var record = new SimpleGroup(schema)
        record.add("a", 1)
        var b = record.addGroup("b")
        b.add("b1", 1)
        b.add("b2", 1)
        var c = record.addGroup("c")
        c.add("c1", 1)
        c.add("c2", 1)
        writer.write(record)

        // No optional values
        record = new SimpleGroup(schema)
        record.add("a", 2)
        b = record.addGroup("b")
        b.add("b1", 2)
        b.add("b2", 2)
        writer.write(record)

        writer.close()
      }

      withTempDir { dir =>
        val path = new Path(dir.toURI.toString, "part-r-0.parquet")
        makeRawParquetFile(path)

        withParquetTable(spark.read.format("parquet").load(path.toString), "complex_types") {
          // Missing optional struct field with nested required field
          // TODO: This produces incorrect results in both native_datafusion and native_iceberg_compat
          val df = sql("select a, c.c1 from complex_types")
          checkSparkAnswer(df)

        }
      }
    }
  }

  test("missing required fields in structs") {
    Seq(true, false).foreach { dictionaryEnabled =>
      def makeRawParquetFile(path: Path): Unit = {
        val schemaStr =
          """message schema {
            |  optional int32 a;
            |  required group b {
            |    optional int32 b1;
            |    required int32 b2;
            |  }
            |  optional group c {
            |    required int32 c1;
            |    optional int32 c2;
            |  }
            |}
        """.stripMargin
        val schema = MessageTypeParser.parseMessageType(schemaStr)
        val writer = createParquetWriter(schema, path, dictionaryEnabled)
        // Missing required field b.b2
        val record = new SimpleGroup(schema)
        record.add("a", 3)
        val b = record.addGroup("b")
        b.add("b1", 4)
        writer.write(record)

        writer.close()
      }

      withTempDir { dir =>
        val path = new Path(dir.toURI.toString, "part-r-0.parquet")
        makeRawParquetFile(path)

        withParquetTable(
          spark.read.format("parquet").load(path.toString),
          "complex_types_missing_struct") {
          // missing required nested field in a struct
          var failed = false
          try {
            val df = sql("select a, b.b2 from complex_types_missing_struct")
            checkSparkAnswer(df)
          } catch {
            case _: Exception => failed = true
          }
          assert(failed)
          // missing required struct field
          failed = false
          try {
            val df = sql("select a, b from complex_types_missing_struct")
            checkSparkAnswer(df)
          } catch {
            case _: Exception => failed = true
          }
          assert(failed)
        }
      }
    }
  }

  test("unsigned int supported") {
    Seq(true, false).foreach { dictionaryEnabled =>
      def makeRawParquetFile(path: Path): Unit = {
        val schemaStr =
          """message root {
            |  required INT32 a(UINT_8);
            |  required INT32 b(UINT_16);
            |  required INT32 c(UINT_32);
            |}
        """.stripMargin
        val schema = MessageTypeParser.parseMessageType(schemaStr)

        val writer = createParquetWriter(schema, path, dictionaryEnabled)

        (0 until 10).foreach { n =>
          val record = new SimpleGroup(schema)
          record.add(0, n.toByte + Byte.MaxValue)
          record.add(1, n.toShort + Short.MaxValue)
          record.add(2, n + Int.MaxValue)
          writer.write(record)
        }
        writer.close()
      }

      withTempDir { dir =>
        val path = new Path(dir.toURI.toString, "part-r-0.parquet")
        makeRawParquetFile(path)
        readParquetFile(path.toString) { df =>
          checkAnswer(
            df,
            (0 until 10).map(n =>
              Row(n.toByte + Byte.MaxValue, n.toShort + Short.MaxValue, n + Int.MaxValue.toLong)))
        }
      }
    }
  }

  test("unsigned long supported") {
    Seq(true, false).foreach { dictionaryEnabled =>
      def makeRawParquetFile(path: Path): Unit = {
        val schemaStr =
          """message root {
            |  required INT64 a(UINT_64);
            |}
        """.stripMargin
        val schema = MessageTypeParser.parseMessageType(schemaStr)

        val writer = createParquetWriter(schema, path, dictionaryEnabled)

        (0 until 10).map(_.toLong).foreach { n =>
          val record = new SimpleGroup(schema)
          record.add(0, n + Long.MaxValue)
          writer.write(record)
        }
        writer.close()
      }

      withTempDir { dir =>
        val path = new Path(dir.toURI.toString, "part-r-0.parquet")
        makeRawParquetFile(path)
        readParquetFile(path.toString) { df =>
          checkAnswer(
            df,
            (0 until 10).map(n =>
              Row(
                new BigDecimal(UnsignedLong.fromLongBits(n + Long.MaxValue).bigIntegerValue()))))
        }
      }
    }
  }

  test("enum support") {
    // https://github.com/apache/parquet-format/blob/master/LogicalTypes.md
    // "enum type should interpret ENUM annotated field as a UTF-8"
    Seq(true, false).foreach { dictionaryEnabled =>
      def makeRawParquetFile(path: Path): Unit = {
        val schemaStr =
          """message root {
            |  required BINARY a(ENUM);
            |}
        """.stripMargin
        val schema = MessageTypeParser.parseMessageType(schemaStr)

        val writer = createParquetWriter(schema, path, dictionaryEnabled)

        (0 until 10).map(_.toLong).foreach { n =>
          val record = new SimpleGroup(schema)
          record.add(0, n.toString)
          writer.write(record)
        }
        writer.close()
      }

      withTempDir { dir =>
        val path = new Path(dir.toURI.toString, "part-r-0.parquet")
        makeRawParquetFile(path)
        readParquetFile(path.toString) { df =>
          checkAnswer(df, (0 until 10).map(n => Row(n.toString)))
        }
      }
    }
  }

  test("FIXED_LEN_BYTE_ARRAY support") {
    Seq(true, false).foreach { dictionaryEnabled =>
      def makeRawParquetFile(path: Path): Unit = {
        val schemaStr =
          """message root {
            |  required FIXED_LEN_BYTE_ARRAY(1) a;
            |  required FIXED_LEN_BYTE_ARRAY(3) b;
            |}
        """.stripMargin
        val schema = MessageTypeParser.parseMessageType(schemaStr)

        val writer = createParquetWriter(schema, path, dictionaryEnabled)

        (0 until 10).map(_.toString).foreach { n =>
          val record = new SimpleGroup(schema)
          record.add(0, n)
          record.add(1, n + n + n)
          writer.write(record)
        }
        writer.close()
      }

      withTempDir { dir =>
        val path = new Path(dir.toURI.toString, "part-r-0.parquet")
        makeRawParquetFile(path)
        readParquetFile(path.toString) { df =>
          checkAnswer(
            df,
            (48 until 58).map(n => // char '0' is 48 in ascii
              Row(Array(n), Array(n, n, n))))
        }
      }
    }
  }

  test("schema evolution") {
    Seq(true, false).foreach { enableSchemaEvolution =>
      Seq(true, false).foreach { useDictionary =>
        {
          withSQLConf(
            CometConf.COMET_SCHEMA_EVOLUTION_ENABLED.key -> enableSchemaEvolution.toString) {
            val data = (0 until 100).map(i => {
              val v = if (useDictionary) i % 5 else i
              (v, v.toFloat)
            })
            val readSchema =
              StructType(
                Seq(StructField("_1", LongType, false), StructField("_2", DoubleType, false)))

            withParquetDataFrame(data, schema = Some(readSchema)) { df =>
              val scan = CometConf.COMET_NATIVE_SCAN_IMPL.get(conf)
              val isNativeDataFusionScan =
                scan == CometConf.SCAN_NATIVE_DATAFUSION || scan == CometConf.SCAN_AUTO
              if (enableSchemaEvolution || isNativeDataFusionScan) {
                // native_datafusion has more permissive schema evolution
                // https://github.com/apache/datafusion-comet/issues/3720
                checkAnswer(df, data.map(Row.fromTuple))
              } else {
                assertThrows[SparkException](df.collect())
              }
            }
          }
        }
      }
    }
  }

  test("native_datafusion rejects string read as non-string/binary type") {
    // Regression guard for https://github.com/apache/datafusion-comet/issues/4088.
    // Spark's vectorized reader rejects reading a Parquet BINARY column as
    // anything except StringType, BinaryType, or a binary-encoded decimal (see
    // TypeUtil.checkParquetType, BINARY case). The native_datafusion scan
    // must do the same in its schema adapter rather than letting DataFusion's
    // cast silently parse the bytes or reinterpret them.
    withSQLConf(
      CometConf.COMET_NATIVE_SCAN_IMPL.key -> CometConf.SCAN_NATIVE_DATAFUSION,
      SQLConf.USE_V1_SOURCE_LIST.key -> "parquet") {
      withTempPath { dir =>
        val path = dir.getCanonicalPath
        Seq("a", "b", "c").toDF("c").write.parquet(path)
        // Cover representative non-string/binary target types: numeric,
        // boolean, date, and timestamp. Each would silently produce wrong
        // results without the schema-adapter guard.
        Seq("int", "bigint", "double", "boolean", "date", "timestamp").foreach { sqlType =>
          val df = spark.read.schema(s"c $sqlType").parquet(path)
          assertThrows[SparkException](df.collect())
        }
      }
    }
  }

  test("native_datafusion rejects incompatible decimal precision/scale") {
    // Regression guard for https://github.com/apache/datafusion-comet/issues/4089.
    // Reading Decimal(10,2) under a Decimal(5,0) read schema is unconditionally
    // lossy: target precision is smaller than source precision and scales differ.
    // Spark's vectorized reader throws SchemaColumnConvertNotSupportedException
    // here on all versions. The native_datafusion scan must reject this in its
    // schema adapter rather than letting Spark Cast silently rescale/truncate.
    withSQLConf(
      CometConf.COMET_NATIVE_SCAN_IMPL.key -> CometConf.SCAN_NATIVE_DATAFUSION,
      SQLConf.USE_V1_SOURCE_LIST.key -> "parquet") {
      withTempPath { dir =>
        val path = dir.getCanonicalPath
        spark
          .sql("select cast('123.45' as decimal(10,2)) as d " +
            "union all select cast('67.89' as decimal(10,2))")
          .write
          .parquet(path)
        val df = spark.read.schema("d decimal(5,0)").parquet(path)
        assertThrows[SparkException](df.collect())
      }
    }
  }

  test("type widening: byte → short/int/long, short → int/long, int → long") {
    withSQLConf(CometConf.COMET_SCHEMA_EVOLUTION_ENABLED.key -> "true") {
      withTempPath { dir =>
        val path = dir.getCanonicalPath
        val values = 1 to 10
        val options: Map[String, String] = Map.empty[String, String]

        // Input types and corresponding DataFrames
        val inputDFs = Seq(
          "byte" -> values.map(_.toByte).toDF("col1"),
          "short" -> values.map(_.toShort).toDF("col1"),
          "int" -> values.map(_.toInt).toDF("col1"))

        // Target Spark read schemas for widening
        val widenTargets = Seq(
          "short" -> values.map(_.toShort).toDF("col1"),
          "int" -> values.map(_.toInt).toDF("col1"),
          "long" -> values.map(_.toLong).toDF("col1"))

        for ((inputType, inputDF) <- inputDFs) {
          val writePath = s"$path/$inputType"
          inputDF.write.format("parquet").options(options).save(writePath)

          for ((targetType, targetDF) <- widenTargets) {
            // Only test valid widenings (e.g., don't test int → short)
            val wideningValid = (inputType, targetType) match {
              case ("byte", "short" | "int" | "long") => true
              case ("short", "int" | "long") => true
              case ("int", "long") => true
              case _ => false
            }

            if (wideningValid) {
              val reader = spark.read
                .schema(s"col1 $targetType")
                .format("parquet")
                .options(options)
                .load(writePath)

              checkAnswer(reader, targetDF)
            }
          }
        }
      }
    }
  }

  test("read byte, int, short, long together") {
    withSQLConf(CometConf.COMET_SCHEMA_EVOLUTION_ENABLED.key -> "true") {
      withTempPath { dir =>
        val path = dir.getCanonicalPath

        val byteDF = (Byte.MaxValue - 2 to Byte.MaxValue).map(_.toByte).toDF("col1")
        val shortDF = (Short.MaxValue - 2 to Short.MaxValue).map(_.toShort).toDF("col1")
        val intDF = (Int.MaxValue - 2 to Int.MaxValue).toDF("col1")
        val longDF = (Long.MaxValue - 2 to Long.MaxValue).toDF("col1")
        val unionDF = byteDF.union(shortDF).union(intDF).union(longDF)

        val byteDir = s"$path${File.separator}part=byte"
        val shortDir = s"$path${File.separator}part=short"
        val intDir = s"$path${File.separator}part=int"
        val longDir = s"$path${File.separator}part=long"

        val options: Map[String, String] = Map.empty[String, String]

        byteDF.write.format("parquet").options(options).save(byteDir)
        shortDF.write.format("parquet").options(options).save(shortDir)
        intDF.write.format("parquet").options(options).save(intDir)
        longDF.write.format("parquet").options(options).save(longDir)

        val df = spark.read
          .schema(unionDF.schema)
          .format("parquet")
          .options(options)
          .load(path)
          .select("col1")

        checkAnswer(df, unionDF)
      }
    }
  }

  test("read dictionary encoded decimals written as FIXED_LEN_BYTE_ARRAY") {
    // In this test, data is encoded using Parquet page v2 format, but with PLAIN encoding
    checkAnswer(
      // Decimal column in this file is encoded using plain dictionary
      readResourceParquetFile("test-data/dec-in-fixed-len.parquet"),
      spark.range(1 << 4).select('id % 10 cast DecimalType(10, 2) as 'fixed_len_dec))
  }

  test("read long decimals with precision <= 9") {
    // decimal32-written-as-64-bit.snappy.parquet was generated using a 3rd-party library. It has
    // 10 rows of Decimal(9, 1) written as LongDecimal instead of an IntDecimal
    var df = readResourceParquetFile("test-data/decimal32-written-as-64-bit.snappy.parquet")
    assert(10 == df.collect().length)
    var first10Df = df.head(10)
    assert(
      Seq(792059492, 986842987, 540247998, null, 357991078, 494131059, 92536396, 426847157,
        -999999999, 204486094)
        .zip(first10Df)
        .forall(d =>
          d._2.isNullAt(0) && d._1 == null ||
            d._1 == d._2.getDecimal(0).unscaledValue().intValue()))

    // decimal32-written-as-64-bit-dict.snappy.parquet was generated using a 3rd-party library. It
    // has 2048 rows of Decimal(3, 1) written as LongDecimal instead of an IntDecimal
    df = readResourceParquetFile("test-data/decimal32-written-as-64-bit-dict.snappy.parquet")
    assert(2048 == df.collect().length)
    first10Df = df.head(10)
    assert(
      Seq(751, 937, 511, null, 337, 467, 84, 403, -999, 190)
        .zip(first10Df)
        .forall(d =>
          d._2.isNullAt(0) && d._1 == null ||
            d._1 == d._2.getDecimal(0).unscaledValue().intValue()))

    val last10Df = df.tail(10)
    assert(
      Seq(866, 20, 492, 76, 824, 604, 343, 820, 864, 243)
        .zip(last10Df)
        .forall(d => d._1 == d._2.getDecimal(0).unscaledValue().intValue()))
  }

  private val actions: Seq[DataFrame => DataFrame] = Seq(
    "_1 = 500",
    "_1 = 500 or _1 = 1500",
    "_1 = 500 or _1 = 501 or _1 = 1500",
    "_1 = 500 or _1 = 501 or _1 = 1000 or _1 = 1500",
    "_1 >= 500 and _1 < 1000",
    "(_1 >= 500 and _1 < 1000) or (_1 >= 1500 and _1 < 1600)").map(f =>
    (df: DataFrame) => df.filter(f))

  test("test lazy materialization when batch size is small") {
    val df = spark.range(0, 2000).selectExpr("id as _1", "cast(id as string) as _11")
    checkParquetDataFrame(df)(actions: _*)
  }

  test("test lazy materialization when batch size is small (dict encode)") {
    val df = spark.range(0, 2000).selectExpr("id as _1", "cast(id % 10 as string) as _11")
    checkParquetDataFrame(df)(actions: _*)
  }

  private def testStandardAndLegacyModes(testName: String)(f: => Unit): Unit = {
    test(s"Standard mode - $testName") {
      withSQLConf(SQLConf.PARQUET_WRITE_LEGACY_FORMAT.key -> "false") { f }
    }

    test(s"Legacy mode - $testName") {
      withSQLConf(SQLConf.PARQUET_WRITE_LEGACY_FORMAT.key -> "true") { f }
    }
  }

  private def checkParquetFile[T <: Product: ClassTag: TypeTag](
      data: Seq[T],
      f: Row => Boolean = _ => true): Unit = {
    withParquetDataFrame(data)(r => checkAnswer(r.filter(f), data.map(Row.fromTuple).filter(f)))
  }

  protected def checkParquetScan[T <: Product: ClassTag: TypeTag](
      data: Seq[T],
      f: Row => Boolean = _ => true): Unit

  /**
   * create parquet file with various page sizes and batch sizes
   */
  private def checkParquetDataFrame(df: DataFrame)(actions: (DataFrame => DataFrame)*): Unit = {
    Seq(true, false).foreach { enableDictionary =>
      Seq(64, 127, 4049).foreach { pageSize =>
        withTempPath(file => {
          df.coalesce(1)
            .write
            .option("parquet.page.size", pageSize.toString)
            .option("parquet.enable.dictionary", enableDictionary.toString)
            .parquet(file.getCanonicalPath)

          Seq(true, false).foreach { useLazyMaterialization =>
            Seq(true, false).foreach { enableCometExec =>
              Seq(4, 13, 4049).foreach { batchSize =>
                withSQLConf(
                  CometConf.COMET_BATCH_SIZE.key -> batchSize.toString,
                  CometConf.COMET_EXEC_ENABLED.key -> enableCometExec.toString,
                  CometConf.COMET_USE_LAZY_MATERIALIZATION.key -> useLazyMaterialization.toString) {
                  readParquetFile(file.getCanonicalPath) { parquetDf =>
                    actions.foreach { action =>
                      checkAnswer(action(parquetDf), action(df))
                    }
                  }
                }
              }
            }
          }
        })
      }
    }
  }

  test("test merge scan range") {
    def makeRawParquetFile(path: Path, n: Int): Seq[Option[Int]] = {
      val dictionaryPageSize = 1024
      val pageRowCount = 500
      val schemaStr =
        """
          |message root {
          |  optional int32   _1(INT_16);
          |  optional int32   _2;
          |  optional int64   _3;
          |}
        """.stripMargin

      val schema = MessageTypeParser.parseMessageType(schemaStr)
      val writer = createParquetWriter(
        schema,
        path,
        dictionaryEnabled = true,
        dictionaryPageSize = dictionaryPageSize,
        pageRowCountLimit = pageRowCount)

      val rand = new scala.util.Random(42)
      val expected = (0 until n).map { i =>
        // use a single value for the first page, to make sure dictionary encoding kicks in
        val value = if (i < pageRowCount) i % 8 else i
        if (rand.nextBoolean()) None
        else Some(value)
      }

      expected.foreach { opt =>
        val record = new SimpleGroup(schema)
        opt match {
          case Some(i) =>
            record.add(0, i.toShort)
            record.add(1, i)
            record.add(2, i.toLong)
          case _ =>
        }
        writer.write(record)
      }

      writer.close()
      expected
    }

    Seq(16, 128).foreach { batchSize =>
      Seq(1024, 1024 * 1024).foreach { mergeRangeDelta =>
        {
          withSQLConf(
            CometConf.COMET_BATCH_SIZE.key -> batchSize.toString,
            CometConf.COMET_IO_MERGE_RANGES.key -> "true",
            CometConf.COMET_IO_MERGE_RANGES_DELTA.key -> mergeRangeDelta.toString) {
            withTempDir { dir =>
              val path = new Path(dir.toURI.toString, "part-r-0.parquet")
              val expected = makeRawParquetFile(path, 10000)
              val schema = StructType(
                Seq(StructField("_1", ShortType, true), StructField("_3", LongType, true)))
              readParquetFile(path.toString, Some(schema)) { df =>
                {
                  // CometScanExec calls sessionState.newHadoopConfWithOptions which copies
                  // the sqlConf and some additional options to the hadoopConf and then
                  // uses the result to create the inputRDD (https://github.com/apache/datafusion-comet/blob/3783faaa01078a35bee93b299368f8c72869198d/spark/src/main/scala/org/apache/spark/sql/comet/CometScanExec.scala#L181).
                  // We don't have access to the created hadoop Conf, but can confirm that the
                  // result does contain the correct configuration
                  assert(
                    df.sparkSession.sessionState
                      .newHadoopConfWithOptions(Map.empty)
                      .get(CometConf.COMET_IO_MERGE_RANGES_DELTA.key)
                      .equals(mergeRangeDelta.toString))
                  checkAnswer(
                    df,
                    expected.map {
                      case None =>
                        Row(null, null)
                      case Some(i) =>
                        Row(i.toShort, i.toLong)
                    })
                }
              }
            }
          }
        }
      }
    }
  }

  def testScanner(
      cometEnabled: String,
      cometNativeScanImpl: String,
      scanner: String,
      v1: Option[String] = None): Unit = {
    withSQLConf(
      CometConf.COMET_ENABLED.key -> cometEnabled,
      CometConf.COMET_EXEC_ENABLED.key -> cometEnabled,
      CometConf.COMET_NATIVE_SCAN_IMPL.key -> cometNativeScanImpl,
      SQLConf.USE_V1_SOURCE_LIST.key -> v1.getOrElse("")) {
      withParquetTable(Seq((Long.MaxValue, 1), (Long.MaxValue, 2)), "tbl") {
        val df = spark.sql("select * from tbl")
        assert(
          stripAQEPlan(df.queryExecution.executedPlan)
            .collectLeaves()
            .head
            .toString()
            .startsWith(scanner))
      }
    }
  }

  private def withId(id: Int) =
    new MetadataBuilder().putLong(ParquetUtils.FIELD_ID_METADATA_KEY, id).build()

  // Based on Spark ParquetIOSuite.test("vectorized reader: array of nested struct")
  test("array of nested struct with and without field id") {
    val nestedSchema = StructType(
      Seq(StructField(
        "_1",
        StructType(Seq(
          StructField("_1", StringType, nullable = true, withId(1)), // Field ID 1
          StructField(
            "_2",
            ArrayType(StructType(Seq(
              StructField("_1", StringType, nullable = true, withId(2)), // Field ID 2
              StructField("_2", StringType, nullable = true, withId(3)) // Field ID 3
            ))),
            nullable = true))),
        nullable = true)))
    val nestedSchemaNoId = StructType(
      Seq(StructField(
        "_1",
        StructType(Seq(
          StructField("_1", StringType, nullable = true),
          StructField(
            "_2",
            ArrayType(StructType(Seq(
              StructField("_1", StringType, nullable = true),
              StructField("_2", StringType, nullable = true)))),
            nullable = true))),
        nullable = true)))
    // data matching the schema
    val data = Seq(
      Row(Row("a", null)),
      Row(Row("b", Seq(Row("c", "d")))),
      Row(null),
      Row(Row("e", Seq(Row("f", null), Row(null, "g")))),
      Row(Row(null, null)),
      Row(Row(null, Seq(null))),
      Row(Row(null, Seq(Row(null, null), Row("h", null), null))),
      Row(Row("i", Seq())),
      Row(null))
    val answer =
      Row(Row("a", null)) ::
        Row(Row("b", Seq(Row("c", "d")))) ::
        Row(null) ::
        Row(Row("e", Seq(Row("f", null), Row(null, "g")))) ::
        Row(Row(null, null)) ::
        Row(Row(null, Seq(null))) ::
        Row(Row(null, Seq(Row(null, null), Row("h", null), null))) ::
        Row(Row("i", Seq())) ::
        Row(null) ::
        Nil

    withSQLConf(SQLConf.PARQUET_FIELD_ID_READ_ENABLED.key -> "true") {
      val df = spark.createDataFrame(spark.sparkContext.parallelize(data), nestedSchema)
      withTempPath { path =>
        df.write.parquet(path.getCanonicalPath)
        readParquetFile(path.getCanonicalPath) { df =>
          checkAnswer(df, answer)
        }
      }
      val df2 = spark.createDataFrame(spark.sparkContext.parallelize(data), nestedSchemaNoId)
      withTempPath { path =>
        df2.write.parquet(path.getCanonicalPath)
        readParquetFile(path.getCanonicalPath) { df =>
          checkAnswer(df, answer)
        }
      }
    }
  }

  // Based on Spark ParquetFieldIdIOSuite.test("Parquet reads infer fields using field ids
  // correctly"). Forces SCAN_NATIVE_DATAFUSION so we can prove that the gate in CometScanRule
  // is removed and that the native_datafusion scan resolves columns by field id rather than by
  // name (the read schema names differ from what is in the file).
  test("native_datafusion: read by Parquet field id when names differ") {
    val writeSchema = StructType(
      Seq(
        StructField("random", IntegerType, nullable = true, withId(1)),
        StructField("name", StringType, nullable = true, withId(0))))
    val readSchema = StructType(
      Seq(
        StructField("a", StringType, nullable = true, withId(0)),
        StructField("b", IntegerType, nullable = true, withId(1))))
    val writeData = Seq(Row(100, "text"), Row(200, "more"))

    withSQLConf(
      CometConf.COMET_NATIVE_SCAN_IMPL.key -> CometConf.SCAN_NATIVE_DATAFUSION,
      SQLConf.PARQUET_FIELD_ID_READ_ENABLED.key -> "true") {
      withTempPath { dir =>
        spark
          .createDataFrame(spark.sparkContext.parallelize(writeData), writeSchema)
          .write
          .mode("overwrite")
          .parquet(dir.getCanonicalPath)
        val df = spark.read.schema(readSchema).parquet(dir.getCanonicalPath)
        checkSparkAnswerAndOperator(df)
      }
    }
  }

  // Based on Spark ParquetFieldIdIOSuite.test("SPARK-38094: absence of field ids: reading nested
  // schema"). Exercises ID matching at every nesting level (struct, array<struct>, map) under
  // SCAN_NATIVE_DATAFUSION. Names differ from the file at every level.
  test("native_datafusion: read nested types by Parquet field id when names differ") {
    val writeSchema = StructType(
      Seq(StructField(
        "outer",
        StructType(Seq(
          StructField("inner_a", IntegerType, nullable = true, withId(11)),
          StructField(
            "inner_arr",
            ArrayType(StructType(Seq(
              StructField("ea", StringType, nullable = true, withId(21)),
              StructField("eb", IntegerType, nullable = true, withId(22))))),
            nullable = true,
            withId(12)),
          StructField(
            "inner_map",
            MapType(StringType, IntegerType, valueContainsNull = true),
            nullable = true,
            withId(13)))),
        nullable = true,
        withId(1))))

    val readSchema = StructType(
      Seq(StructField(
        "renamed_outer",
        StructType(Seq(
          StructField("renamed_a", IntegerType, nullable = true, withId(11)),
          StructField(
            "renamed_arr",
            ArrayType(StructType(Seq(
              StructField("renamed_ea", StringType, nullable = true, withId(21)),
              StructField("renamed_eb", IntegerType, nullable = true, withId(22))))),
            nullable = true,
            withId(12)),
          StructField(
            "renamed_map",
            MapType(StringType, IntegerType, valueContainsNull = true),
            nullable = true,
            withId(13)))),
        nullable = true,
        withId(1))))

    val data = Seq(
      Row(Row(1, Seq(Row("x", 10), Row("y", 20)), Map("k1" -> 100))),
      Row(Row(2, Seq(Row("z", 30)), Map("k2" -> 200, "k3" -> 300))))

    withSQLConf(
      CometConf.COMET_NATIVE_SCAN_IMPL.key -> CometConf.SCAN_NATIVE_DATAFUSION,
      SQLConf.PARQUET_FIELD_ID_READ_ENABLED.key -> "true") {
      withTempPath { dir =>
        spark
          .createDataFrame(spark.sparkContext.parallelize(data), writeSchema)
          .write
          .mode("overwrite")
          .parquet(dir.getCanonicalPath)
        val df = spark.read.schema(readSchema).parquet(dir.getCanonicalPath)
        checkSparkAnswerAndOperator(df)
      }
    }
  }

  // Verbatim port of Spark `ParquetFieldIdIOSuite.test("multiple id matches")`, pinned to
  // `SCAN_NATIVE_DATAFUSION` so the shim error path is exercised on both 3.x and 4.x.
  // The stock suite is the CI signal but it requires the Spark test jars and
  // `withAllParquetReaders`; keeping a copy here lets us iterate locally.
  test("multiple id matches") {
    withSQLConf(
      CometConf.COMET_NATIVE_SCAN_IMPL.key -> CometConf.SCAN_NATIVE_DATAFUSION,
      SQLConf.PARQUET_FIELD_ID_READ_ENABLED.key -> "true") {
      withTempPath { dir =>
        val readSchema =
          new StructType()
            .add("a", IntegerType, true, withId(1))

        val writeSchema =
          new StructType()
            .add("a", IntegerType, true, withId(1))
            .add("rand1", StringType, true, withId(2))
            .add("rand2", StringType, true, withId(1))

        val writeData = Seq(Row(100, "text", "txt"), Row(200, "more", "mr"))
        spark
          .createDataFrame(spark.sparkContext.parallelize(writeData), writeSchema)
          .write
          .mode("overwrite")
          .parquet(dir.getCanonicalPath)

        val cause = intercept[SparkException] {
          spark.read.schema(readSchema).parquet(dir.getCanonicalPath).collect()
        }.getCause
        assert(
          cause.isInstanceOf[RuntimeException] &&
            cause.getMessage.contains("Found duplicate field(s)"))
      }
    }
  }

  // Verbatim port of Spark `ParquetFieldIdIOSuite.test("read parquet file without ids")`,
  // pinned to `SCAN_NATIVE_DATAFUSION` for the same reason as the duplicate-id test above.
  test("read parquet file without ids") {
    withSQLConf(
      CometConf.COMET_NATIVE_SCAN_IMPL.key -> CometConf.SCAN_NATIVE_DATAFUSION,
      SQLConf.PARQUET_FIELD_ID_READ_ENABLED.key -> "true") {
      withTempPath { dir =>
        val readSchema =
          new StructType()
            .add("a", IntegerType, true, withId(1))

        val writeSchema =
          new StructType()
            .add("a", IntegerType, true)
            .add("rand1", StringType, true)
            .add("rand2", StringType, true)

        val writeData = Seq(Row(100, "text", "txt"), Row(200, "more", "mr"))
        spark
          .createDataFrame(spark.sparkContext.parallelize(writeData), writeSchema)
          .write
          .mode("overwrite")
          .parquet(dir.getCanonicalPath)

        Seq(readSchema, readSchema.add("b", StringType, true)).foreach { schema =>
          val cause = intercept[SparkException] {
            spark.read.schema(schema).parquet(dir.getCanonicalPath).collect()
          }.getCause
          assert(
            cause.isInstanceOf[RuntimeException] &&
              cause.getMessage.contains("Parquet file schema doesn't contain any field Ids"))
          val expectedValues = (1 to schema.length).map(_ => null)
          withSQLConf(SQLConf.IGNORE_MISSING_PARQUET_FIELD_ID.key -> "true") {
            checkAnswer(
              spark.read.schema(schema).parquet(dir.getCanonicalPath),
              Row(expectedValues: _*) :: Row(expectedValues: _*) :: Nil)
          }
        }
      }
    }
  }
}

class ParquetReadV1Suite extends ParquetReadSuite with AdaptiveSparkPlanHelper {
  override protected def test(testName: String, testTags: Tag*)(testFun: => Any)(implicit
      pos: Position): Unit = {
    super.test(testName, testTags: _*)(withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "parquet") {
      testFun
    })(pos)
  }

  override def checkParquetScan[T <: Product: ClassTag: TypeTag](
      data: Seq[T],
      f: Row => Boolean = _ => true): Unit = {
    withParquetDataFrame(data) { r =>
      val scans = collect(r.filter(f).queryExecution.executedPlan) {
        case p: CometScanExec =>
          p
        case p: CometNativeScanExec =>
          p
      }
      if (CometConf.COMET_ENABLED.get()) {
        assert(scans.nonEmpty)
      } else {
        assert(scans.isEmpty)
      }
    }
  }

  test("Test V1 parquet scan uses respective scanner") {
    Seq(
      ("false", CometConf.SCAN_NATIVE_DATAFUSION, "FileScan parquet"),
      ("true", CometConf.SCAN_NATIVE_DATAFUSION, "CometNativeScan"),
      ("true", CometConf.SCAN_NATIVE_ICEBERG_COMPAT, "CometScan [native_iceberg_compat] parquet"))
      .foreach { case (cometEnabled, cometNativeScanImpl, expectedScanner) =>
        testScanner(
          cometEnabled,
          cometNativeScanImpl,
          scanner = expectedScanner,
          v1 = Some("parquet"))
      }
  }

  test("test V1 parquet native scan -- case insensitive") {
    withTempPath { path =>
      spark.range(10).toDF("a").write.parquet(path.toString)
      Seq(CometConf.SCAN_NATIVE_DATAFUSION, CometConf.SCAN_NATIVE_ICEBERG_COMPAT).foreach(
        scanMode => {
          withSQLConf(CometConf.COMET_NATIVE_SCAN_IMPL.key -> scanMode) {
            withTable("test") {
              sql("create table test (A long) using parquet options (path '" + path + "')")
              val df = sql("select A from test")
              checkSparkAnswer(df)
              // TODO: pushed down filters do not used schema adapter in datafusion, will cause empty result
              // val df = sql("select * from test where A > 5")
              // checkSparkAnswer(df)
            }
          }
        })
    }
  }

  test("test V1 parquet scan filter pushdown of primitive types uses native_iceberg_compat") {
    withTempPath { dir =>
      val path = new Path(dir.toURI.toString, "test1.parquet")
      val rows = 1000
      withSQLConf(
        CometConf.COMET_NATIVE_SCAN_IMPL.key -> CometConf.SCAN_NATIVE_ICEBERG_COMPAT,
        CometConf.COMET_PARQUET_UNSIGNED_SMALL_INT_CHECK.key -> "true") {
        makeParquetFileAllPrimitiveTypes(
          path,
          dictionaryEnabled = false,
          0,
          rows,
          nullEnabled = false)
      }
      Seq(CometConf.SCAN_NATIVE_DATAFUSION, CometConf.SCAN_NATIVE_ICEBERG_COMPAT).foreach {
        scanMode =>
          Seq(true, false).foreach { pushDown =>
            breakable {
              withSQLConf(
                CometConf.COMET_NATIVE_SCAN_IMPL.key -> scanMode,
                SQLConf.PARQUET_FILTER_PUSHDOWN_ENABLED.key -> pushDown.toString) {
                Seq(
                  ("_1 = true", Math.ceil(rows.toDouble / 2)), // Boolean
                  ("_2 = 1", Math.ceil(rows.toDouble / 256)), // Byte
                  ("_3 = 1", 1), // Short
                  ("_4 = 1", 1), // Integer
                  ("_5 = 1", 1), // Long
                  ("_6 = 1.0", 1), // Float
                  ("_7 = 1.0", 1), // Double
                  (s"_8 = '${1.toString * 48}'", 1), // String
                  ("_21 = to_binary('1', 'utf-8')", 1), // binary
                  ("_15 = 0.0", 1), // DECIMAL(5, 2)
                  ("_16 = 0.0", 1), // DECIMAL(18, 10)
                  (
                    s"_17 = ${new BigDecimal(new BigInteger(("1" * 16).getBytes), 37).toString}",
                    Math.ceil(rows.toDouble / 10)
                  ), // DECIMAL(38, 37)
                  (s"_19 = TIMESTAMP '${DateTimeUtils.toJavaTimestamp(1)}'", 1), // Timestamp
                  ("_20 = DATE '1970-01-02'", 1) // Date
                ).foreach { case (whereCause, expectedRows) =>
                  val df = spark.read
                    .parquet(path.toString)
                    .where(whereCause)
                  val (_, cometPlan) = checkSparkAnswer(df)
                  val scan = collect(cometPlan) {
                    case p: CometScanExec =>
                      assert(p.dataFilters.nonEmpty)
                      p
                    case p: CometNativeScanExec =>
                      assert(p.dataFilters.nonEmpty)
                      p
                  }
                  assert(scan.size == 1)

                  if (pushDown) {
                    assert(scan.head.metrics("output_rows").value == expectedRows)
                  } else {
                    assert(scan.head.metrics("output_rows").value == rows)
                  }
                }
              }
            }
          }
      }
    }
  }

  test("read basic complex types") {
    Seq(true, false).foreach(dictionaryEnabled => {
      withTempPath { dir =>
        val path = new Path(dir.toURI.toString, "complex_types.parquet")
        makeParquetFileComplexTypes(path, dictionaryEnabled, 10)
        withParquetTable(path.toUri.toString, "complex_types") {
          Seq(CometConf.SCAN_NATIVE_DATAFUSION, CometConf.SCAN_NATIVE_ICEBERG_COMPAT).foreach(
            scanMode => {
              withSQLConf(CometConf.COMET_NATIVE_SCAN_IMPL.key -> scanMode) {
                checkSparkAnswerAndOperator(sql("select * from complex_types"))
                // First level
                checkSparkAnswerAndOperator(sql(
                  "select optional_array, array_of_struct, optional_map, complex_map from complex_types"))
                // second nested level
                checkSparkAnswerAndOperator(
                  sql(
                    "select optional_array[0], " +
                      "array_of_struct[0].field1, " +
                      "array_of_struct[0].optional_nested_array, " +
                      "optional_map.key, " +
                      "optional_map.value, " +
                      "map_keys(complex_map), " +
                      "map_entries(complex_map), " +
                      "map_values(complex_map) " +
                      "from complex_types"))
                // leaf fields
                checkSparkAnswerAndOperator(
                  sql(
                    "select optional_array[0], " +
                      "array_of_struct[0].field1, " +
                      "array_of_struct[0].optional_nested_array[0], " +
                      "optional_map.key, " +
                      "optional_map.value, " +
                      "map_keys(complex_map)[0].key_field1, " +
                      "map_keys(complex_map)[0].key_field2, " +
                      "map_entries(complex_map)[0].key, " +
                      "map_entries(complex_map)[0].value, " +
                      "map_values(complex_map)[0].value_field1, " +
                      "map_values(complex_map)[0].value_field2 " +
                      "from complex_types"))
              }
            })
        }
      }
    })
  }

  test("reading ancient dates before 1582") {
    // Verify that legacy dates (before 1582-10-15) are read without error.
    // Comet does not support datetime rebasing, so these dates are read as if they were
    // written using the Proleptic Gregorian calendar (no rebase, no exception).
    val file =
      getResourceParquetFilePath("test-data/before_1582_date_v3_2_0.snappy.parquet")

    Seq(CometConf.SCAN_NATIVE_ICEBERG_COMPAT, CometConf.SCAN_NATIVE_DATAFUSION).foreach {
      scanImpl =>
        withSQLConf(CometConf.COMET_NATIVE_SCAN_IMPL.key -> scanImpl) {
          val df = spark.read.parquet(file)

          // Verify Comet scan is in the plan
          val plan = df.queryExecution.executedPlan
          checkCometOperators(plan)

          // Verify all 8 rows are read and contain dates before 1582
          val rows = df.collect()
          assert(rows.length == 8, s"Expected 8 rows with $scanImpl, got ${rows.length}")
          rows.foreach { row =>
            val date = row.getDate(0)
            assert(
              date.toLocalDate.getYear < 1582,
              s"Expected date before 1582 with $scanImpl, got $date")
          }
        }
    }
  }

}
