[BUG] from_json ArrayIndexOutOfBoundsException in 24.02

[andygrove]

Describe the bug

Calling from_json for json with nested structs can cause an ArrayIndexOutOfBoundsException if the provided schema for the nested struct has fewer fields than are present in the json.

Steps/Code to reproduce bug

Input file test.json

{ "x": { "a": 54321, "b": 1, "c": 2} }
{ "x": { "a": 54321, "b": 1, "c": 2} }
{ "x": { "a": 54321, "b": 1, "c": 2} }
{ "x": { "a": 54321, "b": 1, "c": 2} }

Test Setup

scala> import org.apache.spark.sql.types._

scala> val df = spark.read.text("test.json")

scala> df.write.parquet("test.parquet")

scala> val df = spark.read.parquet("test.parquet")

scala> val df2 = df.withColumn("mystruct", from_json(col("value"), new StructType(Array(StructField("x", new StructType(Array(StructField("a", DataTypes.IntegerType, false))))))))

scala> df2.show

CPU

+--------------------+---------+
|               value| mystruct|
+--------------------+---------+
|{ "x": { "a": 543...|{{54321}}|
|{ "x": { "a": 543...|{{54321}}|
|{ "x": { "a": 543...|{{54321}}|
|{ "x": { "a": 543...|{{54321}}|
+--------------------+---------+

GPU

Caused by: java.lang.ArrayIndexOutOfBoundsException: Index 1 out of bounds for length 1
  at org.apache.spark.sql.types.StructType.apply(StructType.scala:423)
  at com.nvidia.spark.rapids.GpuCast$.$anonfun$castStructToStruct$2(GpuCast.scala:1584)
  at com.nvidia.spark.rapids.GpuCast$.$anonfun$castStructToStruct$2$adapted(GpuCast.scala:1580)
  at scala.collection.immutable.Range.foreach(Range.scala:158)
  at com.nvidia.spark.rapids.GpuCast$.$anonfun$castStructToStruct$1(GpuCast.scala:1580)
  at com.nvidia.spark.rapids.Arm$.withResource(Arm.scala:66)
  at com.nvidia.spark.rapids.GpuCast$.castStructToStruct(GpuCast.scala:1579)
  at com.nvidia.spark.rapids.GpuCast$.doCast(GpuCast.scala:584)
  at org.apache.spark.sql.rapids.GpuJsonToStructs.$anonfun$doColumnar$7(GpuJsonToStructs.scala:227)

Cause

The following code uses index as an index into both from and to types and assumes that they are the same size.

  private def castStructToStruct(
      from: StructType,
      to: StructType,
      input: ColumnView,
      options: CastOptions): ColumnVector = {
    withResource(new ArrayBuffer[ColumnVector](from.length)) { childColumns =>
      from.indices.foreach { index =>
        childColumns += doCast(
          input.getChildColumnView(index),
          from(index).dataType,
          to(index).dataType, options)
      }

Expected behavior
Match Spark behavior

Environment details (please complete the following information)

Additional context

[andygrove]

I could not reproduce the issue with latest from branch-24.04

scala> spark.conf.set("spark.rapids.sql.expression.JsonToStructs", true)

scala> df2.show
24/04/02 23:48:01 WARN GpuOverrides: 
!Exec <CollectLimitExec> cannot run on GPU because the Exec CollectLimitExec has been disabled, and is disabled by default because Collect Limit replacement can be slower on the GPU, if huge number of rows in a batch it could help by limiting the number of rows transferred from GPU to CPU. Set spark.rapids.sql.exec.CollectLimitExec to true if you wish to enable it
  @Partitioning <SinglePartition$> could run on GPU
  *Exec <ProjectExec> will run on GPU
    *Expression <Alias> cast(from_json(StructField(x,StructType(StructField(a,IntegerType,false)),true), value#0, Some(UTC)) as string) AS mystruct#18 will run on GPU
      *Expression <Cast> cast(from_json(StructField(x,StructType(StructField(a,IntegerType,false)),true), value#0, Some(UTC)) as string) will run on GPU
        *Expression <JsonToStructs> from_json(StructField(x,StructType(StructField(a,IntegerType,false)),true), value#0, Some(UTC)) will run on GPU
    *Exec <FileSourceScanExec> will run on GPU

+--------------------+---------+
|               value| mystruct|
+--------------------+---------+
|{ "x": { "a": 543...|{{54321}}|
|{ "x": { "a": 543...|{{54321}}|
|{ "x": { "a": 543...|{{54321}}|
|{ "x": { "a": 543...|{{54321}}|
+--------------------+---------+

[thirtiseven]

The test can pass only after this commit in 24.04. It looks like the commit fixed this issue.

There are some schema handling in the commit that might solve the indexing issue in castStructToStruct, but I don't have much context about from_json.

[revans2]

The issue still exists, at least in a few situations. The problem shows up when we ask CUDF to return data with a specific schema, but CUDF sees something that does not match the schema it expects and decides to return a struct or a list instead of an actual string.

rapidsai/cudf#15278

Is the remaining case that I know causes something like this.

[revans2]

I also manually verified that this is working now for 24.04, once you turn on JsonToStructs. I think we can close this as there are other issues to track the remaining JSON work.