Fix BigDiffy NPE on parquet schema mismatch in array fields

.circleci/config.yml

@@ -60,15 +60,21 @@ jobs:
     steps:
       - build
   diffy_212:
-    executor: scala_212
+    docker:
+      - image: cimg/openjdk:17.0
+    resource_class: large
     environment:
-      TEST_PROJECT=ratatoolDiffy
+      SCALA_VERSION: 2.12.20
+      TEST_PROJECT: ratatoolDiffy
     steps:
       - build
   sampling_212:
-    executor: scala_212
+    docker:
+      - image: cimg/openjdk:17.0
+    resource_class: large
     environment:
-      TEST_PROJECT=ratatoolSampling
+      SCALA_VERSION: 2.12.20
+      TEST_PROJECT: ratatoolSampling
     steps:
       - build
   shapeless_212:
@@ -102,15 +108,21 @@ jobs:
     steps:
       - build
   diffy_213:
-    executor: scala_213
+    docker:
+      - image: cimg/openjdk:17.0
+    resource_class: large
     environment:
-      TEST_PROJECT=ratatoolDiffy
+      SCALA_VERSION: 2.13.16
+      TEST_PROJECT: ratatoolDiffy
     steps:
       - build
   sampling_213:
-    executor: scala_213
+    docker:
+      - image: cimg/openjdk:17.0
+    resource_class: large
     environment:
-      TEST_PROJECT=ratatoolSampling
+      SCALA_VERSION: 2.13.16
+      TEST_PROJECT: ratatoolSampling
     steps:
       - build
   shapeless_213:

ratatool-diffy/src/test/scala/com/spotify/ratatool/diffy/BigDiffyTest.scala

@@ -445,6 +445,75 @@ class BigDiffyTest extends PipelineSpec {
     sc.run()
   }
 
