[Coral-Schema] Restore ancestor nullability when a nested field is projected out

coral-schema/src/main/java/com/linkedin/coral/schema/avro/RelToAvroSchemaConverter.java

@@ -574,9 +574,19 @@ private void handleFieldAccess(RexFieldAccess rexFieldAccess, RexInputRef refere
       String newFieldName = SchemaUtilities.getFieldName(oldFieldName, suggestNewFieldName);
 
       Schema topSchema = inputSchema.getFields().get(referenceExpr.getIndex()).schema();
-      Schema.Field accessedField = getFieldFromTopSchema(topSchema, oldFieldName, innerRecordNames);
-      assert accessedField != null;
-      SchemaUtilities.appendField(newFieldName, accessedField, fieldAssembler);
+      AccessedField accessedField = getFieldFromTopSchema(topSchema, oldFieldName, innerRecordNames);
+      assert accessedField.field != null;
+
+      // Projecting a leaf out of a nullable ancestor flattens two levels of nullability into one column, so the
+      // column has to carry the union of both: it is null whenever the leaf is null OR any ancestor is absent.
+      // Without this the leaf's own `required` would be published as the column's nullability, producing a schema
+      // that the data violates for every row where an ancestor is null.
+      Schema fieldSchema = accessedField.field.schema();
+      if (accessedField.nullableAncestor) {
+        fieldSchema = SchemaUtilities.reorderOptionIfRequired(SchemaUtilities.makeNullable(fieldSchema, false),
+            SchemaUtilities.defaultValue(accessedField.field));
+      }
+      SchemaUtilities.appendField(newFieldName, accessedField.field, fieldSchema, fieldAssembler);
     }
 
     private void handleUDFFieldAccess(RexFieldAccess rexFieldAccess, RexCall referenceExpr) {
@@ -626,9 +636,13 @@ private void appendField(RelDataType fieldType, boolean isNullable, String doc)
      * Get the field named `fieldName` in the given `topSchema` which might be nested
      * @param innerRecordNames contains inner record names, if `topSchema` contains inner record, we need to retrieve the inner record which
      *                         is the ancestor of field named `fieldName`
+     * @return the resolved field together with whether any ancestor traversed on the way to it was nullable
      */
-    private Schema.Field getFieldFromTopSchema(@Nonnull Schema topSchema, @Nonnull String fieldName,
+    private AccessedField getFieldFromTopSchema(@Nonnull Schema topSchema, @Nonnull String fieldName,
         @Nonnull Deque<String> innerRecordNames) {
+      // Every `extractIfOption` below discards an ancestor's null branch. Record that it happened so the caller can
+      // put the nullability back on the flattened field; otherwise it is lost for good.
+      boolean nullableAncestor = AvroSerdeUtils.isNullableType(topSchema);
       topSchema = SchemaUtilities.extractIfOption(topSchema);
 
       while (topSchema.getType() != Schema.Type.RECORD
@@ -652,6 +666,7 @@ private Schema.Field getFieldFromTopSchema(@Nonnull Schema topSchema, @Nonnull S
           default:
             throw new IllegalArgumentException("Unsupported topSchema type: " + topSchema.getType());
         }
+        nullableAncestor = nullableAncestor || AvroSerdeUtils.isNullableType(topSchema);
         topSchema = SchemaUtilities.extractIfOption(topSchema);
       }
 
@@ -662,7 +677,22 @@ private Schema.Field getFieldFromTopSchema(@Nonnull Schema topSchema, @Nonnull S
           break;
         }
       }
-      return targetField;
+      return new AccessedField(targetField, nullableAncestor);
+    }
+  }
+
+  /**
+   * A field resolved through a (possibly nested) field access, plus whether any ancestor traversed to reach it
+   * was nullable. A leaf may be `required` inside its parent record while the flattened column that projects it
+   * is still nullable, because the parent itself may be absent.
+   */
+  private static class AccessedField {
+    private final Schema.Field field;
+    private final boolean nullableAncestor;
+
+    private AccessedField(Schema.Field field, boolean nullableAncestor) {
+      this.field = field;
+      this.nullableAncestor = nullableAncestor;
     }
   }
 }

coral-schema/src/main/java/com/linkedin/coral/schema/avro/SchemaUtilities.java

