NullFieldSplitter extends SplitterTransform<StructuredRecord, StructuredRecord> { @VisibleForTesting static Schema getNonNullSchema(Schema nullableSchema, String fieldName, @Nullable FailureCollector collector) { List<Schema.Field> fields = new ArrayList<>(nullableSchema.getFields().size()); for (Schema.Field field : nullableSchema.getFields()) { Schema fieldSchema = field.getSchema(); if (!field.getName().equals(fieldName) || fieldSchema.getType() != Schema.Type.UNION) { fields.add(field); continue; } List<Schema> unionSchemas = fieldSchema.getUnionSchemas(); List<Schema> fieldSchemas = new ArrayList<>(unionSchemas.size()); for (Schema unionSchema : unionSchemas) { if (unionSchema.getType() != Schema.Type.NULL) { fieldSchemas.add(unionSchema); } } if (fieldSchemas.isEmpty()) { if (collector != null) { collector.addFailure( String.format("Field '%s' does not contain a non-null type in its union schema.", fieldName), null) .withInputSchemaField(fieldName); } else { throw new IllegalArgumentException( String.format("Field '%s' does not contain a non-null type in its union schema.", fieldName)); } } else if (fieldSchemas.size() == 1) { fields.add(Schema.Field.of(fieldName, fieldSchemas.iterator().next())); } else { fields.add(Schema.Field.of(fieldName, Schema.unionOf(fieldSchemas))); } } return Schema.recordOf(nullableSchema.getRecordName() + ".nonnull", fields); } NullFieldSplitter(Conf conf); @Override void configurePipeline(MultiOutputPipelineConfigurer multiOutputPipelineConfigurer); @Override void initialize(TransformContext context); @Override void transform(StructuredRecord record, MultiOutputEmitter<StructuredRecord> emitter); static final String NULL_PORT; static final String NON_NULL_PORT; }
@Test public void testGetNonNullSchema() { Schema schema = Schema.recordOf("abc", Schema.Field.of("a", Schema.nullableOf(Schema.of(Schema.Type.STRING))), Schema.Field.of("b", Schema.nullableOf(Schema.of(Schema.Type.DOUBLE))), Schema.Field.of("c", Schema.unionOf( Schema.of(Schema.Type.NULL), Schema.of(Schema.Type.BYTES), Schema.of(Schema.Type.INT), Schema.of(Schema.Type.LONG)))); Schema expected = Schema.recordOf("abc.nonnull", Schema.Field.of("a", Schema.of(Schema.Type.STRING)), Schema.Field.of("b", Schema.nullableOf(Schema.of(Schema.Type.DOUBLE))), Schema.Field.of("c", Schema.unionOf( Schema.of(Schema.Type.NULL), Schema.of(Schema.Type.BYTES), Schema.of(Schema.Type.INT), Schema.of(Schema.Type.LONG)))); Assert.assertEquals(expected, NullFieldSplitter.getNonNullSchema(schema, "a", null)); expected = Schema.recordOf("abc.nonnull", Schema.Field.of("a", Schema.nullableOf(Schema.of(Schema.Type.STRING))), Schema.Field.of("b", Schema.of(Schema.Type.DOUBLE)), Schema.Field.of("c", Schema.unionOf( Schema.of(Schema.Type.NULL), Schema.of(Schema.Type.BYTES), Schema.of(Schema.Type.INT), Schema.of(Schema.Type.LONG)))); Assert.assertEquals(expected, NullFieldSplitter.getNonNullSchema(schema, "b", null)); expected = Schema.recordOf("abc.nonnull", Schema.Field.of("a", Schema.nullableOf(Schema.of(Schema.Type.STRING))), Schema.Field.of("b", Schema.nullableOf(Schema.of(Schema.Type.DOUBLE))), Schema.Field.of("c", Schema.unionOf( Schema.of(Schema.Type.BYTES), Schema.of(Schema.Type.INT), Schema.of(Schema.Type.LONG)))); Assert.assertEquals(expected, NullFieldSplitter.getNonNullSchema(schema, "c", null)); }