Add UNSHREDDED Parquet writer and read support for Variant

[the-other-tim-brown]

Task Description

What needs to be done:
Add support for writing parquet files with Variant data types.

Additionally, set up Spark functional tests (or augment existing ones) for the variant type.

Also ensure that files with this data type can be read back into Hoodie Records or engine specific row types in MOR and COW tables

Update functional tests in spark to ensure that the data can be read back properly.

Why this task is needed:
This sets up the foundation for supporting Variant in Hudi

After this is completed, we should be able to have a simple CoW and Mor workflow with variant for unshredded data.

Task Type

Code improvement/refactoring

Related Issues

Parent feature issue: (if applicable )
Related issues:
NOTE: Use Relationships button to add parent/blocking issues after issue is created.

[voonhous]

Test code to debug variant read path, I added this for debugging to allow for finer grain control. It is almost identical to TestVariantDataType.scala. Just that this instantiates a CloseableInternalRowIterator for row reading.

Place it in: hudi/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/hudi/common/table/read/

Details

/*
 * 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.hudi.common.table.read

import org.apache.hudi.{HoodieSparkUtils, SparkAdapterSupport}
import org.apache.hudi.common.table.HoodieTableMetaClient
import org.apache.hudi.common.testutils.{HoodieTestTable, HoodieTestUtils}
import org.apache.hudi.common.util.{Option => HOption}
import org.apache.hudi.storage.StorageConfiguration
import org.apache.hudi.storage.hadoop.HadoopStorageConfiguration
import org.apache.hudi.util.CloseableInternalRowIterator

import org.apache.hadoop.conf.Configuration
import org.apache.spark.{HoodieSparkKryoRegistrar, SparkConf}
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.internal.SQLConf.LEGACY_RESPECT_NULLABILITY_IN_TEXT_DATASET_CONVERSION
import org.apache.spark.sql.types.StructType
import org.junit.jupiter.api.{AfterEach, BeforeEach, Test}
import org.junit.jupiter.api.Assertions.{assertArrayEquals, assertEquals}
import org.junit.jupiter.api.Assumptions.assumeTrue

import java.lang.Exception
import java.nio.file.{Files, Path}

class TestHoodieFileGroupReaderOnSparkVariant extends SparkAdapterSupport {
  var spark: SparkSession = _
  var tempDir: Path = _

  @BeforeEach
  def setup(): Unit = {
    val sparkConf = new SparkConf
    sparkConf.set("spark.app.name", getClass.getName)
    sparkConf.set("spark.master", "local[8]")
    sparkConf.set("spark.default.parallelism", "4")
    sparkConf.set("spark.sql.shuffle.partitions", "4")
    sparkConf.set("spark.driver.maxResultSize", "2g")
    sparkConf.set("spark.serializer", "org.apache.spark.serializer.KryoSerializer")
    sparkConf.set("spark.kryo.registrator", "org.apache.spark.HoodieSparkKryoRegistrar")
    sparkConf.set("spark.sql.extensions", "org.apache.spark.sql.hudi.HoodieSparkSessionExtension")
    sparkConf.set("spark.sql.parquet.enableVectorizedReader", "false")
    sparkConf.set("spark.sql.orc.enableVectorizedReader", "false")
    sparkConf.set(LEGACY_RESPECT_NULLABILITY_IN_TEXT_DATASET_CONVERSION.key, "true")
    HoodieSparkKryoRegistrar.register(sparkConf)
    spark = SparkSession.builder.config(sparkConf).getOrCreate
    tempDir = Files.createTempDirectory("test_variant_")
  }

  @AfterEach
  def teardown(): Unit = {
    if (spark != null) {
      spark.stop()
    }
  }

  def getStorageConf: StorageConfiguration[_] = {
    HoodieTestUtils.getDefaultStorageConf.getInline
  }

  def getBasePath: String = {
    tempDir.toAbsolutePath.toUri.toString
  }

  @Test
  def testReadVariantDataType(): Unit = {
    // Variant type is only supported in Spark 4.0+
    assumeTrue(HoodieSparkUtils.gteqSpark4_0, "Variant type requires Spark 4.0 or higher")

    val tableName = "test_variant_table"

    // Create table with variant column
    spark.sql(
      s"""
         |create table $tableName (
         |  id int,
         |  name string,
         |  v variant,
         |  b binary,
         |  ts long
         |) using hudi
         | location '$getBasePath'
         | tblproperties (
         |  primaryKey = 'id',
         |  type = 'mor',
         |  preCombineField = 'ts'
         | )
       """.stripMargin)

    // Insert variant data
    spark.sql(
      s"""
         |insert into $tableName
         |values
         |  (1, 'row1', parse_json('{"key": "value1", "num": 1}'), X'0102030405', 1000),
         |  (2, 'row2', parse_json('{"key": "value2", "list": [1, 2, 3]}'), X'0504030201', 1000)
       """.stripMargin)

    // Update variant data
    spark.sql(
      s"""
         |update $tableName
         |set v = parse_json('{"updated": true, "new_field": 123}')
         |where id = 1
       """.stripMargin)

    // Get metaClient and base files
    val metaClient = HoodieTableMetaClient.builder()
      .setConf(getStorageConf)
      .setBasePath(getBasePath)
      .build()
    val allBaseFiles = HoodieTestTable.of(metaClient).listAllBaseFiles
    assertEquals(1, allBaseFiles.size())

    // Create parquet reader
    val hadoopConf = new Configuration(spark.sparkContext.hadoopConfiguration)
    val reader = sparkAdapter.createParquetFileReader(
      vectorized = false,
      spark.sessionState.conf,
      Map.empty,
      hadoopConf)

    // Get the schema for the table
    val dataSchema = spark.table(tableName).schema

    // Create partitioned file for reading
    val fileInfo = sparkAdapter.getSparkPartitionedFileUtils.createPartitionedFile(
      InternalRow.empty,
      allBaseFiles.get(0).getPath,
      0,
      allBaseFiles.get(0).getLength)

    // Read using CloseableInternalRowIterator
    val iterator = new CloseableInternalRowIterator(
      reader.read(
        fileInfo,
        dataSchema,
        StructType(Seq.empty),
        HOption.empty(),
        Seq.empty,
        new HadoopStorageConfiguration(hadoopConf)))

    // Validate we can read variant data for id = 1
    while (iterator.hasNext) {
      val row = iterator.next()
      val id = row.getInt(5)

      if (id == 1) {
        val name = row.getUTF8String(6).toString
        val variantStr = row.get(7, dataSchema(7).dataType).toString
        val binaryVal = row.getBinary(8)
        val ts = row.getLong(9)

        // Verify row 1 contents
        assertEquals(1, id)
        assertEquals("row1", name)
        // Verify Binary Bytes [1, 2, 3, 4, 5]
        val expectedBytes = Array[Byte](1, 2, 3, 4, 5)
        assertArrayEquals(expectedBytes, binaryVal)
        assertEquals(1000L, ts)
        // Variant data should not be null
        assertEquals("{\"key\":\"value1\",\"num\":1}", variantStr)
      }
    }
    iterator.close()

    // Verify the actual variant contents by casting to string
    val resultDf = spark.sql(s"select id, name, cast(v as string) as v_str, ts from $tableName where id = 1")
    val rows = resultDf.collect()
    assertEquals(1, rows.length, "Should have exactly one row with id=1")

    val resultRow = rows(0)
    assertEquals(1, resultRow.getInt(0))
    assertEquals("row1", resultRow.getString(1))
    val variantStr = resultRow.getString(2)
    assertEquals("{\"new_field\":123,\"updated\":true}", variantStr)
    assertEquals(1000L, resultRow.getLong(3))
  }
}