@@ -340,13 +340,28 @@ static boolean isFieldNullable(@Nonnull RexCall rexCall, @Nonnull Schema inputSc
 
   static void appendField(@Nonnull String fieldName, @Nonnull Schema.Field field,
       @Nonnull SchemaBuilder.FieldAssembler<Schema> fieldAssembler) {
+    appendField(fieldName, field, field.schema(), fieldAssembler);
+  }
+
+  /**
+   * Appends {@code field} under the name {@code fieldName}, but writes {@code fieldSchema} in place of the
+   * schema the field carries.
+   *
+   * <p>This is needed when a field is lifted out of its original position and its nullability has to change
+   * as a result - for example when a nested field is projected out of a nullable ancestor, where the leaf is
+   * {@code required} within its parent record but the flattened column can still be null because the parent
+   * itself may be absent.
+   */
+  static void appendField(@Nonnull String fieldName, @Nonnull Schema.Field field, @Nonnull Schema fieldSchema,
+      @Nonnull SchemaBuilder.FieldAssembler<Schema> fieldAssembler) {
     Preconditions.checkNotNull(fieldName);
     Preconditions.checkNotNull(field);
+    Preconditions.checkNotNull(fieldSchema);
     Preconditions.checkNotNull(fieldAssembler);
 
     Object defaultValue = defaultValue(field);
 
-    SchemaBuilder.GenericDefault genericDefault = fieldAssembler.name(fieldName).doc(field.doc()).type(field.schema());
+    SchemaBuilder.GenericDefault genericDefault = fieldAssembler.name(fieldName).doc(field.doc()).type(fieldSchema);
     if (defaultValue != null) {
       genericDefault.withDefault(defaultValue);
     } else {

coral-schema/src/test/java/com/linkedin/coral/schema/avro/TestUtils.java

@@ -109,6 +109,7 @@ private static void initializeTables() {
     String basePrimitive = loadSchema("base-primitive.avsc");
     String baseComplexNestedStructSameName = loadSchema("base-complex-nested-struct-same-name.avsc");
     String baseComplexMixedNullabilities = loadSchema("base-complex-mixed-nullabilities.avsc");
+    String baseNullableStructRequiredLeaf = loadSchema("base-nullable-struct-required-leaf.avsc");
 
     executeCreateTableQuery("default", "basecomplex", baseComplexSchema);
     executeCreateTableQuery("default", "basecomplexunioncompatible", baseComplexUnionCompatible);
@@ -125,6 +126,7 @@ private static void initializeTables() {
     executeCreateTableQuery("default", "basecomplexlowercase", baseComplexLowercase);
     executeCreateTableQuery("default", "baseprimitive", basePrimitive);
     executeCreateTableQuery("default", "basecomplexnestedstructsamename", baseComplexNestedStructSameName);
+    executeCreateTableQuery("default", "basenullablestructrequiredleaf", baseNullableStructRequiredLeaf);
     executeCreateTableWithPartitionQuery("default", "basecasepreservation", baseCasePreservation);
     executeCreateTableWithPartitionFieldSchemaQuery("default", "basecomplexfieldschema", baseComplexFieldSchema);
     executeCreateTableWithPartitionQuery("default", "basenestedcomplex", baseNestedComplexSchema);

coral-schema/src/test/java/com/linkedin/coral/schema/avro/ViewToAvroSchemaConverterTests.java

@@ -780,6 +780,112 @@ public void testSelectNestedStructFieldFromNestComplex() {
         TestUtils.loadSchema("testSelectNestedStructFieldFromNestComplex-expected.avsc"));
   }
 
+  @Test
+  public void testFlattenedFieldFromNullableStructIsNullable() {
+    String viewSql = "CREATE VIEW v AS SELECT Job.Region job_region_col, Job.Country job_country_col, "
+        + "Product.Lob product_lob_col FROM basenullablestructrequiredleaf";
+    TestUtils.executeCreateViewQuery("default", "v", viewSql);
+
+    ViewToAvroSchemaConverter viewToAvroSchemaConverter = ViewToAvroSchemaConverter.create(hiveMetastoreClient);
+    Schema actualSchema = viewToAvroSchemaConverter.toAvroSchema("default", "v");
+
+    // `Region` is required inside `JobRecord`, but `Job` itself is nullable. Flattening collapses both levels into
+    // one column, so that column is null on every row where `Job` is absent and has to be declared nullable.
+    Assert.assertTrue(AvroSerdeUtils.isNullableType(actualSchema.getField("job_region_col").schema()));
+    Assert.assertTrue(AvroSerdeUtils.isNullableType(actualSchema.getField("job_country_col").schema()));
+
+    // `Product` is not nullable, so a required leaf projected out of it must stay required - the fix must not
+    // blanket-nullify every flattened field.
+    Assert.assertFalse(AvroSerdeUtils.isNullableType(actualSchema.getField("product_lob_col").schema()));
+    Assert.assertEquals(actualSchema.getField("product_lob_col").schema().getType(), Schema.Type.STRING);
+  }
+
+  @Test
+  public void testFlattenedFieldFromNullableGrandparentIsNullable() {
+    String viewSql = "CREATE VIEW v AS SELECT Wrapper.Mid.Leaf outer_mid_leaf_col FROM basenullablestructrequiredleaf";
+    TestUtils.executeCreateViewQuery("default", "v", viewSql);
+
+    ViewToAvroSchemaConverter viewToAvroSchemaConverter = ViewToAvroSchemaConverter.create(hiveMetastoreClient);
+    Schema actualSchema = viewToAvroSchemaConverter.toAvroSchema("default", "v");
+
+    // `Leaf` is required and its immediate parent `Mid` is not nullable, but `Wrapper` is. Nullability has to
+    // propagate from any ancestor in the chain, not just the immediate parent.
+    Assert.assertTrue(AvroSerdeUtils.isNullableType(actualSchema.getField("outer_mid_leaf_col").schema()));
+  }
+
+  @Test
+  public void testFlattenedFieldFromNullableMidAncestorIsNullable() {
+    String viewSql = "CREATE VIEW v AS SELECT Deep.Mid.Nested.Leaf deep_leaf_col, "
+        + "DeepReq.Mid.Nested.Leaf deep_req_leaf_col FROM basenullablestructrequiredleaf";
+    TestUtils.executeCreateViewQuery("default", "v", viewSql);
+
+    ViewToAvroSchemaConverter viewToAvroSchemaConverter = ViewToAvroSchemaConverter.create(hiveMetastoreClient);
+    Schema actualSchema = viewToAvroSchemaConverter.toAvroSchema("default", "v");
+
+    // `Deep` is required and so is every level below the nullable `Mid`, so the traversal starts with
+    // `nullableAncestor == false`, turns it true on the second hop, and then walks a third required hop before
+    // reaching the leaf. Only an accumulating OR survives that: replacing `||=` with a plain assignment leaves the
+    // final iteration's `false` in place and this assertion fails. `Wrapper.Mid.Leaf` cannot catch that, because
+    // `Wrapper` is nullable at the top and seeds the flag before the loop runs at all.
+    Assert.assertTrue(AvroSerdeUtils.isNullableType(actualSchema.getField("deep_leaf_col").schema()));
+
+    // The same depth with no nullable ancestor anywhere must stay required, so the accumulation cannot be
+    // satisfied by simply nullifying anything reached through more than one hop.
+    Assert.assertFalse(AvroSerdeUtils.isNullableType(actualSchema.getField("deep_req_leaf_col").schema()));
+    Assert.assertEquals(actualSchema.getField("deep_req_leaf_col").schema().getType(), Schema.Type.INT);
+  }
+
+  @Test
+  public void testFlattenedFieldFromNullableStructKeepsDefaultValid() {
+    String viewSql = "CREATE VIEW v AS SELECT Job.Tier job_tier_col FROM basenullablestructrequiredleaf";
+    TestUtils.executeCreateViewQuery("default", "v", viewSql);
+
+    ViewToAvroSchemaConverter viewToAvroSchemaConverter = ViewToAvroSchemaConverter.create(hiveMetastoreClient);
+    Schema actualSchema = viewToAvroSchemaConverter.toAvroSchema("default", "v");
+
+    // `Tier` is required and carries a default. Avro requires a field's default to match the *first* branch of its
+    // union, so naively emitting ["null", "string"] with default "unknown" would produce a schema that no longer
+    // parses. The null branch has to go second instead.
+    Schema tierSchema = actualSchema.getField("job_tier_col").schema();
+    Assert.assertTrue(AvroSerdeUtils.isNullableType(tierSchema));
+    Assert.assertEquals(tierSchema.getTypes().get(0).getType(), Schema.Type.STRING);
+    Assert.assertEquals(tierSchema.getTypes().get(1).getType(), Schema.Type.NULL);
+
+    // The emitted schema must survive a round trip through the Avro parser.
+    Assert.assertEquals(new Schema.Parser().parse(actualSchema.toString()), actualSchema);
+  }
+
+  @Test
+  public void testFlattenedFieldFromNullableArrayAndMapIsNullable() {
+    String viewSql = "CREATE VIEW v AS SELECT Arr[0].Arr_Leaf arr_leaf_col, Mp['x'].Map_Leaf map_leaf_col "
+        + "FROM basenullablestructrequiredleaf";
+    TestUtils.executeCreateViewQuery("default", "v", viewSql);
+
+    ViewToAvroSchemaConverter viewToAvroSchemaConverter = ViewToAvroSchemaConverter.create(hiveMetastoreClient);
+    Schema actualSchema = viewToAvroSchemaConverter.toAvroSchema("default", "v");
+
+    // The traversal walks array element and map value types the same way it walks records, so a nullable array or
+    // map anywhere in the chain has to propagate too.
+    Assert.assertTrue(AvroSerdeUtils.isNullableType(actualSchema.getField("arr_leaf_col").schema()));
+    Assert.assertTrue(AvroSerdeUtils.isNullableType(actualSchema.getField("map_leaf_col").schema()));
+  }
+
+  @Test
+  public void testSelectingNullableStructWholesaleKeepsRequiredLeaf() {
+    String viewSql = "CREATE VIEW v AS SELECT Job FROM basenullablestructrequiredleaf";
+    TestUtils.executeCreateViewQuery("default", "v", viewSql);
+
+    ViewToAvroSchemaConverter viewToAvroSchemaConverter = ViewToAvroSchemaConverter.create(hiveMetastoreClient);
+    Schema actualSchema = viewToAvroSchemaConverter.toAvroSchema("default", "v");
+
+    // No flattening here: the `["null", ...]` envelope still carries the "absent" case, so `Region` stays required
+    // inside the record. Only projecting a leaf out of the struct loses that envelope.
+    Schema jobSchema = actualSchema.getField("Job").schema();
+    Assert.assertTrue(AvroSerdeUtils.isNullableType(jobSchema));
+    Assert.assertFalse(
+        AvroSerdeUtils.isNullableType(SchemaUtilities.extractIfOption(jobSchema).getField("Region").schema()));
+  }
+
   @Test
   public void testSelectDeepNestStructFieldFromDeepNestComplex() {
     String viewSql = "CREATE VIEW v AS SELECT struct_col_1.struct_col_2.struct_col_3.int_field_1 Int_Field_1, "

coral-schema/src/test/resources/base-nullable-struct-required-leaf.avsc

@@ -0,0 +1,157 @@
+{
+  "type": "record",
+  "name": "BookingsRecord",
+  "fields": [
+    {
+      "name": "Job",
+      "type": [ "null", {
+        "type": "record",
+        "name": "JobRecord",
+        "fields": [
+          {
+            "name": "Region",
+            "type": "string"
+          },
+          {
+            "name": "Tier",
+            "type": "string",
+            "default": "unknown"
+          },
+          {
+            "name": "Country",
+            "type": [ "null", "string" ]
+          }
+        ]
+      } ]
+    },
+    {
+      "name": "Product",
+      "type": {
+        "type": "record",
+        "name": "ProductRecord",
+        "fields": [
+          {
+            "name": "Lob",
+            "type": "string"
+          }
+        ]
+      }
+    },
+    {
+      "name": "Wrapper",
+      "type": [ "null", {
+        "type": "record",
+        "name": "WrapperRecord",
+        "fields": [
+          {
+            "name": "Mid",
+            "type": {
+              "type": "record",
+              "name": "MidRecord",
+              "fields": [
+                {
+                  "name": "Leaf",
+                  "type": "int"
+                }
+              ]
+            }
+          }
+        ]
+      } ]
+    },
+    {
+      "name": "Deep",
+      "type": {
+        "type": "record",
+        "name": "DeepARecord",
+        "fields": [
+          {
+            "name": "Mid",
+            "type": [ "null", {
+              "type": "record",
+              "name": "DeepBRecord",
+              "fields": [
+                {
+                  "name": "Nested",
+                  "type": {
+                    "type": "record",
+                    "name": "DeepCRecord",
+                    "fields": [
+                      {
+                        "name": "Leaf",
+                        "type": "int"
+                      }
+                    ]
+                  }
+                }
+              ]
+            } ]
+          }
+        ]
+      }
+    },
+    {
+      "name": "DeepReq",
+      "type": {
+        "type": "record",
+        "name": "DeepReqARecord",
+        "fields": [
+          {
+            "name": "Mid",
+            "type": {
+              "type": "record",
+              "name": "DeepReqBRecord",
+              "fields": [
+                {
+                  "name": "Nested",
+                  "type": {
+                    "type": "record",
+                    "name": "DeepReqCRecord",
+                    "fields": [
+                      {
+                        "name": "Leaf",
+                        "type": "int"
+                      }
+                    ]
+                  }
+                }
+              ]
+            }
+          }
+        ]
+      }
+    },
+    {
+      "name": "Arr",
+      "type": [ "null", {
+        "type": "array",
+        "items": {
+          "type": "record",
+          "name": "ArrayElement",
+          "fields": [
+            {
+              "name": "Arr_Leaf",
+              "type": "int"
+            }
+          ]
+        }
+      } ]
+    },
+    {
+      "name": "Mp",
+      "type": [ "null", {
+        "type": "map",
+        "values": {
+          "type": "record",
+          "name": "MapValue",
+          "fields": [
+            {
+              "name": "Map_Leaf",
+              "type": "int"
+            }
+          ]
+        }
+      } ]
+    }
+  ]
+}