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

import java.io.{ByteArrayOutputStream, PrintStream}
import java.nio.file.Files
import java.time.{DateTimeException, LocalTime}
import java.util.Properties

import scala.collection.mutable
import scala.concurrent.{ExecutionContext, Future}
import scala.concurrent.duration.{DurationInt, FiniteDuration}
import scala.jdk.CollectionConverters._

import io.grpc.{CallOptions, Channel, ClientCall, ClientInterceptor, ForwardingClientCall, ForwardingClientCallListener, MethodDescriptor}
import org.apache.commons.io.output.TeeOutputStream
import org.scalactic.TolerantNumerics
import org.scalatest.PrivateMethodTester

import org.apache.spark.{SparkArithmeticException, SparkException, SparkUpgradeException}
import org.apache.spark.SparkBuildInfo.{spark_version => SPARK_VERSION}
import org.apache.spark.connect.proto
import org.apache.spark.internal.config.ConfigBuilder
import org.apache.spark.sql.{functions, AnalysisException, Observation, Row, SaveMode}
import org.apache.spark.sql.catalyst.analysis.{NamespaceAlreadyExistsException, NoSuchNamespaceException, TableAlreadyExistsException, TempTableAlreadyExistsException}
import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.StringEncoder
import org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema
import org.apache.spark.sql.catalyst.parser.ParseException
import org.apache.spark.sql.connect.ConnectConversions._
import org.apache.spark.sql.connect.client.{PlanCompressionOptions, RetryPolicy, SparkConnectClient, SparkResult}
import org.apache.spark.sql.connect.test.{ConnectFunSuite, IntegrationTestUtils, QueryTest, RemoteSparkSession, SQLHelper}
import org.apache.spark.sql.connect.test.SparkConnectServerUtils.{createSparkSession, port}
import org.apache.spark.sql.functions._
import org.apache.spark.sql.internal.SqlApiConf
import org.apache.spark.sql.types._
import org.apache.spark.util.{SparkFileUtils, SparkSystemUtils, SparkThreadUtils}