+  it should "handle Parquet schema evolution with required array field" in {
+    val schemaBase = new Schema.Parser().parse(
+      """|{"type":"record",
+       |"name":"ParquetRecord",
+       |"namespace":"com.spotify.ratatool.diffy",
+       |"fields":[{"name":"id","type":"int"}]}
+       """.stripMargin
+    )
+    val schemaWithArray = new Schema.Parser().parse(
+      """|{"type":"record",
+       |"name":"ParquetRecord",
+       |"namespace":"com.spotify.ratatool.diffy",
+       |"fields":[
+       |{"name":"id","type":"int"},
+       |{"name":"tags","type":{"type":"array","items":"string"},"default":[]}]}
+       """.stripMargin
+    )
+
+    def toGenericRecord(schema: Schema, fields: Map[String, _]): GenericRecord = {
+      val gr = new GenericData.Record(schema)
+      fields.foreach { case (k, v) => gr.put(k, v) }
+      gr
+    }
+    val (lhsPath, rhsPath) = (
+      ParquetTestData.createTempDir("lhs-array") + "/out.parquet",
+      ParquetTestData.createTempDir("rhs-array") + "/out.parquet"
+    )
+
+    ParquetIO.writeToFile(
+      (1 to 5).map(i =>
+        toGenericRecord(
+          schemaWithArray,
+          Map("id" -> i, "tags" -> java.util.Collections.singletonList(s"tag$i"))
+        )
+      ),
+      schemaWithArray,
+      lhsPath
+    )
+
+    ParquetIO.writeToFile(
+      (1 to 5).map(i => toGenericRecord(schemaBase, Map("id" -> i))),
+      schemaBase,
+      rhsPath
+    )
+
+    // Schema transformation: tags field should be wrapped in union with null
+    val compatSchema = ParquetIO.getCompatibleSchemaForFiles(lhsPath, rhsPath)
+    val tagsField = compatSchema.getField("tags")
+    tagsField should not be null
+    tagsField.schema().getType shouldBe Schema.Type.UNION
+    tagsField.schema().getTypes.asScala.map(_.getType) should contain(Schema.Type.NULL)
+    tagsField.schema().getTypes.asScala.map(_.getType) should contain(Schema.Type.ARRAY)
+    tagsField.hasDefaultValue shouldBe true
+
+    // Simulate what the parquet reader produces for a record missing the tags column:
+    // a GenericRecord with the compatible schema where tags is null.
+    // Without the nullable union fix, encoding this record NPEs in
+    // GenericDatumWriter.getArraySize because it calls ((Collection) null).size().
+    val beamCoder = CoderMaterializer.beamWithDefault(avroGenericRecordCoder(compatSchema))
+    val simulatedRhsRecord = new GenericData.Record(compatSchema)
+    simulatedRhsRecord.put("id", 1)
+    simulatedRhsRecord.put("tags", null)
+
+    val encoded = CoderUtils.encodeToByteArray(beamCoder, simulatedRhsRecord)
+    val decoded = CoderUtils.decodeFromByteArray(beamCoder, encoded)
+    decoded.get("id") shouldBe 1
+    decoded.get("tags") shouldBe null
+  }
+
   "mergeTableSchema" should "merge two schemas" in {
 
     def jl[T](x: T*): java.util.List[T] = List(x: _*).asJava

ratatool-sampling/src/main/scala/com/spotify/ratatool/io/ParquetIO.scala

@@ -62,7 +62,9 @@ object ParquetIO {
         .checkReaderWriterCompatibility(s1, s2)
         .getType == SchemaCompatibilityType.COMPATIBLE
 
-    if (isReadCompatible(schemaLhs, schemaRhs)) {
+    val compatSchema = if (schemaLhs == schemaRhs) {
+      schemaLhs
+    } else if (isReadCompatible(schemaLhs, schemaRhs)) {
       schemaLhs
     } else if (isReadCompatible(schemaRhs, schemaLhs)) {
       schemaRhs
@@ -72,8 +74,58 @@ object ParquetIO {
           s"$path2: $schemaLhs not compatible with $schemaRhs"
       )
     }
+
+    makeCollectionFieldsNullable(compatSchema)
   }
 
+  // Parquet can produce null values for non-nullable array/map fields (e.g. when
+  // reading with a schema that has more fields than the file, or when the underlying
+  // data simply contains nulls). Avro's GenericDatumWriter.getArraySize calls
+  // ((Collection) array).size() with no null guard, causing an NPE. This wraps all
+  // non-nullable collection fields in a union with null so the coder tolerates it.
+  private[ratatool] def makeCollectionFieldsNullable(schema: Schema): Schema = {
+    val collectionTypes = Set(Schema.Type.ARRAY, Schema.Type.MAP)
+    val hasFieldsToFix = schema.getFields.asScala.exists { field =>
+      collectionTypes.contains(field.schema().getType) ||
+      (field.schema().getType == Schema.Type.RECORD)
+    }
+
+    if (!hasFieldsToFix) return schema
+
+    val newFields = schema.getFields.asScala.map { field =>
+      val fieldType = field.schema().getType
+      if (collectionTypes.contains(fieldType) && !isNullableSchema(field.schema())) {
+        // Original type first so any existing default value stays valid;
+        // null second so the coder tolerates nulls from the parquet reader.
+        val nullableType =
+          Schema.createUnion(field.schema(), Schema.create(Schema.Type.NULL))
+        new Schema.Field(
+          field.name(),
+          nullableType,
+          field.doc(),
+          field.defaultVal()
+        )
+      } else if (fieldType == Schema.Type.RECORD) {
+        val nested = makeCollectionFieldsNullable(field.schema())
+        new Schema.Field(field.name(), nested, field.doc(), field.defaultVal())
+      } else {
+        new Schema.Field(field.name(), field.schema(), field.doc(), field.defaultVal())
+      }
+    }
+
+    Schema.createRecord(
+      schema.getName,
+      schema.getDoc,
+      schema.getNamespace,
+      schema.isError,
+      newFields.asJava
+    )
+  }
+
+  private def isNullableSchema(schema: Schema): Boolean =
+    schema.getType == Schema.Type.UNION &&
+      schema.getTypes.asScala.exists(_.getType == Schema.Type.NULL)
+
   private[ratatool] def genericRecordReadConfig(schema: Schema, path: String): Configuration = {
     val job = Job.getInstance(new Configuration())
 

ratatool-sampling/src/test/scala/com/spotify/ratatool/io/ParquetIOTest.scala

@@ -32,6 +32,7 @@ import org.scalatest.matchers.should.Matchers
 
 import java.io.File
 import java.nio.file.Files
+import scala.jdk.CollectionConverters._
 
 class ParquetIOTest extends AnyFlatSpec with Matchers with BeforeAndAfterAll {
   private lazy val (typedOut, avroOut) =
@@ -64,6 +65,58 @@ class ParquetIOTest extends AnyFlatSpec with Matchers with BeforeAndAfterAll {
     ) shouldEqual ParquetTestData.avroSchema
   }
 
+  it should "make non-nullable array and map fields nullable" in {
+    val schema = new Schema.Parser().parse(
+      """|{"type":"record","name":"TestRecord","namespace":"com.spotify.ratatool.io",
+       |"fields":[
+       |{"name":"id","type":"int"},
+       |{"name":"tags","type":{"type":"array","items":"string"},"default":[]},
+       |{"name":"meta","type":{"type":"map","values":"string"},"default":{}}
+       |]}""".stripMargin
+    )
+
+    val result = ParquetIO.makeCollectionFieldsNullable(schema)
+
+    result.getField("id").schema().getType shouldBe Schema.Type.INT
+
+    val tagsField = result.getField("tags")
+    tagsField.schema().getType shouldBe Schema.Type.UNION
+    tagsField.schema().getTypes.asScala.map(_.getType) should contain(Schema.Type.NULL)
+    tagsField.schema().getTypes.asScala.map(_.getType) should contain(Schema.Type.ARRAY)
+    tagsField.hasDefaultValue shouldBe true
+
+    val metaField = result.getField("meta")
+    metaField.schema().getType shouldBe Schema.Type.UNION
+    metaField.schema().getTypes.asScala.map(_.getType) should contain(Schema.Type.NULL)
+    metaField.schema().getTypes.asScala.map(_.getType) should contain(Schema.Type.MAP)
+    metaField.hasDefaultValue shouldBe true
+  }
+
+  it should "preserve already-nullable fields in makeCollectionFieldsNullable" in {
+    val schema = new Schema.Parser().parse(
+      """|{"type":"record","name":"TestRecord","namespace":"com.spotify.ratatool.io",
+       |"fields":[
+       |{"name":"id","type":"int"},
+       |{"name":"label","type":["null","string"],"default":null}
+       |]}""".stripMargin
+    )
+
+    val result = ParquetIO.makeCollectionFieldsNullable(schema)
+
+    val labelField = result.getField("label")
+    labelField.schema().getType shouldBe Schema.Type.UNION
+    labelField.schema().getTypes should have size 2
+  }
+
+  it should "return original schema when no collection fields exist" in {
+    val schema = new Schema.Parser().parse(
+      """|{"type":"record","name":"TestRecord","namespace":"com.spotify.ratatool.io",
+       |"fields":[{"name":"id","type":"int"},{"name":"name","type":"string"}]}""".stripMargin
+    )
+
+    ParquetIO.makeCollectionFieldsNullable(schema) shouldBe schema
+  }
+
   it should "write parquet-avro as GenericRecords to file" in {
     val writePath = ParquetTestData.createTempDir("avro-write") + "/out.parquet"