class ClientE2ETestSuite
    extends QueryTest
    with ConnectFunSuite
    with RemoteSparkSession
    with SQLHelper
    with PrivateMethodTester {

  test("throw SparkException with null filename in stack trace elements") {
    withSQLConf("spark.sql.connect.enrichError.enabled" -> "true") {
      val session = spark
      import session.implicits._

      val throwException =
        udf((_: String) => {
          val testError = new SparkException("test")
          val stackTrace = testError.getStackTrace()
          stackTrace(0) = new StackTraceElement(
            stackTrace(0).getClassName,
            stackTrace(0).getMethodName,
            null,
            stackTrace(0).getLineNumber)
          testError.setStackTrace(stackTrace)
          throw testError
        })

      val ex = intercept[SparkException] {
        Seq("1").toDS().withColumn("udf_val", throwException($"value")).collect()
      }

      assert(ex.getCause.isInstanceOf[SparkException])
      assert(ex.getCause.getStackTrace().length > 0)
      assert(ex.getCause.getStackTrace()(0).getFileName == null)
    }
  }

  for (enrichErrorEnabled <- Seq(false, true)) {
    test(s"cause exception - ${enrichErrorEnabled}") {
      withSQLConf(
        "spark.sql.connect.enrichError.enabled" -> enrichErrorEnabled.toString,
        "spark.sql.legacy.timeParserPolicy" -> "EXCEPTION") {
        val ex = intercept[SparkUpgradeException] {
          spark
            .sql("""
                |select from_json(
                |  '{"d": "02-29"}',
                |  'd date',
                |  map('dateFormat', 'MM-dd'))
                |""".stripMargin)
            .collect()
        }
        assert(
          ex.getCondition ===
            "INCONSISTENT_BEHAVIOR_CROSS_VERSION.PARSE_DATETIME_BY_NEW_PARSER")
        assert(
          ex.getMessageParameters.asScala == Map(
            "datetime" -> "'02-29'",
            "config" -> "\"spark.sql.legacy.timeParserPolicy\""))
        if (enrichErrorEnabled) {
          assert(ex.getCause.isInstanceOf[DateTimeException])
        } else {
          assert(ex.getCause == null)
        }
      }
    }
  }

  test("throw SparkException with large cause exception") {
    withSQLConf("spark.sql.connect.enrichError.enabled" -> "true") {
      val session = spark
      import session.implicits._

      val throwException =
        udf((_: String) => throw new SparkException("test".repeat(10000)))

      val ex = intercept[SparkException] {
        Seq("1").toDS().withColumn("udf_val", throwException($"value")).collect()
      }

      assert(ex.getCondition != null)
      assert(!ex.getMessageParameters.isEmpty)
      assert(ex.getCause.isInstanceOf[SparkException])

      val cause = ex.getCause.asInstanceOf[SparkException]
      assert(cause.getCondition == null)
      assert(cause.getMessageParameters.isEmpty)
      assert(cause.getMessage.contains("test".repeat(10000)))
    }
  }

  for (isServerStackTraceEnabled <- Seq(false, true)) {
    test(s"server-side stack trace is set in exceptions - ${isServerStackTraceEnabled}") {
      withSQLConf(
        "spark.sql.connect.serverStacktrace.enabled" -> isServerStackTraceEnabled.toString,
        "spark.sql.pyspark.jvmStacktrace.enabled" -> "false") {
        val ex = intercept[AnalysisException] {
          spark.sql("select x").collect()
        }
        assert(ex.getCondition != null)
        assert(!ex.messageParameters.isEmpty)
        assert(ex.getSqlState != null)
        assert(!ex.isInternalError)
        assert(ex.getQueryContext.length == 1)
        assert(ex.getQueryContext.head.startIndex() == 7)
        assert(ex.getQueryContext.head.stopIndex() == 7)
        assert(ex.getQueryContext.head.fragment() == "x")
        assert(
          ex.getStackTrace
            .find(_.getClassName.contains("org.apache.spark.sql.catalyst.analysis.CheckAnalysis"))
            .isDefined)
      }
    }
  }

  test("throw SparkArithmeticException") {
    withSQLConf("spark.sql.ansi.enabled" -> "true") {
      intercept[SparkArithmeticException] {
        spark.sql("select 1/0").collect()
      }
    }
  }

  test("throw NoSuchNamespaceException") {
    val ex = intercept[NoSuchNamespaceException] {
      spark.sql("use database123")
    }
    assert(ex.getCondition != null)
  }

  test("table not found for spark.catalog.getTable") {
    val ex = intercept[AnalysisException] {
      spark.catalog.getTable("test_table")
    }
    assert(ex.getCondition != null)
  }

  test("throw NamespaceAlreadyExistsException") {
    try {
      spark.sql("create database test_db")
      val ex = intercept[NamespaceAlreadyExistsException] {
        spark.sql("create database test_db")
      }
      assert(ex.getCondition != null)
    } finally {
      spark.sql("drop database test_db")
    }
  }

  test("throw TempTableAlreadyExistsException") {
    try {
      spark.sql("create temporary view test_view as select 1")
      val ex = intercept[TempTableAlreadyExistsException] {
        spark.sql("create temporary view test_view as select 1")
      }
      assert(ex.getCondition != null)
    } finally {
      spark.sql("drop view test_view")
    }
  }

  test("throw TableAlreadyExistsException") {
    withTable("testcat.test_table") {
      spark.sql(s"create table testcat.test_table (id int)")
      val ex = intercept[TableAlreadyExistsException] {
        spark.sql(s"create table testcat.test_table (id int)")
      }
      assert(ex.getCondition != null)
    }
  }

  test("throw ParseException") {
    val ex = intercept[ParseException] {
      spark.sql("selet 1").collect()
    }
    assert(ex.getCondition != null)
    assert(!ex.messageParameters.isEmpty)
    assert(ex.getSqlState != null)
    assert(!ex.isInternalError)
  }

  test("spark deep recursion") {
    var recursionDepth = if (SparkSystemUtils.osArch == "s390x") 400 else 500
    var df = spark.range(1)
    for (a <- 1 to recursionDepth) {
      df = df.union(spark.range(a, a + 1))
    }
    assert(df.collect().length == recursionDepth + 1)
  }

  test("handle unknown exception") {
    var df = spark.range(1)
    val limit = spark.conf.get("spark.connect.grpc.marshallerRecursionLimit").toInt + 1
    for (a <- 1 to limit) {
      df = df.union(spark.range(a, a + 1))
    }
    val ex = intercept[SparkException] {
      df.collect()
    }
    assert(ex.getMessage.contains("io.grpc.StatusRuntimeException: UNKNOWN"))
  }

  test("many tables") {
    withSQLConf("spark.sql.execution.arrow.maxRecordsPerBatch" -> "10") {
      val numTables = 20
      try {
        for (i <- 0 to numTables) {
          spark.sql(s"create table testcat.table${i} (id int)")
        }
        assert(spark.sql("show tables in testcat").collect().length == numTables + 1)
      } finally {
        for (i <- 0 to numTables) {
          spark.sql(s"drop table if exists testcat.table${i}")
        }
      }
    }
  }

  // Spark Result
  test("spark result schema") {
    val df = spark.sql("select val from (values ('Hello'), ('World')) as t(val)")
    df.withResult { result =>
      val schema = result.schema
      assert(schema == StructType(StructField("val", StringType, nullable = false) :: Nil))
    }
  }

  test("spark result array") {
    val df = spark.sql("select val from (values ('Hello'), ('World')) as t(val)")
    val result = df.collect()
    assert(result.length == 2)
    assert(result(0).getString(0) == "Hello")
    assert(result(1).getString(0) == "World")
  }

  test("eager execution of sql") {
    assume(IntegrationTestUtils.isSparkHiveJarAvailable)
    withTable("test_martin") {
      // Fails, because table does not exist.
      assertThrows[AnalysisException] {
        spark.sql("select * from test_martin").collect()
      }
      // Execute eager, DML
      spark.sql("create table test_martin (id int)")
      // Execute read again.
      val rows = spark.sql("select * from test_martin").collect()
      assert(rows.length == 0)
      spark.sql("insert into test_martin values (1), (2)")
      val rows_new = spark.sql("select * from test_martin").collect()
      assert(rows_new.length == 2)
    }
  }

  test("simple dataset") {
    val df = spark.range(10).limit(3)
    val result = df.collect()
    assert(result.length == 3)
    assert(result(0) == 0)
    assert(result(1) == 1)
    assert(result(2) == 2)
  }

  test("read and write") {
    assume(IntegrationTestUtils.isSparkHiveJarAvailable)
    val testDataPath = java.nio.file.Paths
      .get(
        IntegrationTestUtils.sparkHome,
        "sql",
        "connect",
        "common",
        "src",
        "test",
        "resources",
        "query-tests",
        "test-data",
        "people.csv")
      .toAbsolutePath
    val df = spark.read
      .format("csv")
      .option("path", testDataPath.toString)
      .options(Map("header" -> "true", "delimiter" -> ";"))
      .schema(
        StructType(
          StructField("name", StringType) ::
            StructField("age", IntegerType) ::
            StructField("job", StringType) :: Nil))
      .load()
    val outputFolderPath = Files.createTempDirectory("output").toAbsolutePath

    df.write
      .format("csv")
      .mode("overwrite")
      .options(Map("header" -> "true", "delimiter" -> ";"))
      .save(outputFolderPath.toString)

    // We expect only one csv file saved.
    val outputFile = outputFolderPath.toFile
      .listFiles()
      .filter(file => file.getPath.endsWith(".csv"))(0)

    assert(SparkFileUtils.contentEquals(testDataPath.toFile, outputFile))
  }

  test("read path collision") {
    val testDataPath = java.nio.file.Paths
      .get(
        IntegrationTestUtils.sparkHome,
        "sql",
        "connect",
        "common",
        "src",
        "test",
        "resources",
        "query-tests",
        "test-data",
        "people.csv")
      .toAbsolutePath
    val df = spark.read
      .format("csv")
      .option("path", testDataPath.toString)
      .options(Map("header" -> "true", "delimiter" -> ";"))
      .schema(
        StructType(
          StructField("name", StringType) ::
            StructField("age", IntegerType) ::
            StructField("job", StringType) :: Nil))
      .csv(testDataPath.toString)
    // Failed because the path cannot be provided both via option and load method (csv).
    assertThrows[AnalysisException] {
      df.collect()
    }
  }

  test("textFile") {
    assume(IntegrationTestUtils.isSparkHiveJarAvailable)
    val testDataPath = java.nio.file.Paths
      .get(
        IntegrationTestUtils.sparkHome,
        "sql",
        "connect",
        "common",
        "src",
        "test",
        "resources",
        "query-tests",
        "test-data",
        "people.txt")
      .toAbsolutePath
    val result = spark.read.textFile(testDataPath.toString).collect()
    val expected = Array("Michael, 29", "Andy, 30", "Justin, 19")
    assert(result.length == 3)
    assert(result === expected)
  }

  test("write table") {
    assume(IntegrationTestUtils.isSparkHiveJarAvailable)
    withTable("myTable") {
      val df = spark.range(10).limit(3)
      df.write.mode(SaveMode.Overwrite).saveAsTable("myTable")
      spark.range(2).write.insertInto("myTable")
      val result = spark.sql("select * from myTable").sort("id").collect()
      assert(result.length == 5)
      assert(result(0).getLong(0) == 0)
      assert(result(1).getLong(0) == 0)
      assert(result(2).getLong(0) == 1)
      assert(result(3).getLong(0) == 1)
      assert(result(4).getLong(0) == 2)
    }
  }

  test("different spark session join/union") {
    val df = spark.range(10).limit(3)

    val spark2 = SparkSession
      .builder()
      .client(
        SparkConnectClient
          .builder()
          .port(port)
          .build())
      .create()

    val df2 = spark2.range(10).limit(3)

    assertThrows[SparkException] {
      df.union(df2).collect()
    }

    assertThrows[SparkException] {
      df.unionByName(df2).collect()
    }

    assertThrows[SparkException] {
      df.join(df2).collect()
    }

  }

  test("write without table or path") {
    assume(IntegrationTestUtils.isSparkHiveJarAvailable)
    // Should receive no error to write noop
    spark.range(10).write.format("noop").mode("append").save()
  }

  test("write jdbc") {
    assume(IntegrationTestUtils.isSparkHiveJarAvailable)
    val url = "jdbc:derby:memory:1234"
    val table = "t1"
    try {
      spark.range(10).write.jdbc(url = s"$url;create=true", table, new Properties())
      val result = spark.read.jdbc(url = url, table, new Properties()).collect()
      assert(result.length == 10)
    } finally {
      // clean up
      assertThrows[SparkException] {
        spark.read.jdbc(url = s"$url;drop=true", table, new Properties()).collect()
      }
    }
  }

  test("writeTo with create") {
    withTable("testcat.myTableV2") {

      val rows = Seq(Row(1L, "a"), Row(2L, "b"), Row(3L, "c"))

      val schema = StructType(Array(StructField("id", LongType), StructField("data", StringType)))

      spark.createDataFrame(rows.asJava, schema).writeTo("testcat.myTableV2").create()

      val outputRows = spark.table("testcat.myTableV2").collect()
      assert(outputRows.length == 3)
    }
  }

  test("writeTo with create and using") {
    withTable("testcat.myTableV2") {
      val rows = Seq(Row(1L, "a"), Row(2L, "b"), Row(3L, "c"))

      val schema = StructType(Array(StructField("id", LongType), StructField("data", StringType)))

      spark.createDataFrame(rows.asJava, schema).writeTo("testcat.myTableV2").create()
      val outputRows = spark.table("testcat.myTableV2").collect()
      assert(outputRows.length == 3)

      val columns = spark.table("testcat.myTableV2").columns
      assert(columns.length == 2)

      val sqlOutputRows = spark.sql("select * from testcat.myTableV2").collect()
      assert(outputRows.length == 3)
      assert(sqlOutputRows(0).schema == schema)
      assert(sqlOutputRows(1).getString(1) == "b")
    }
  }

  test("writeTo with create and append") {
    withTable("testcat.myTableV2") {

      val rows = Seq(Row(1L, "a"), Row(2L, "b"), Row(3L, "c"))

      val schema = StructType(Array(StructField("id", LongType), StructField("data", StringType)))

      spark.sql("CREATE TABLE testcat.myTableV2 (id bigint, data string) USING foo")

      assert(spark.table("testcat.myTableV2").collect().isEmpty)

      spark.createDataFrame(rows.asJava, schema).writeTo("testcat.myTableV2").append()
      val outputRows = spark.table("testcat.myTableV2").collect()
      assert(outputRows.length == 3)
    }
  }

  test("WriteTo with overwrite") {
    withTable("testcat.myTableV2") {

      val rows1 = (1L to 3L).map { i =>
        Row(i, "" + (i - 1 + 'a'))
      }
      val rows2 = (4L to 7L).map { i =>
        Row(i, "" + (i - 1 + 'a'))
      }

      val schema = StructType(Array(StructField("id", LongType), StructField("data", StringType)))

      spark.sql(
        "CREATE TABLE testcat.myTableV2 (id bigint, data string) USING foo PARTITIONED BY (id)")

      assert(spark.table("testcat.myTableV2").collect().isEmpty)

      spark.createDataFrame(rows1.asJava, schema).writeTo("testcat.myTableV2").append()
      val outputRows = spark.table("testcat.myTableV2").collect()
      assert(outputRows.length == 3)

      spark
        .createDataFrame(rows2.asJava, schema)
        .writeTo("testcat.myTableV2")
        .overwrite(functions.expr("true"))
      val outputRows2 = spark.table("testcat.myTableV2").collect()
      assert(outputRows2.length == 4)

    }
  }

  test("WriteTo with overwritePartitions") {
    withTable("testcat.myTableV2") {

      val rows = (4L to 7L).map { i =>
        Row(i, "" + (i - 1 + 'a'))
      }

      val schema = StructType(Array(StructField("id", LongType), StructField("data", StringType)))

      spark.sql(
        "CREATE TABLE testcat.myTableV2 (id bigint, data string) USING foo PARTITIONED BY (id)")

      assert(spark.table("testcat.myTableV2").collect().isEmpty)

      spark
        .createDataFrame(rows.asJava, schema)
        .writeTo("testcat.myTableV2")
        .overwritePartitions()
      val outputRows = spark.table("testcat.myTableV2").collect()
      assert(outputRows.length == 4)

    }
  }

  test("write path collision") {
    val df = spark.range(10)
    val outputFolderPath = Files.createTempDirectory("output").toAbsolutePath
    // Failed because the path cannot be provided both via option and save method.
    assertThrows[AnalysisException] {
      df.write.option("path", outputFolderPath.toString).save(outputFolderPath.toString)
    }
  }

  // TODO test large result when we can create table or view
  // test("test spark large result")
  private def captureStdOut(block: => Unit): String = {
    val currentOut = Console.out
    val capturedOut = new ByteArrayOutputStream()
    val newOut = new PrintStream(new TeeOutputStream(currentOut, capturedOut))
    Console.withOut(newOut) {
      block
    }
    capturedOut.toString
  }

  private def checkFragments(result: String, fragmentsToCheck: Seq[String]): Unit = {
    fragmentsToCheck.foreach { fragment =>
      assert(result.contains(fragment))
    }
  }

  private def testCapturedStdOut(block: => Unit, fragmentsToCheck: String*): Unit = {
    checkFragments(captureStdOut(block), fragmentsToCheck)
  }

  private def testCapturedStdOut(
      block: => Unit,
      expectedNumLines: Int,
      expectedMaxWidth: Int,
      fragmentsToCheck: String*): Unit = {
    val result = captureStdOut(block)
    val lines = result.split('\n')
    assert(lines.length === expectedNumLines)
    assert(lines.map((s: String) => s.length).max <= expectedMaxWidth)
    checkFragments(result, fragmentsToCheck)
  }

  private val simpleSchema = new StructType().add("id", "long", nullable = false)

  // Dataset tests
  test("Dataset inspection") {
    val df = spark.range(10)
    val local = spark.newDataFrame { builder =>
      builder.getLocalRelationBuilder.setSchema(simpleSchema.catalogString)
    }
    assert(!df.isLocal)
    assert(local.isLocal)
    assert(!df.isStreaming)
    assert(df.toString.contains("[id: bigint]"))
    assert(df.inputFiles.isEmpty)
  }

  test("Dataset schema") {
    val df = spark.range(10)
    assert(df.schema === simpleSchema)
    assert(df.dtypes === Array(("id", "LongType")))
    assert(df.columns === Array("id"))
    testCapturedStdOut(df.printSchema(), simpleSchema.treeString)
    testCapturedStdOut(df.printSchema(5), simpleSchema.treeString(5))
  }

  test("Dataframe schema") {
    val df = spark.sql("select * from range(10)")
    val expectedSchema = new StructType().add("id", "long", nullable = false)
    assert(df.schema === expectedSchema)
    assert(df.dtypes === Array(("id", "LongType")))
    assert(df.columns === Array("id"))
    testCapturedStdOut(df.printSchema(), expectedSchema.treeString)
    testCapturedStdOut(df.printSchema(5), expectedSchema.treeString(5))
  }

  test("Dataset explain") {
    val df = spark.range(10)
    val simpleExplainFragments = Seq("== Physical Plan ==")
    testCapturedStdOut(df.explain(), simpleExplainFragments: _*)
    testCapturedStdOut(df.explain(false), simpleExplainFragments: _*)
    testCapturedStdOut(df.explain("simple"), simpleExplainFragments: _*)
    val extendedExplainFragments = Seq(
      "== Parsed Logical Plan ==",
      "== Analyzed Logical Plan ==",
      "== Optimized Logical Plan ==") ++
      simpleExplainFragments
    testCapturedStdOut(df.explain(true), extendedExplainFragments: _*)
    testCapturedStdOut(df.explain("extended"), extendedExplainFragments: _*)
    testCapturedStdOut(
      df.explain("cost"),
      simpleExplainFragments :+ "== Optimized Logical Plan ==": _*)
    testCapturedStdOut(df.explain("codegen"), "WholeStageCodegen subtrees.")
    testCapturedStdOut(df.explain("formatted"), "Range", "Arguments: ")
  }

  test("Dataset result collection") {
    def checkResult(rows: IterableOnce[java.lang.Long], expectedValues: Long*): Unit = {
      rows.iterator.zipAll(expectedValues.iterator, null, null).foreach {
        case (actual, expected) => assert(actual === expected)
      }
    }
    val df = spark.range(10)
    checkResult(df.head() :: Nil, 0L)
    checkResult(df.head(5), 0L, 1L, 2L, 3L, 4L)
    checkResult(df.first() :: Nil, 0L)
    assert(!df.isEmpty)
    assert(df.filter("id > 100").isEmpty)
    checkResult(df.take(3), 0L, 1L, 2L)
    checkResult(df.tail(3), 7L, 8L, 9L)
    checkResult(df.takeAsList(10).asScala, 0L, 1L, 2L, 3L, 4L, 5L, 6L, 7L, 8L, 9L)
    checkResult(df.filter("id % 3 = 0").collect(), 0L, 3L, 6L, 9L)
    checkResult(df.filter("id < 3").collectAsList().asScala, 0L, 1L, 2L)
    val iterator = df.filter("id > 5 and id < 9").toLocalIterator()
    try {
      checkResult(iterator.asScala, 6L, 7L, 8L)
    } finally {
      iterator.asInstanceOf[AutoCloseable].close()
    }
  }

  test("Dataset show") {
    val df = spark.range(20)
    testCapturedStdOut(df.show(), 24, 5, "+---+", "| id|", "|  0|", "| 19|")
    testCapturedStdOut(
      df.show(10),
      15,
      24,
      "+---+",
      "| id|",
      "|  0|",
      "|  9|",
      "only showing top 10 rows")
    val wideDf =
      spark.range(4).selectExpr("id", "concat('very_very_very_long_string', id) as val")
    testCapturedStdOut(
      wideDf.show(true),
      8,
      26,
      "+---+--------------------+",
      "| id|                 val|",
      "|  0|very_very_very_lo...|")
    testCapturedStdOut(
      wideDf.show(false),
      8,
      33,
      "+---+---------------------------+",
      "|id |val                        |",
      "|2  |very_very_very_long_string2|")
    testCapturedStdOut(
      wideDf.show(2, truncate = false),
      7,
      33,
      "+---+---------------------------+",
      "|id |val                        |",
      "|1  |very_very_very_long_string1|",
      "only showing top 2 rows")
    testCapturedStdOut(
      df.show(8, 10, vertical = true),
      17,
      23,
      "-RECORD 3--",
      "id  | 7",
      "only showing top 8 rows")
  }

  test("Dataset randomSplit") {
    implicit val tolerance = TolerantNumerics.tolerantDoubleEquality(0.01)

    val df = spark.range(100)
    def checkSample(
        ds: Dataset[java.lang.Long],
        lower: Double,
        upper: Double,
        seed: Long): Unit = {
      assert(ds.plan.getRoot.hasSample)
      val sample = ds.plan.getRoot.getSample
      assert(sample.getSeed === seed)
      assert(sample.getLowerBound === lower)
      assert(sample.getUpperBound === upper)
    }
    val Array(ds1, ds2, ds3) = df.randomSplit(Array(8, 9, 7), 123L)
    checkSample(ds1, 0, 8.0 / 24.0, 123L)
    checkSample(ds2, 8.0 / 24.0, 17.0 / 24.0, 123L)
    checkSample(ds3, 17.0 / 24.0, 1.0, 123L)

    val datasets = df.randomSplitAsList(Array(1, 2, 3, 4), 9L)
    assert(datasets.size() === 4)
    checkSample(datasets.get(0), 0, 1.0 / 10.0, 9L)
    checkSample(datasets.get(1), 1.0 / 10.0, 3.0 / 10.0, 9L)
    checkSample(datasets.get(2), 3.0 / 10.0, 6.0 / 10.0, 9L)
    checkSample(datasets.get(3), 6.0 / 10.0, 1.0, 9L)
  }

  test("Dataset count") {
    assert(spark.range(10).count() === 10)
  }

  test("Dataset collect tuple") {
    val session = spark
    import session.implicits._
    val result = session
      .range(3)
      .select(col("id"), (col("id") % 2).cast("int").as("a"), (col("id") / lit(10.0d)).as("b"))
      .as[(Long, Int, Double)]
      .collect()
    result.zipWithIndex.foreach { case ((id, a, b), i) =>
      assert(id == i)
      assert(a == id % 2)
      assert(b == id / 10.0d)
    }
  }

  private val generateMyTypeColumns = Seq(
    (col("id") / lit(10.0d)).as("b"),
    col("id"),
    lit("world").as("d"),
    (col("id") % 2).as("a"))

  private def validateMyTypeResult(result: Array[MyType]): Unit = {
    result.zipWithIndex.foreach { case (MyType(id, a, b), i) =>
      assert(id == i)
      assert(a == id % 2)
      assert(b == id / 10.0d)
    }
  }

  private def validateMyTypeResult(result: Array[(MyType, MyType, MyType)]): Unit = {
    result.zipWithIndex.foreach { case (row, i) =>
      val t1 = row._1
      val t2 = row._2
      val t3 = row._3
      assert(t1 === t2)
      assert(t2 === t3)
      assert(t1.id == i)
      assert(t1.a == t1.id % 2)
      assert(t1.b == t1.id / 10.0d)
    }
  }

  test("Dataset collect complex type") {
    val session = spark
    import session.implicits._
    val result = session
      .range(3)
      .select(generateMyTypeColumns: _*)
      .as[MyType]
      .collect()
    validateMyTypeResult(result)
  }

  test("Dataset typed select - simple column") {
    val numRows = spark.range(1000).select(count("id")).first()
    assert(numRows === 1000)
  }

  test("Dataset typed select - multiple columns") {
    val result = spark.range(1000).select(count("id"), sum("id")).first()
    assert(result.getLong(0) === 1000)
    assert(result.getLong(1) === 499500)
  }

  test("Dataset typed select - complex column") {
    val session = spark
    import session.implicits._
    val ds = session
      .range(3)
      .select(struct(generateMyTypeColumns: _*).as[MyType])
    validateMyTypeResult(ds.collect())
  }

  test("Dataset typed select - multiple complex columns") {
    val session = spark
    import session.implicits._
    val s = struct(generateMyTypeColumns: _*).as[MyType]
    val ds = session
      .range(3)
      .select(s, s, s)
    validateMyTypeResult(ds.collect())
  }

  test("lambda functions") {
    // This test is mostly to validate lambda variables are properly resolved.
    val result = spark
      .range(3)
      .select(
        col("id"),
        array(sequence(col("id"), lit(10)), sequence(col("id") * 2, lit(10))).as("data"))
      .select(col("id"), transform(col("data"), x => transform(x, x => x + 1)).as("data"))
      .select(
        col("id"),
        transform(col("data"), x => aggregate(x, lit(0L), (x, y) => x + y)).as("summaries"))
      .collect()
    val expected = Array(Row(0L, Seq(66L, 66L)), Row(1L, Seq(65L, 63L)), Row(2L, Seq(63L, 56L)))
    assert(result === expected)
  }

  test("shuffle array") {
    // We cannot do structural tests for shuffle because its random seed will always change.
    val result = spark
      .sql("select 1")
      .select(shuffle(array(lit(1), lit(2), lit(3), lit(74))))
      .head()
      .getSeq[Int](0)
    assert(result.toSet === Set(1, 2, 3, 74))
  }

  test("ambiguous joins") {
    val left = spark.range(100).select(col("id"), rand(10).as("a"))
    val right = spark.range(100).select(col("id"), rand(12).as("a"))
    val joined = left.join(right, left("id") === right("id")).select(left("id"), right("a"))
    assert(joined.schema.catalogString === "struct<id:bigint,a:double>")

    val joined2 = left
      .join(right, left.colRegex("id") === right.colRegex("id"))
      .select(left("id"), right("a"))
    assert(joined2.schema.catalogString === "struct<id:bigint,a:double>")
  }

  test("join with dataframe star") {
    val left = spark.range(100)
    val right = spark.range(100).select(col("id"), rand(12).as("a"))
    val join1 = left.join(right, left("id") === right("id"))
    assert(
      join1.select(join1.col("*")).schema.catalogString ===
        "struct<id:bigint,id:bigint,a:double>")
    assert(join1.select(left.col("*")).schema.catalogString === "struct<id:bigint>")
    assert(join1.select(right.col("*")).schema.catalogString === "struct<id:bigint,a:double>")

    val join2 = left.join(right)
    assert(
      join2.select(join2.col("*")).schema.catalogString ===
        "struct<id:bigint,id:bigint,a:double>")
    assert(join2.select(left.col("*")).schema.catalogString === "struct<id:bigint>")
    assert(join2.select(right.col("*")).schema.catalogString === "struct<id:bigint,a:double>")

    val join3 = left.join(right, "id")
    assert(
      join3.select(join3.col("*")).schema.catalogString ===
        "struct<id:bigint,a:double>")
    assert(join3.select(left.col("*")).schema.catalogString === "struct<id:bigint>")
    assert(join3.select(right.col("*")).schema.catalogString === "struct<id:bigint,a:double>")
  }

  test("SPARK-45509: ambiguous column reference") {
    val session = spark
    import session.implicits._
    val df1 = Seq(1 -> "a").toDF("i", "j")
    val df1_filter = df1.filter(df1("i") > 0)
    val df2 = Seq(2 -> "b").toDF("i", "y")

    checkSameResult(
      Seq(Row(1)),
      // df1("i") is not ambiguous, and it's still valid in the filtered df.
      df1_filter.select(df1("i")))

    val e1 = intercept[AnalysisException] {
      // df1("i") is not ambiguous, but it's not valid in the projected df.
      df1.select((df1("i") + 1).as("plus")).select(df1("i")).collect()
    }
    assert(e1.getMessage.contains("UNRESOLVED_COLUMN.WITH_SUGGESTION"))

    checkSameResult(
      Seq(Row(1, "a")),
      // All these column references are not ambiguous and are still valid after join.
      df1.join(df2, df1("i") + 1 === df2("i")).sort(df1("i").desc).select(df1("i"), df1("j")))

    val e2 = intercept[AnalysisException] {
      // df1("i") is ambiguous as df1 appears in both join sides.
      df1.join(df1, df1("i") === 1).collect()
    }
    assert(e2.getMessage.contains("AMBIGUOUS_COLUMN_REFERENCE"))

    val e3 = intercept[AnalysisException] {
      // df1("i") is ambiguous as df1 appears in both join sides.
      df1.join(df1).select(df1("i")).collect()
    }
    assert(e3.getMessage.contains("AMBIGUOUS_COLUMN_REFERENCE"))

    // TODO(SPARK-47749): Dataframe.collect should accept duplicated column names
    assert(
      // df1.join(df1_filter, df1("i") === 1) fails in classic spark due to:
      // org.apache.spark.sql.AnalysisException: Column i#24 are ambiguous
      df1.join(df1_filter, df1("i") === 1).columns ===
        Array("i", "j", "i", "j"))

    checkSameResult(
      Seq(Row("a")),
      // df1_filter("i") is not ambiguous as df1_filter does not exist in the join left side.
      df1.join(df1_filter, df1_filter("i") === 1).select(df1_filter("j")))

    val e5 = intercept[AnalysisException] {
      // df1("i") is ambiguous as df1 appears in both sides of the first join.
      df1.join(df1_filter, df1_filter("i") === 1).join(df2, df1("i") === 1).collect()
    }
    assert(e5.getMessage.contains("AMBIGUOUS_COLUMN_REFERENCE"))

    checkSameResult(
      Seq(Row("a")),
      // df1_filter("i") is not ambiguous as df1_filter only appears once.
      df1.join(df1_filter).join(df2, df1_filter("i") === 1).select(df1_filter("j")))
  }

  test("broadcast join") {
    withSQLConf("spark.sql.autoBroadcastJoinThreshold" -> "-1") {
      val left = spark.range(100).select(col("id"), rand(10).as("a"))
      val right = spark.range(100).select(col("id"), rand(12).as("a"))
      val joined =
        left.join(broadcast(right), left("id") === right("id")).select(left("id"), right("a"))
      assert(joined.schema.catalogString === "struct<id:bigint,a:double>")
      testCapturedStdOut(joined.explain(), "BroadcastHashJoin")
    }
  }

  test("test temp view") {
    try {
      spark.range(100).createTempView("test1")
      assert(spark.sql("SELECT * FROM test1").count() == 100)
      spark.range(1000).createOrReplaceTempView("test1")
      assert(spark.sql("SELECT * FROM test1").count() == 1000)
      spark.range(100).createGlobalTempView("view1")
      assert(spark.sql("SELECT * FROM global_temp.view1").count() == 100)
      spark.range(1000).createOrReplaceGlobalTempView("view1")
      assert(spark.sql("SELECT * FROM global_temp.view1").count() == 1000)
    } finally {
      spark.sql("DROP VIEW IF EXISTS test1")
      spark.sql("DROP VIEW IF EXISTS global_temp.view1")
    }
  }

  test("time") {
    val timeFragments = Seq("Time taken: ", " ms")
    testCapturedStdOut(spark.time(spark.sql("select 1").collect()), timeFragments: _*)
  }

  test("RuntimeConfig") {
    intercept[NoSuchElementException](spark.conf.get("foo.bar"))
    assert(spark.conf.getOption("foo.bar").isEmpty)
    assert(spark.conf.get("foo.bar", "nope") == "nope")
    assert(spark.conf.get("foo.bar", null) == null)
    spark.conf.set("foo.bar", value = true)
    assert(spark.conf.getOption("foo.bar") === Option("true"))
    assert(spark.conf.get("foo.bar", "nope") === "true")
    assert(spark.conf.get("foo.bar", null) === "true")
    spark.conf.set("foo.bar.numBaz", 100L)
    assert(spark.conf.get("foo.bar.numBaz") === "100")
    spark.conf.set("foo.bar.name", "donkey")
    assert(spark.conf.get("foo.bar.name") === "donkey")
    spark.conf.unset("foo.bar.name")
    val allKeyValues = spark.conf.getAll
    assert(allKeyValues("foo.bar") === "true")
    assert(allKeyValues("foo.bar.numBaz") === "100")
    assert(!spark.conf.isModifiable("foo.bar")) // This is a bit odd.
    assert(spark.conf.isModifiable("spark.sql.ansi.enabled"))
    assert(!spark.conf.isModifiable("spark.sql.globalTempDatabase"))
    intercept[Exception](spark.conf.set("spark.sql.globalTempDatabase", "/dev/null"))

    val entry = ConfigBuilder("my.simple.conf").intConf.createOptional
    intercept[NoSuchElementException](spark.conf.get(entry.key))
    assert(spark.conf.get(entry).isEmpty)
    assert(spark.conf.get(entry, Option(55)) === Option(55))
    spark.conf.set(entry, Option(33))
    assert(spark.conf.get(entry.key) === "33")
    assert(spark.conf.get(entry) === Option(33))
    assert(spark.conf.get(entry, Option(55)) === Option(33))

    val entryWithDefault = ConfigBuilder("my.important.conf").intConf.createWithDefault(10)
    intercept[NoSuchElementException](spark.conf.get(entryWithDefault.key))
    assert(spark.conf.get(entryWithDefault) === 10)
    assert(spark.conf.get(entryWithDefault, 11) === 11)
    spark.conf.set(entryWithDefault, 12)
    assert(spark.conf.get(entryWithDefault.key) === "12")
    assert(spark.conf.get(entryWithDefault) === 12)
    assert(spark.conf.get(entryWithDefault, 11) === 12)

    assert(spark.conf.contains(entryWithDefault.key))
    assert(!spark.conf.contains("nope"))
  }

  test("SparkVersion") {
    assert(spark.version.nonEmpty)
    assert(spark.version == SPARK_VERSION)
  }

  private def checkSameResult[E](expected: scala.collection.Seq[E], dataset: Dataset[E]): Unit = {
    dataset.withResult { result =>
      val iter = result.iterator
      try {
        assert(expected === iter.toBuffer)
      } finally {
        iter.close()
      }
    }
  }

  test("Local Relation implicit conversion") {
    val session = spark
    import session.implicits._

    val simpleValues = Seq(1, 24, 3)
    checkSameResult(simpleValues, simpleValues.toDS())
    checkSameResult(simpleValues.map(v => Row(v)), simpleValues.toDF())

    val complexValues = Seq((5, "a"), (6, "b"))
    checkSameResult(complexValues, complexValues.toDS())
    checkSameResult(
      complexValues.map(kv => KV(kv._2, kv._1)),
      complexValues.toDF("value", "key").as[KV])
  }

  test("SparkSession.createDataFrame - row") {
    val rows = java.util.Arrays.asList(Row("bob", 99), Row("Club", 5), Row("Bag", 5))
    val schema = new StructType().add("key", "string").add("value", "int")
    checkSameResult(rows.asScala, spark.createDataFrame(rows, schema))
  }

  test("SparkSession.createDataFrame - bean") {
    def bean(v: String): SimpleBean = {
      val bean = new SimpleBean
      bean.setValue(v)
      bean
    }
    val beans = java.util.Arrays.asList(bean("x"), bean("s"), bean("d"))
    checkSameResult(
      beans.asScala.map(b => Row(b.value)),
      spark.createDataFrame(beans, classOf[SimpleBean]))
  }

  test("SparkSession typed createDataSet/createDataframe") {
    val session = spark
    import session.implicits._
    val list = java.util.Arrays.asList(KV("bob", 99), KV("Club", 5), KV("Bag", 5))
    checkSameResult(list.asScala, session.createDataset(list))
    checkSameResult(
      list.asScala.map(kv => Row(kv.key, kv.value)),
      session.createDataFrame(list.asScala.toSeq))
  }

  test("SparkSession newSession") {
    val oldId = spark.sql("SELECT 1").analyze.getSessionId
    val newId = spark.newSession().sql("SELECT 1").analyze.getSessionId
    assert(oldId != newId)
  }

  test("createDataFrame from complex type schema") {
    val schema = new StructType()
      .add(
        "c1",
        new StructType()
          .add("c1-1", StringType)
          .add("c1-2", StringType))
    val data = Seq(Row(Row(null, "a2")), Row(Row("b1", "b2")), Row(null))
    val result = spark.createDataFrame(data.asJava, schema).collect()
    assert(result === data)
  }

  test("SameSemantics") {
    val plan = spark.sql("select 1")
    val otherPlan = spark.sql("select 1")
    assert(plan.sameSemantics(otherPlan))
  }

  test("sameSemantics and semanticHash") {
    val df1 = spark.createDataFrame(Seq((1, 2), (4, 5)))
    val df2 = spark.createDataFrame(Seq((1, 2), (4, 5)))
    val df3 = spark.createDataFrame(Seq((0, 2), (4, 5)))
    val df4 = spark.createDataFrame(Seq((0, 2), (4, 5)))

    assert(df1.sameSemantics(df2) === true)
    assert(df1.sameSemantics(df3) === false)
    assert(df3.sameSemantics(df4) === true)

    assert(df1.semanticHash() === df2.semanticHash())
    assert(df1.semanticHash() !== df3.semanticHash())
    assert(df3.semanticHash() === df4.semanticHash())
  }

  test("toJSON") {
    val expected = Array(
      """{"b":0.0,"id":0,"d":"world","a":0}""",
      """{"b":0.1,"id":1,"d":"world","a":1}""",
      """{"b":0.2,"id":2,"d":"world","a":0}""")
    val result = spark
      .range(3)
      .select(generateMyTypeColumns: _*)
      .toJSON
      .collect()
    assert(result sameElements expected)
  }

  test("json from Dataset[String] inferSchema") {
    val session = spark
    import session.implicits._
    val expected = Seq(
      new GenericRowWithSchema(
        Array(73, "Shandong", "Kong"),
        new StructType().add("age", LongType).add("city", StringType).add("name", StringType)))
    val ds = Seq("""{"name":"Kong","age":73,"city":'Shandong'}""").toDS()
    val result = spark.read.option("allowSingleQuotes", "true").json(ds)
    checkSameResult(expected, result)
  }

  test("json from Dataset[String] with schema") {
    val session = spark
    import session.implicits._
    val schema = new StructType().add("city", StringType).add("name", StringType)
    val expected = Seq(new GenericRowWithSchema(Array("Shandong", "Kong"), schema))
    val ds = Seq("""{"name":"Kong","age":73,"city":'Shandong'}""").toDS()
    val result = spark.read.schema(schema).option("allowSingleQuotes", "true").json(ds)
    checkSameResult(expected, result)
  }

  test("json from Dataset[String] with invalid schema") {
    val message = intercept[ParseException] {
      spark.read.schema("123").json(spark.createDataset(Seq.empty[String])(StringEncoder))
    }.getMessage
    assert(message.contains("PARSE_SYNTAX_ERROR"))
  }

  test("csv from Dataset[String] inferSchema") {
    val session = spark
    import session.implicits._
    val expected = Seq(
      new GenericRowWithSchema(
        Array("Meng", 84, "Shandong"),
        new StructType().add("name", StringType).add("age", LongType).add("city", StringType)))
    val ds = Seq("name,age,city", """"Meng",84,"Shandong"""").toDS()
    val result = spark.read
      .option("header", "true")
      .option("inferSchema", "true")
      .csv(ds)
    checkSameResult(expected, result)
  }

  test("csv from Dataset[String] with schema") {
    val session = spark
    import session.implicits._
    val schema = new StructType().add("name", StringType).add("age", LongType)
    val expected = Seq(new GenericRowWithSchema(Array("Meng", 84), schema))
    val ds = Seq(""""Meng",84,"Shandong"""").toDS()
    val result = spark.read.schema(schema).csv(ds)
    checkSameResult(expected, result)
  }

  test("csv from Dataset[String] with invalid schema") {
    val message = intercept[ParseException] {
      spark.read.schema("123").csv(spark.createDataset(Seq.empty[String])(StringEncoder))
    }.getMessage
    assert(message.contains("PARSE_SYNTAX_ERROR"))
  }

  test("Dataset result destructive iterator") {
    // Helper methods for accessing private field `idxToBatches` from SparkResult
    val getResultMap =
      PrivateMethod[mutable.Map[Int, Any]](Symbol("resultMap"))

    def assertResultsMapEmpty(result: SparkResult[_]): Unit = {
      val resultMap = result invokePrivate getResultMap()
      assert(resultMap.isEmpty)
    }

    val df = spark
      .range(0, 10, 1, 10)
      .filter("id > 5 and id < 9")

    df.withResult { result =>
      // build and verify the destructive iterator
      val iterator = result.destructiveIterator
      try {
        // resultMap Map is empty before traversing the result iterator
        assertResultsMapEmpty(result)
        val buffer = mutable.Set.empty[Long]
        while (iterator.hasNext) {
          // resultMap is empty during iteration because results get removed immediately on access.
          assertResultsMapEmpty(result)
          buffer += iterator.next()
        }
        // resultMap Map is empty afterward because all results have been removed.
        assertResultsMapEmpty(result)

        val expectedResult = Set(6L, 7L, 8L)
        assert(buffer.size === 3 && expectedResult == buffer)
      } finally {
        iterator.close()
      }
    }
  }

  test("SparkSession.createDataFrame - large data set") {
    val threshold = 1024 * 1024
    withSQLConf(SqlApiConf.LOCAL_RELATION_CACHE_THRESHOLD_KEY -> threshold.toString) {
      val count = 2
      val suffix = "abcdef"
      val str = scala.util.Random.alphanumeric.take(1024 * 1024).mkString + suffix
      val data = Seq.tabulate(count)(i => (i, str))
      for (_ <- 0 until 2) {
        val df = spark.createDataFrame(data)
        assert(df.count() === count)
        assert(!df.filter(df("_2").endsWith(suffix)).isEmpty)
      }
    }
  }

  test("sql() with positional parameters") {
    val result0 = spark.sql("select 1", Array.empty).collect()
    assert(result0.length == 1 && result0(0).getInt(0) === 1)

    val result1 = spark.sql("select ?", Array(1)).collect()
    assert(result1.length == 1 && result1(0).getInt(0) === 1)

    val result2 = spark.sql("select ?, ?", Array(1, "abc")).collect()
    assert(result2.length == 1)
    assert(result2(0).getInt(0) === 1)
    assert(result2(0).getString(1) === "abc")

    val result3 = spark.sql("select element_at(?, 1)", Array(array(lit(1)))).collect()
    assert(result3.length == 1 && result3(0).getInt(0) === 1)
  }

  test("sql() with named parameters") {
    val result0 = spark.sql("select 1", Map.empty[String, Any]).collect()
    assert(result0.length == 1 && result0(0).getInt(0) === 1)

    val result1 = spark.sql("select :abc", Map("abc" -> 1)).collect()
    assert(result1.length == 1 && result1(0).getInt(0) === 1)

    val result2 = spark.sql("select :c0 limit :l0", Map("l0" -> 1, "c0" -> "abc")).collect()
    assert(result2.length == 1 && result2(0).getString(0) === "abc")

    val result3 =
      spark.sql("select element_at(:m, 'a')", Map("m" -> map(lit("a"), lit(1)))).collect()
    assert(result3.length == 1 && result3(0).getInt(0) === 1)
  }

  test("joinWith, flat schema") {
    val session: SparkSession = spark
    import session.implicits._
    val ds1 = Seq(1, 2, 3).toDS().as("a")
    val ds2 = Seq(1, 2).toDS().as("b")

    val joined = ds1.joinWith(ds2, $"a.value" === $"b.value", "inner")

    val expectedSchema = StructType(
      Seq(
        StructField("_1", IntegerType, nullable = false),
        StructField("_2", IntegerType, nullable = false)))

    assert(joined.schema === expectedSchema)

    val expected = Seq((1, 1), (2, 2))
    checkSameResult(expected, joined)
  }

  test("joinWith tuple with primitive, expression") {
    val session: SparkSession = spark
    import session.implicits._
    val ds1 = Seq(1, 1, 2).toDS()
    val ds2 = Seq(("a", 1), ("b", 2)).toDS()

    val joined = ds1.joinWith(ds2, $"value" === $"_2")

    // This is an inner join, so both outputs fields are non-nullable
    val expectedSchema = StructType(
      Seq(
        StructField("_1", IntegerType, nullable = false),
        StructField(
          "_2",
          StructType(
            Seq(StructField("_1", StringType), StructField("_2", IntegerType, nullable = false))),
          nullable = false)))
    assert(joined.schema === expectedSchema)

    checkSameResult(Seq((1, ("a", 1)), (1, ("a", 1)), (2, ("b", 2))), joined)
  }

  test("joinWith tuple with primitive, rows") {
    val session: SparkSession = spark
    import session.implicits._
    val ds1 = Seq(1, 1, 2).toDF()
    val ds2 = Seq(("a", 1), ("b", 2)).toDF()

    val joined = ds1.joinWith(ds2, $"value" === $"_2")

    checkSameResult(
      Seq((Row(1), Row("a", 1)), (Row(1), Row("a", 1)), (Row(2), Row("b", 2))),
      joined)
  }

  test("joinWith class with primitive, toDF") {
    val session: SparkSession = spark
    import session.implicits._
    val ds1 = Seq(1, 1, 2).toDS()
    val ds2 = Seq(ClassData("a", 1), ClassData("b", 2)).toDS()

    val df = ds1
      .joinWith(ds2, $"value" === $"b")
      .toDF()
      .select($"_1", $"_2.a", $"_2.b")
    checkSameResult(Seq(Row(1, "a", 1), Row(1, "a", 1), Row(2, "b", 2)), df)
  }

  test("multi-level joinWith") {
    val session: SparkSession = spark
    import session.implicits._
    val ds1 = Seq(("a", 1), ("b", 2)).toDS().as("a")
    val ds2 = Seq(("a", 1), ("b", 2)).toDS().as("b")
    val ds3 = Seq(("a", 1), ("b", 2)).toDS().as("c")

    val joined = ds1
      .joinWith(ds2, $"a._2" === $"b._2")
      .as("ab")
      .joinWith(ds3, $"ab._1._2" === $"c._2")

    checkSameResult(
      Seq(((("a", 1), ("a", 1)), ("a", 1)), ((("b", 2), ("b", 2)), ("b", 2))),
      joined)
  }

  test("multi-level joinWith, rows") {
    val session: SparkSession = spark
    import session.implicits._
    val ds1 = Seq(("a", 1), ("b", 2)).toDF().as("a")
    val ds2 = Seq(("a", 1), ("b", 2)).toDF().as("b")
    val ds3 = Seq(("a", 1), ("b", 2)).toDF().as("c")

    val joined = ds1
      .joinWith(ds2, $"a._2" === $"b._2")
      .as("ab")
      .joinWith(ds3, $"ab._1._2" === $"c._2")

    checkSameResult(
      Seq(((Row("a", 1), Row("a", 1)), Row("a", 1)), ((Row("b", 2), Row("b", 2)), Row("b", 2))),
      joined)
  }

  test("self join") {
    val session: SparkSession = spark
    import session.implicits._
    val ds = Seq("1", "2").toDS().as("a")
    val joined = ds.joinWith(ds, lit(true), "cross")
    checkSameResult(Seq(("1", "1"), ("1", "2"), ("2", "1"), ("2", "2")), joined)
  }

  test("SPARK-11894: Incorrect results are returned when using null") {
    val session: SparkSession = spark
    import session.implicits._
    val nullInt = null.asInstanceOf[java.lang.Integer]
    val ds1 = Seq((nullInt, "1"), (java.lang.Integer.valueOf(22), "2")).toDS()
    val ds2 = Seq((nullInt, "1"), (java.lang.Integer.valueOf(22), "2")).toDS()

    checkSameResult(
      Seq(
        ((nullInt, "1"), (nullInt, "1")),
        ((nullInt, "1"), (java.lang.Integer.valueOf(22), "2")),
        ((java.lang.Integer.valueOf(22), "2"), (nullInt, "1")),
        ((java.lang.Integer.valueOf(22), "2"), (java.lang.Integer.valueOf(22), "2"))),
      ds1.joinWith(ds2, lit(true), "cross"))
  }

  test("SPARK-15441: Dataset outer join") {
    val session: SparkSession = spark
    import session.implicits._
    val left = Seq(ClassData("a", 1), ClassData("b", 2)).toDS().as("left")
    val right = Seq(ClassData("x", 2), ClassData("y", 3)).toDS().as("right")
    val joined = left.joinWith(right, $"left.b" === $"right.b", "left")

    val expectedSchema = StructType(
      Seq(
        StructField(
          "_1",
          StructType(
            Seq(StructField("a", StringType), StructField("b", IntegerType, nullable = false))),
          nullable = false),
        // This is a left join, so the right output is nullable:
        StructField(
          "_2",
          StructType(
            Seq(StructField("a", StringType), StructField("b", IntegerType, nullable = false))))))
    assert(joined.schema === expectedSchema)

    val result = joined.collect().toSet
    assert(result == Set(ClassData("a", 1) -> null, ClassData("b", 2) -> ClassData("x", 2)))
  }

  test("SPARK-37829: DataFrame outer join") {
    // Same as "SPARK-15441: Dataset outer join" but using DataFrames instead of Datasets
    val session: SparkSession = spark
    import session.implicits._
    val left = Seq(ClassData("a", 1), ClassData("b", 2)).toDF().as("left")
    val right = Seq(ClassData("x", 2), ClassData("y", 3)).toDF().as("right")
    val joined = left.joinWith(right, $"left.b" === $"right.b", "left")

    val leftFieldSchema = StructType(
      Seq(StructField("a", StringType), StructField("b", IntegerType, nullable = false)))
    val rightFieldSchema = StructType(
      Seq(StructField("a", StringType), StructField("b", IntegerType, nullable = false)))
    val expectedSchema = StructType(
      Seq(
        StructField("_1", leftFieldSchema, nullable = false),
        // This is a left join, so the right output is nullable:
        StructField("_2", rightFieldSchema)))
    assert(joined.schema === expectedSchema)

    val result = joined.collect().toSet
    val expected = Set(
      new GenericRowWithSchema(Array("a", 1), leftFieldSchema) ->
        null,
      new GenericRowWithSchema(Array("b", 2), leftFieldSchema) ->
        new GenericRowWithSchema(Array("x", 2), rightFieldSchema))
    assert(result == expected)
  }

  test("SPARK-24762: joinWith on Option[Product]") {
    val session: SparkSession = spark
    import session.implicits._
    val ds1 = Seq(Some((1, 2)), Some((2, 3)), None).toDS().as("a")
    val ds2 = Seq(Some((1, 2)), Some((2, 3)), None).toDS().as("b")
    val joined = ds1.joinWith(ds2, $"a.value._1" === $"b.value._2", "inner")
    checkSameResult(Seq((Some((2, 3)), Some((1, 2)))), joined)
  }

  test("dropDuplicatesWithinWatermark not supported in batch DataFrame") {
    def testAndVerify(df: Dataset[_]): Unit = {
      val exc = intercept[AnalysisException] {
        df.write.format("noop").mode(SaveMode.Append).save()
      }

      assert(exc.getMessage.contains("dropDuplicatesWithinWatermark is not supported"))
      assert(exc.getMessage.contains("batch DataFrames/DataSets"))
    }

    val result = spark.range(10).dropDuplicatesWithinWatermark()
    testAndVerify(result)

    val result2 = spark
      .range(10)
      .withColumn("newcol", col("id"))
      .dropDuplicatesWithinWatermark("newcol")
    testAndVerify(result2)
  }

  test("Dataset.metadataColumn") {
    val session: SparkSession = spark
    import session.implicits._
    withTempPath { file =>
      val path = file.getAbsoluteFile.toURI.toString
      spark
        .range(0, 100, 1, 1)
        .withColumn("_metadata", concat(lit("lol_"), col("id")))
        .write
        .parquet(file.toPath.toAbsolutePath.toString)

      val df = spark.read.parquet(path)
      val (filepath, rc) = df
        .groupBy(df.metadataColumn("_metadata").getField("file_path"))
        .count()
        .as[(String, Long)]
        .head()
      assert(filepath.startsWith(path))
      assert(rc == 100)
    }
  }

  test("SPARK-45216: Non-deterministic functions with seed") {
    val session: SparkSession = spark
    import session.implicits._

    val df = Seq(Array.range(0, 10)).toDF("a")

    val r = rand()
    val r2 = randn()
    val r3 = random()
    val r4 = uuid()
    val r5 = shuffle(col("a"))
    df.select(r, r.as("r"), r2, r2.as("r2"), r3, r3.as("r3"), r4, r4.as("r4"), r5, r5.as("r5"))
      .collect()
      .foreach { row =>
        (0 until 5).foreach(i => assert(row.get(i * 2) === row.get(i * 2 + 1)))
      }
  }

  test("Observable metrics") {
    val df = spark.range(99).withColumn("extra", col("id") - 1)
    val ob1 = new Observation("ob1")
    val observedDf = df.observe(ob1, min("id"), avg("id"), max("id"))
    val observedObservedDf = observedDf.observe("ob2", min("extra"), avg("extra"), max("extra"))

    val ob1Schema = new StructType()
      .add("min(id)", LongType)
      .add("avg(id)", DoubleType)
      .add("max(id)", LongType)
    val ob2Schema = new StructType()
      .add("min(extra)", LongType)
      .add("avg(extra)", DoubleType)
      .add("max(extra)", LongType)
    val ob1Metrics = Map("ob1" -> new GenericRowWithSchema(Array(0, 49, 98), ob1Schema))
    val ob2Metrics = Map("ob2" -> new GenericRowWithSchema(Array(-1, 48, 97), ob2Schema))

    val obMetrics = observedDf.withResult(_.getObservedMetrics)
    assert(df.withResult(_.getObservedMetrics) === Map.empty)
    assert(observedDf.withResult(_.getObservedMetrics) === ob1Metrics)
    assert(obMetrics.map(_._2.schema) === Seq(ob1Schema))

    val obObMetrics = observedObservedDf.withResult(_.getObservedMetrics)
    assert(obObMetrics === ob1Metrics ++ ob2Metrics)
    assert(obObMetrics.map(_._2.schema).exists(_.equals(ob1Schema)))
    assert(obObMetrics.map(_._2.schema).exists(_.equals(ob2Schema)))
  }

  for (collectFunc <- Seq(
      ("collect", (df: DataFrame) => df.collect()),
      ("collectAsList", (df: DataFrame) => df.collectAsList()),
      ("collectResult", (df: DataFrame) => df.withResult(_.length)),
      ("write", (df: DataFrame) => df.write.format("noop").mode("append").save())))
    test(
      "Observation.get is blocked until the query is finished, " +
        s"collect using method ${collectFunc._1}") {
      val df = spark.range(99).withColumn("extra", col("id") - 1)
      val ob1 = new Observation("ob1")
      val ob2 = new Observation("ob2")
      val observedDf = df.observe(ob1, min("id"), avg("id"), max("id"))
      val observedObservedDf = observedDf.observe(ob2, min("extra"), avg("extra"), max("extra"))
      // Start new threads to get observations
      val future1 = Future(ob1.get)(ExecutionContext.global)
      val future2 = Future(ob2.get)(ExecutionContext.global)
      // make sure the threads are blocked right now
      val e1 = intercept[java.util.concurrent.TimeoutException] {
        SparkThreadUtils.awaitResult(future1, 2.seconds)
      }
      assert(e1.getMessage.contains("timed out after"))
      val e2 = intercept[java.util.concurrent.TimeoutException] {
        SparkThreadUtils.awaitResult(future2, 2.seconds)
      }
      assert(e2.getMessage.contains("timed out after"))
      collectFunc._2(observedObservedDf)
      // make sure the threads are unblocked after the query is finished
      val metrics1 = SparkThreadUtils.awaitResult(future1, 5.seconds)
      assert(metrics1 === Map("min(id)" -> 0, "avg(id)" -> 49, "max(id)" -> 98))
      val metrics2 = SparkThreadUtils.awaitResult(future2, 5.seconds)
      assert(metrics2 === Map("min(extra)" -> -1, "avg(extra)" -> 48, "max(extra)" -> 97))
    }

  test("SPARK-48852: trim function on a string column returns correct results") {
    val session: SparkSession = spark
    import session.implicits._
    val df = Seq("  a  ", "b  ", "   c").toDF("col")
    val result = df.select(trim(col("col"), " ").as("trimmed_col")).collect()
    assert(result sameElements Array(Row("a"), Row("b"), Row("c")))
  }

  test("SPARK-49673: new batch size, multiple batches") {
    val maxBatchSize = spark.conf.get("spark.connect.grpc.arrow.maxBatchSize").dropRight(1).toInt
    // Adjust client grpcMaxMessageSize to maxBatchSize (10MiB; set in RemoteSparkSession config)
    val sparkWithLowerMaxMessageSize = SparkSession
      .builder()
      .client(
        SparkConnectClient
          .builder()
          .userId("test")
          .port(port)
          .grpcMaxMessageSize(maxBatchSize)
          .retryPolicy(RetryPolicy
            .defaultPolicy()
            .copy(maxRetries = Some(10), maxBackoff = Some(FiniteDuration(30, "s"))))
          .build())
      .create()
    assert(sparkWithLowerMaxMessageSize.range(maxBatchSize).collect().length == maxBatchSize)
  }

  test("SPARK-50965: Multiple positional parameterized nodes in the parsed logical plan") {
    var df = spark.sql("SELECT ?", Array(0))
    for (i <- 1 until 3) {
      val temp = spark.sql("SELECT ?", Array(i))
      df = df.union(temp)
    }
    checkAnswer(df, (0 until 3).map(i => Row(i)))
  }

  test("SPARK-50965: Multiple named parameterized nodes in the parsed logical plan") {
    var df = spark.sql("SELECT :key", args = Map("key" -> 0))
    for (i <- 1 until 3) {
      val temp = spark.sql("SELECT :key", args = Map("key" -> i))
      df = df.union(temp)
    }
    checkAnswer(df, (0 until 3).map(i => Row(i)))
  }

  test(
    "SPARK-50965: Multiple named and positional parameterized nodes in the parsed logical plan") {
    var df = spark.sql("SELECT ?", Array(0))
    df = df.union(spark.sql("SELECT :key", args = Map("key" -> 1)))
    for (i <- 1 until 3) {
      df = df.union(spark.sql("SELECT ?", Array(2 * i)))
      df = df.union(spark.sql("SELECT :key", args = Map("key" -> (2 * i + 1))))
    }
    checkAnswer(df, (0 until 6).map(i => Row(i)))
  }

  test("SPARK-52770: Support Time type") {
    val df = spark.sql("SELECT TIME '12:13:14'")

    checkAnswer(df, Row(LocalTime.of(12, 13, 14)))
  }

  test("SPARK-53054: DataFrameReader defaults to spark.sql.sources.default") {
    withTempPath { file =>
      val path = file.getAbsoluteFile.toURI.toString
      spark.range(100).write.parquet(file.toPath.toAbsolutePath.toString)

      spark.conf.set("spark.sql.sources.default", "parquet")

      val df = spark.read.load(path)
      assert(df.count() == 100)
    }
  }

  test("SPARK-52930: the nullability of arrays should be preserved using typedlit") {
    val arrays = Seq(
      (typedlit(Array[Int]()), false),
      (typedlit(Array[Int](1)), false),
      (typedlit(Array[Integer]()), true),
      (typedlit(Array[Integer](1)), true))
    for ((array, containsNull) <- arrays) {
      val df = spark.sql("select 1").select(array)
      df.createOrReplaceTempView("test_array_nullability")
      val schema = spark.sql("select * from test_array_nullability").schema
      assert(schema.fields.head.dataType.asInstanceOf[ArrayType].containsNull === containsNull)
    }
  }

  test("SPARK-52930: the nullability of map values should be preserved using typedlit") {
    val maps = Seq(
      (typedlit(Map[String, Int]()), false),
      (typedlit(Map[String, Int]("a" -> 1)), false),
      (typedlit(Map[String, Integer]()), true),
      (typedlit(Map[String, Integer]("a" -> 1)), true))
    for ((map, valueContainsNull) <- maps) {
      val df = spark.sql("select 1").select(map)
      df.createOrReplaceTempView("test_map_nullability")
      val schema = spark.sql("select * from test_map_nullability").schema
      assert(
        schema.fields.head.dataType.asInstanceOf[MapType].valueContainsNull === valueContainsNull)
    }
  }

  test("SPARK-54043: DirectShufflePartitionID should be supported") {
    val df = spark.range(100).withColumn("expected_p_id", col("id") % 10)
    val repartitioned = df.repartitionById(10, col("expected_p_id").cast("int"))
    val result = repartitioned.withColumn("actual_p_id", spark_partition_id())

    assert(result.filter(col("expected_p_id") =!= col("actual_p_id")).count() == 0)

    val negativeDf = spark.range(10).toDF("id")
    val negativeRepartitioned = negativeDf.repartitionById(10, (col("id") - 5).cast("int"))
    val negativeResult =
      negativeRepartitioned
        .withColumn("actual_p_id", spark_partition_id())
        .collect()

    assert(negativeResult.forall(row => {
      val actualPartitionId = row.getAs[Int]("actual_p_id")
      val id = row.getAs[Long]("id")
      val expectedPartitionId = {
        val mod = (id - 5) % 10
        if (mod < 0) mod + 10 else mod
      }.toInt
      actualPartitionId == expectedPartitionId
    }))

    val nullDf = spark.range(10).toDF("id")
    val nullExpr = when(col("id") < 5, col("id")).otherwise(lit(null)).cast("int")
    val nullRepartitioned = nullDf.repartitionById(10, nullExpr)
    val nullResult = nullRepartitioned.withColumn("actual_p_id", spark_partition_id()).collect()

    val nullRows = nullResult.filter(_.getAs[Long]("id") >= 5)
    assert(nullRows.forall(_.getAs[Int]("actual_p_id") == 0))
  }

  test("SPARK-53490: struct type in observed metrics") {
    val observation = Observation("struct")
    spark
      .range(10)
      .observe(observation, struct(count(lit(1)).as("rows"), max("id").as("maxid")).as("struct"))
      .collect()
    val expectedSchema =
      StructType(Seq(StructField("rows", LongType), StructField("maxid", LongType)))
    val expectedValue = new GenericRowWithSchema(Array(10, 9), expectedSchema)
    assert(observation.get.size === 1)
    assert(observation.get.contains("struct"))
    assert(observation.get("struct") === expectedValue)
  }

  test("SPARK-53490: array type in observed metrics") {
    val observation = Observation("array")
    spark
      .range(10)
      .observe(observation, array(count(lit(1))).as("array"))
      .collect()
    assert(observation.get.size === 1)
    assert(observation.get.contains("array"))
    assert(observation.get("array") === Array(10))
  }

  test("SPARK-53490: map type in observed metrics") {
    val observation = Observation("map")
    spark
      .range(10)
      .observe(observation, map(lit("count"), count(lit(1))).as("map"))
      .collect()
    assert(observation.get.size === 1)
    assert(observation.get.contains("map"))
    assert(observation.get("map") === Map("count" -> 10))
  }

  test("SPARK-53553: null value handling in literals") {
    val df = spark.sql("select 1").select(typedlit(Array[Integer](1, null)).as("arr_col"))
    val result = df.collect()
    assert(result.length === 1)
    assert(result(0).getAs[Array[Integer]]("arr_col") === Array(1, null))
  }

  // SQL Scripting tests
  test("SQL Script result") {
    withSQLConf("spark.sql.scripting.enabled" -> "true") {
      val df = spark.sql("""BEGIN
          |  IF 1=1 THEN
          |    SELECT 1;
          |  ELSE
          |    SELECT 3;
          |  END IF;
          |END
          |""".stripMargin)
      checkAnswer(df, Seq(Row(1)))
    }
  }

  test("SQL Script schema") {
    withSQLConf("spark.sql.scripting.enabled" -> "true") {
      withTable("script_tbl") {
        val df = spark.sql("""BEGIN
            |  CREATE TABLE script_tbl (a INT, b STRING);
            |  INSERT INTO script_tbl VALUES (1, 'Hello'), (2, 'World');
            |  SELECT * FROM script_tbl;
            |END
            |""".stripMargin)
        assert(
          df.schema == StructType(
            StructField("a", IntegerType, nullable = true)
              :: StructField("b", StringType, nullable = true)
              :: Nil))
      }
    }
  }

  test("SQL Script empty result") {
    withSQLConf("spark.sql.scripting.enabled" -> "true") {
      withTable("script_tbl") {
        val df = spark.sql("""BEGIN
            |  CREATE TABLE script_tbl (a INT, b STRING);
            |  SELECT * FROM script_tbl;
            |END
            |""".stripMargin)
        assert(
          df.schema == StructType(
            StructField("a", IntegerType, nullable = true)
              :: StructField("b", StringType, nullable = true)
              :: Nil))
        checkAnswer(df, Seq.empty)
      }
    }
  }

  test("SQL Script no result") {
    withSQLConf("spark.sql.scripting.enabled" -> "true") {
      withTable("script_tbl") {
        val df = spark.sql("""BEGIN
            |  CREATE TABLE script_tbl (a INT, b STRING);
            |END
            |""".stripMargin)
        assert(df.schema == StructType(Nil))
        checkAnswer(df, Seq.empty)
      }
    }
  }

  // Helper class to capture Arrow batch chunk information from gRPC responses
  private class ArrowBatchInterceptor extends ClientInterceptor {
    case class BatchInfo(
        batchIndex: Int,
        rowCount: Long,
        startOffset: Long,
        chunks: Seq[ChunkInfo]) {
      def totalChunks: Int = chunks.length
    }

    case class ChunkInfo(
        batchIndex: Int,
        chunkIndex: Int,
        numChunksInBatch: Int,
        rowCount: Long,
        startOffset: Long,
        dataSize: Int)

    private val batches: mutable.Buffer[BatchInfo] = mutable.Buffer.empty
    private var currentBatchIndex: Int = 0
    private val currentBatchChunks: mutable.Buffer[ChunkInfo] = mutable.Buffer.empty

    override def interceptCall[ReqT, RespT](
        method: MethodDescriptor[ReqT, RespT],
        callOptions: CallOptions,
        next: Channel): ClientCall[ReqT, RespT] = {
      new ForwardingClientCall.SimpleForwardingClientCall[ReqT, RespT](
        next.newCall(method, callOptions)) {
        override def start(
            responseListener: ClientCall.Listener[RespT],
            headers: io.grpc.Metadata): Unit = {
          super.start(
            new ForwardingClientCallListener.SimpleForwardingClientCallListener[RespT](
              responseListener) {
              override def onMessage(message: RespT): Unit = {
                message match {
                  case response: proto.ExecutePlanResponse if response.hasArrowBatch =>
                    val arrowBatch = response.getArrowBatch
                    // Track chunk information for every chunk
                    currentBatchChunks += ChunkInfo(
                      batchIndex = currentBatchIndex,
                      chunkIndex = arrowBatch.getChunkIndex.toInt,
                      numChunksInBatch = arrowBatch.getNumChunksInBatch.toInt,
                      rowCount = arrowBatch.getRowCount,
                      startOffset = arrowBatch.getStartOffset,
                      dataSize = arrowBatch.getData.size())
                    // When we receive the last chunk, create the BatchInfo
                    if (currentBatchChunks.length == arrowBatch.getNumChunksInBatch) {
                      batches += BatchInfo(
                        batchIndex = currentBatchIndex,
                        rowCount = arrowBatch.getRowCount,
                        startOffset = arrowBatch.getStartOffset,
                        chunks = currentBatchChunks.toList)
                      currentBatchChunks.clear()
                      currentBatchIndex += 1
                    }
                  case _ => // Not an ExecutePlanResponse with ArrowBatch, ignore
                }
                super.onMessage(message)
              }
            },
            headers)
        }
      }
    }

    // Get all batch information
    def getBatchInfos: Seq[BatchInfo] = batches.toSeq

    def clear(): Unit = {
      currentBatchIndex = 0
      currentBatchChunks.clear()
      batches.clear()
    }
  }

  test("Arrow batch result chunking") {
    // This test validates that the client can correctly reassemble chunked Arrow batches
    // using SequenceInputStream as implemented in SparkResult.processResponses

    // Two cases are tested here:
    // (a) client preferred chunk size is set: the server should respect it
    // (b) client preferred chunk size is not set: the server should use its own max chunk size
    Seq((Some(1024), None), (None, Some(1024))).foreach {
      case (preferredChunkSizeOpt, maxChunkSizeOpt) =>
        // Create interceptor to capture chunk information
        val arrowBatchInterceptor = new ArrowBatchInterceptor()

        try {
          // Set preferred chunk size if specified and add interceptor
          preferredChunkSizeOpt match {
            case Some(size) =>
              spark = createSparkSession(
                _.preferredArrowChunkSize(Some(size)).interceptor(arrowBatchInterceptor))
            case None =>
              spark = createSparkSession(_.interceptor(arrowBatchInterceptor))
          }
          // Set server max chunk size if specified
          maxChunkSizeOpt.foreach { size =>
            spark.conf.set("spark.connect.session.resultChunking.maxChunkSize", size.toString)
          }

          val sqlQuery =
            "select id, CAST(id + 0.5 AS DOUBLE) as double_val from range(0, 2000, 1, 4)"

          // Execute the query using withResult to access SparkResult object
          spark.sql(sqlQuery).withResult { result =>
            // Verify the results are correct and complete
            assert(result.length == 2000)

            // Get batch information from interceptor
            val batchInfos = arrowBatchInterceptor.getBatchInfos

            // Assert there are 4 batches (partitions) in total
            assert(batchInfos.length == 4)

            // Validate chunk information for each batch
            val maxChunkSize = preferredChunkSizeOpt.orElse(maxChunkSizeOpt).get
            batchInfos.foreach { batch =>
              // In this example, the max chunk size is set to a small value,
              // so each Arrow batch should be split into multiple chunks
              assert(batch.totalChunks > 5)
              assert(batch.chunks.nonEmpty)
              assert(batch.chunks.length == batch.totalChunks)
              batch.chunks.zipWithIndex.foreach { case (chunk, expectedIndex) =>
                assert(chunk.chunkIndex == expectedIndex)
                assert(chunk.numChunksInBatch == batch.totalChunks)
                assert(chunk.rowCount == batch.rowCount)
                assert(chunk.startOffset == batch.startOffset)
                assert(chunk.dataSize > 0)
                assert(chunk.dataSize <= maxChunkSize)
              }
            }

            // Validate data integrity across the range to ensure chunking didn't corrupt anything
            val rows = result.toArray
            var expectedId = 0L
            rows.foreach { row =>
              assert(row.getLong(0) == expectedId)
              val expectedDouble = expectedId + 0.5
              val actualDouble = row.getDouble(1)
              assert(math.abs(actualDouble - expectedDouble) < 0.001)
              expectedId += 1
            }
          }
        } finally {
          // Clean up configurations
          maxChunkSizeOpt.foreach { _ =>
            spark.conf.unset("spark.connect.session.resultChunking.maxChunkSize")
          }
          arrowBatchInterceptor.clear()
        }
    }
  }

  test("Plan compression works correctly") {
    val originalPlanCompressionOptions = spark.client.getPlanCompressionOptions
    assert(originalPlanCompressionOptions.nonEmpty)
    assert(originalPlanCompressionOptions.get.thresholdBytes > 0)
    assert(originalPlanCompressionOptions.get.algorithm == "ZSTD")
    try {
      spark.client.setPlanCompressionOptions(Some(PlanCompressionOptions(1000, "ZSTD")))
      // Execution should work
      assert(spark.sql(s"select '${"Apache Spark" * 10000}' as value").collect().length == 1)
      // Analysis should work
      assert(spark.sql(s"select '${"Apache Spark" * 10000}' as value").columns.length == 1)
    } finally {
      spark.client.setPlanCompressionOptions(originalPlanCompressionOptions)
    }
  }
}

private[sql] case class ClassData(a: String, b: Int)

private[sql] case class MyType(id: Long, a: Double, b: Double)
private[sql] case class KV(key: String, value: Int)
private[sql] class SimpleBean {
  @scala.beans.BeanProperty
  var value: String = _
}
