@Test public void testSplitOnField() throws Exception { Schema rec1Schema = Schema.recordOf( "rec1", Schema.Field.of("x", Schema.of(Schema.Type.INT)), Schema.Field.of("y", Schema.of(Schema.Type.INT))); Schema rec2Schema = Schema.recordOf( "rec2", Schema.Field.of("y", Schema.of(Schema.Type.INT)), Schema.Field.of("z", Schema.of(Schema.Type.INT))); StructuredRecord rec1 = StructuredRecord.builder(rec1Schema).set("x", 0).set("y", 1).build(); StructuredRecord rec2 = StructuredRecord.builder(rec2Schema).set("y", 2).set("z", 3).build(); Schema inputSchema = Schema.recordOf( "union", Schema.Field.of("a", Schema.of(Schema.Type.LONG)), Schema.Field.of("b", Schema.unionOf(Schema.of(Schema.Type.NULL), Schema.of(Schema.Type.BOOLEAN), Schema.of(Schema.Type.INT), Schema.of(Schema.Type.LONG), Schema.of(Schema.Type.FLOAT), Schema.of(Schema.Type.DOUBLE), Schema.of(Schema.Type.STRING), rec1Schema, rec2Schema))); Schema nullSchema = Schema.recordOf("union.null", Schema.Field.of("a", Schema.of(Schema.Type.LONG)), Schema.Field.of("b", Schema.of(Schema.Type.NULL))); Schema boolSchema = Schema.recordOf("union.boolean", Schema.Field.of("a", Schema.of(Schema.Type.LONG)), Schema.Field.of("b", Schema.of(Schema.Type.BOOLEAN))); Schema intSchema = Schema.recordOf("union.int", Schema.Field.of("a", Schema.of(Schema.Type.LONG)), Schema.Field.of("b", Schema.of(Schema.Type.INT))); Schema longSchema = Schema.recordOf("union.long", Schema.Field.of("a", Schema.of(Schema.Type.LONG)), Schema.Field.of("b", Schema.of(Schema.Type.LONG))); Schema floatSchema = Schema.recordOf("union.float", Schema.Field.of("a", Schema.of(Schema.Type.LONG)), Schema.Field.of("b", Schema.of(Schema.Type.FLOAT))); Schema doubleSchema = Schema.recordOf("union.double", Schema.Field.of("a", Schema.of(Schema.Type.LONG)), Schema.Field.of("b", Schema.of(Schema.Type.DOUBLE))); Schema stringSchema = Schema.recordOf("union.string", Schema.Field.of("a", Schema.of(Schema.Type.LONG)), Schema.Field.of("b", Schema.of(Schema.Type.STRING))); Schema withRec1Schema = Schema.recordOf("union.rec1", Schema.Field.of("a", Schema.of(Schema.Type.LONG)), Schema.Field.of("b", rec1Schema)); Schema withRec2Schema = Schema.recordOf("union.rec2", Schema.Field.of("a", Schema.of(Schema.Type.LONG)), Schema.Field.of("b", rec2Schema)); UnionSplitter unionSplitter = new UnionSplitter(new UnionSplitter.Conf("b", true)); TransformContext context = new MockTransformContext(); unionSplitter.initialize(context); MockMultiOutputEmitter<StructuredRecord> mockEmitter = new MockMultiOutputEmitter<>(); unionSplitter.transform(StructuredRecord.builder(inputSchema) .set("a", 0L) .set("b", null) .build(), mockEmitter); unionSplitter.transform(StructuredRecord.builder(inputSchema) .set("a", 0L) .set("b", true) .build(), mockEmitter); unionSplitter.transform(StructuredRecord.builder(inputSchema) .set("a", 0L) .set("b", new byte[] { 5 }) .build(), mockEmitter); unionSplitter.transform(StructuredRecord.builder(inputSchema) .set("a", 0L) .set("b", 5) .build(), mockEmitter); unionSplitter.transform(StructuredRecord.builder(inputSchema) .set("a", 0L) .set("b", 5L) .build(), mockEmitter); unionSplitter.transform(StructuredRecord.builder(inputSchema) .set("a", 0L) .set("b", 5.5f) .build(), mockEmitter); unionSplitter.transform(StructuredRecord.builder(inputSchema) .set("a", 0L) .set("b", 5.5d) .build(), mockEmitter); unionSplitter.transform(StructuredRecord.builder(inputSchema) .set("a", 0L) .set("b", "5") .build(), mockEmitter); unionSplitter.transform(StructuredRecord.builder(inputSchema) .set("a", 0L) .set("b", rec1) .build(), mockEmitter); unionSplitter.transform(StructuredRecord.builder(inputSchema) .set("a", 0L) .set("b", rec2) .build(), mockEmitter); Map<String, List<StructuredRecord>> expected = new HashMap<>(); expected.put("null", ImmutableList.of(StructuredRecord.builder(nullSchema) .set("a", 0L) .set("b", null) .build())); expected.put("boolean", ImmutableList.of(StructuredRecord.builder(boolSchema) .set("a", 0L) .set("b", true) .build())); expected.put("int", ImmutableList.of(StructuredRecord.builder(intSchema) .set("a", 0L) .set("b", 5) .build())); expected.put("long", ImmutableList.of(StructuredRecord.builder(longSchema) .set("a", 0L) .set("b", 5L) .build())); expected.put("float", ImmutableList.of(StructuredRecord.builder(floatSchema) .set("a", 0L) .set("b", 5.5f) .build())); expected.put("double", ImmutableList.of(StructuredRecord.builder(doubleSchema) .set("a", 0L) .set("b", 5.5d) .build())); expected.put("string", ImmutableList.of(StructuredRecord.builder(stringSchema) .set("a", 0L) .set("b", "5") .build())); expected.put("rec1", ImmutableList.of(StructuredRecord.builder(withRec1Schema) .set("a", 0L) .set("b", rec1) .build())); expected.put("rec2", ImmutableList.of(StructuredRecord.builder(withRec2Schema) .set("a", 0L) .set("b", rec2) .build())); Map<String, List<Object>> actual = mockEmitter.getEmitted(); Assert.assertEquals(expected, actual); Assert.assertEquals(0, context.getFailureCollector().getValidationFailures().size()); }
@Override public void transform(StructuredRecord record, MultiOutputEmitter<StructuredRecord> emitter) { if (conf.unionField == null) { emitter.emit(record.getSchema().getRecordName(), record); return; } Schema.Field schemaField = record.getSchema().getField(conf.unionField); if (schemaField == null) { emitter.emitError(new InvalidEntry<>(100, String.format("Field '%s' does not exist.", conf.unionField), record)); return; } Schema fieldSchema = schemaField.getSchema(); if (fieldSchema.getType() != Schema.Type.UNION) { emitter.emitError(new InvalidEntry<>(200, String.format("Field '%s' is not of type union, but is of type '%s'.", conf.unionField, fieldSchema.getType()), record)); return; } Object val = record.get(conf.unionField); Schema valSchema; if (val == null) { valSchema = Schema.of(Schema.Type.NULL); } else if (val instanceof Boolean) { valSchema = Schema.of(Schema.Type.BOOLEAN); } else if (val instanceof ByteBuffer || val instanceof byte[] || val instanceof Byte[]) { valSchema = Schema.of(Schema.Type.BYTES); } else if (val instanceof Integer) { valSchema = Schema.of(Schema.Type.INT); } else if (val instanceof Long) { valSchema = Schema.of(Schema.Type.LONG); } else if (val instanceof Float) { valSchema = Schema.of(Schema.Type.FLOAT); } else if (val instanceof Double) { valSchema = Schema.of(Schema.Type.DOUBLE); } else if (val instanceof String) { valSchema = Schema.of(Schema.Type.STRING); } else if (val instanceof StructuredRecord) { valSchema = ((StructuredRecord) val).getSchema(); } else if (val.getClass().isEnum()) { emitter.emitError( new InvalidEntry<>(300, String.format("Field '%s' is an Enum, which is not supported.", conf.unionField), record)); return; } else if (val instanceof Map) { emitter.emitError( new InvalidEntry<>(301, String.format("Field '%s' is a Map, which is not supported.", conf.unionField), record)); return; } else if (val instanceof Collection) { emitter.emitError( new InvalidEntry<>(302, String.format("Field '%s' is an array, which is not supported.", conf.unionField), record)); return; } else { emitter.emitError( new InvalidEntry<>(303, String.format("Could not determine type for field '%s' with value of class '%s'.", conf.unionField, val.getClass().getName()), record)); return; } boolean foundSchema = false; for (Schema unionSchema : fieldSchema.getUnionSchemas()) { if (unionSchema.equals(valSchema)) { foundSchema = true; } } if (!foundSchema) { emitter.emitError( new InvalidEntry<>(400, String.format("Field '%s' has schema '%s', which is not in its union schema.", conf.unionField, valSchema), record)); return; } Schema inputSchema = record.getSchema(); List<Schema.Field> fields = new ArrayList<>(inputSchema.getFields().size()); for (Schema.Field inputSchemaField : inputSchema.getFields()) { String fieldName = inputSchemaField.getName(); if (fieldName.equals(conf.unionField)) { fields.add(Schema.Field.of(fieldName, valSchema)); } else { fields.add(inputSchemaField); } } Schema.Type valType = valSchema.getType(); String port = valType == Schema.Type.RECORD ? valSchema.getRecordName() : valType.name().toLowerCase(); Schema outputSchema = conf.modifySchema ? Schema.recordOf(inputSchema.getRecordName() + "." + port, fields) : inputSchema; StructuredRecord.Builder builder = StructuredRecord.builder(outputSchema); for (Schema.Field inputSchemaField : inputSchema.getFields()) { String fieldName = inputSchemaField.getName(); builder.set(fieldName, record.get(fieldName)); } emitter.emit(port, builder.build()); }
UnionSplitter extends SplitterTransform<StructuredRecord, StructuredRecord> { @Override public void transform(StructuredRecord record, MultiOutputEmitter<StructuredRecord> emitter) { if (conf.unionField == null) { emitter.emit(record.getSchema().getRecordName(), record); return; } Schema.Field schemaField = record.getSchema().getField(conf.unionField); if (schemaField == null) { emitter.emitError(new InvalidEntry<>(100, String.format("Field '%s' does not exist.", conf.unionField), record)); return; } Schema fieldSchema = schemaField.getSchema(); if (fieldSchema.getType() != Schema.Type.UNION) { emitter.emitError(new InvalidEntry<>(200, String.format("Field '%s' is not of type union, but is of type '%s'.", conf.unionField, fieldSchema.getType()), record)); return; } Object val = record.get(conf.unionField); Schema valSchema; if (val == null) { valSchema = Schema.of(Schema.Type.NULL); } else if (val instanceof Boolean) { valSchema = Schema.of(Schema.Type.BOOLEAN); } else if (val instanceof ByteBuffer || val instanceof byte[] || val instanceof Byte[]) { valSchema = Schema.of(Schema.Type.BYTES); } else if (val instanceof Integer) { valSchema = Schema.of(Schema.Type.INT); } else if (val instanceof Long) { valSchema = Schema.of(Schema.Type.LONG); } else if (val instanceof Float) { valSchema = Schema.of(Schema.Type.FLOAT); } else if (val instanceof Double) { valSchema = Schema.of(Schema.Type.DOUBLE); } else if (val instanceof String) { valSchema = Schema.of(Schema.Type.STRING); } else if (val instanceof StructuredRecord) { valSchema = ((StructuredRecord) val).getSchema(); } else if (val.getClass().isEnum()) { emitter.emitError( new InvalidEntry<>(300, String.format("Field '%s' is an Enum, which is not supported.", conf.unionField), record)); return; } else if (val instanceof Map) { emitter.emitError( new InvalidEntry<>(301, String.format("Field '%s' is a Map, which is not supported.", conf.unionField), record)); return; } else if (val instanceof Collection) { emitter.emitError( new InvalidEntry<>(302, String.format("Field '%s' is an array, which is not supported.", conf.unionField), record)); return; } else { emitter.emitError( new InvalidEntry<>(303, String.format("Could not determine type for field '%s' with value of class '%s'.", conf.unionField, val.getClass().getName()), record)); return; } boolean foundSchema = false; for (Schema unionSchema : fieldSchema.getUnionSchemas()) { if (unionSchema.equals(valSchema)) { foundSchema = true; } } if (!foundSchema) { emitter.emitError( new InvalidEntry<>(400, String.format("Field '%s' has schema '%s', which is not in its union schema.", conf.unionField, valSchema), record)); return; } Schema inputSchema = record.getSchema(); List<Schema.Field> fields = new ArrayList<>(inputSchema.getFields().size()); for (Schema.Field inputSchemaField : inputSchema.getFields()) { String fieldName = inputSchemaField.getName(); if (fieldName.equals(conf.unionField)) { fields.add(Schema.Field.of(fieldName, valSchema)); } else { fields.add(inputSchemaField); } } Schema.Type valType = valSchema.getType(); String port = valType == Schema.Type.RECORD ? valSchema.getRecordName() : valType.name().toLowerCase(); Schema outputSchema = conf.modifySchema ? Schema.recordOf(inputSchema.getRecordName() + "." + port, fields) : inputSchema; StructuredRecord.Builder builder = StructuredRecord.builder(outputSchema); for (Schema.Field inputSchemaField : inputSchema.getFields()) { String fieldName = inputSchemaField.getName(); builder.set(fieldName, record.get(fieldName)); } emitter.emit(port, builder.build()); } }
UnionSplitter extends SplitterTransform<StructuredRecord, StructuredRecord> { @Override public void transform(StructuredRecord record, MultiOutputEmitter<StructuredRecord> emitter) { if (conf.unionField == null) { emitter.emit(record.getSchema().getRecordName(), record); return; } Schema.Field schemaField = record.getSchema().getField(conf.unionField); if (schemaField == null) { emitter.emitError(new InvalidEntry<>(100, String.format("Field '%s' does not exist.", conf.unionField), record)); return; } Schema fieldSchema = schemaField.getSchema(); if (fieldSchema.getType() != Schema.Type.UNION) { emitter.emitError(new InvalidEntry<>(200, String.format("Field '%s' is not of type union, but is of type '%s'.", conf.unionField, fieldSchema.getType()), record)); return; } Object val = record.get(conf.unionField); Schema valSchema; if (val == null) { valSchema = Schema.of(Schema.Type.NULL); } else if (val instanceof Boolean) { valSchema = Schema.of(Schema.Type.BOOLEAN); } else if (val instanceof ByteBuffer || val instanceof byte[] || val instanceof Byte[]) { valSchema = Schema.of(Schema.Type.BYTES); } else if (val instanceof Integer) { valSchema = Schema.of(Schema.Type.INT); } else if (val instanceof Long) { valSchema = Schema.of(Schema.Type.LONG); } else if (val instanceof Float) { valSchema = Schema.of(Schema.Type.FLOAT); } else if (val instanceof Double) { valSchema = Schema.of(Schema.Type.DOUBLE); } else if (val instanceof String) { valSchema = Schema.of(Schema.Type.STRING); } else if (val instanceof StructuredRecord) { valSchema = ((StructuredRecord) val).getSchema(); } else if (val.getClass().isEnum()) { emitter.emitError( new InvalidEntry<>(300, String.format("Field '%s' is an Enum, which is not supported.", conf.unionField), record)); return; } else if (val instanceof Map) { emitter.emitError( new InvalidEntry<>(301, String.format("Field '%s' is a Map, which is not supported.", conf.unionField), record)); return; } else if (val instanceof Collection) { emitter.emitError( new InvalidEntry<>(302, String.format("Field '%s' is an array, which is not supported.", conf.unionField), record)); return; } else { emitter.emitError( new InvalidEntry<>(303, String.format("Could not determine type for field '%s' with value of class '%s'.", conf.unionField, val.getClass().getName()), record)); return; } boolean foundSchema = false; for (Schema unionSchema : fieldSchema.getUnionSchemas()) { if (unionSchema.equals(valSchema)) { foundSchema = true; } } if (!foundSchema) { emitter.emitError( new InvalidEntry<>(400, String.format("Field '%s' has schema '%s', which is not in its union schema.", conf.unionField, valSchema), record)); return; } Schema inputSchema = record.getSchema(); List<Schema.Field> fields = new ArrayList<>(inputSchema.getFields().size()); for (Schema.Field inputSchemaField : inputSchema.getFields()) { String fieldName = inputSchemaField.getName(); if (fieldName.equals(conf.unionField)) { fields.add(Schema.Field.of(fieldName, valSchema)); } else { fields.add(inputSchemaField); } } Schema.Type valType = valSchema.getType(); String port = valType == Schema.Type.RECORD ? valSchema.getRecordName() : valType.name().toLowerCase(); Schema outputSchema = conf.modifySchema ? Schema.recordOf(inputSchema.getRecordName() + "." + port, fields) : inputSchema; StructuredRecord.Builder builder = StructuredRecord.builder(outputSchema); for (Schema.Field inputSchemaField : inputSchema.getFields()) { String fieldName = inputSchemaField.getName(); builder.set(fieldName, record.get(fieldName)); } emitter.emit(port, builder.build()); } UnionSplitter(Conf conf); }
UnionSplitter extends SplitterTransform<StructuredRecord, StructuredRecord> { @Override public void transform(StructuredRecord record, MultiOutputEmitter<StructuredRecord> emitter) { if (conf.unionField == null) { emitter.emit(record.getSchema().getRecordName(), record); return; } Schema.Field schemaField = record.getSchema().getField(conf.unionField); if (schemaField == null) { emitter.emitError(new InvalidEntry<>(100, String.format("Field '%s' does not exist.", conf.unionField), record)); return; } Schema fieldSchema = schemaField.getSchema(); if (fieldSchema.getType() != Schema.Type.UNION) { emitter.emitError(new InvalidEntry<>(200, String.format("Field '%s' is not of type union, but is of type '%s'.", conf.unionField, fieldSchema.getType()), record)); return; } Object val = record.get(conf.unionField); Schema valSchema; if (val == null) { valSchema = Schema.of(Schema.Type.NULL); } else if (val instanceof Boolean) { valSchema = Schema.of(Schema.Type.BOOLEAN); } else if (val instanceof ByteBuffer || val instanceof byte[] || val instanceof Byte[]) { valSchema = Schema.of(Schema.Type.BYTES); } else if (val instanceof Integer) { valSchema = Schema.of(Schema.Type.INT); } else if (val instanceof Long) { valSchema = Schema.of(Schema.Type.LONG); } else if (val instanceof Float) { valSchema = Schema.of(Schema.Type.FLOAT); } else if (val instanceof Double) { valSchema = Schema.of(Schema.Type.DOUBLE); } else if (val instanceof String) { valSchema = Schema.of(Schema.Type.STRING); } else if (val instanceof StructuredRecord) { valSchema = ((StructuredRecord) val).getSchema(); } else if (val.getClass().isEnum()) { emitter.emitError( new InvalidEntry<>(300, String.format("Field '%s' is an Enum, which is not supported.", conf.unionField), record)); return; } else if (val instanceof Map) { emitter.emitError( new InvalidEntry<>(301, String.format("Field '%s' is a Map, which is not supported.", conf.unionField), record)); return; } else if (val instanceof Collection) { emitter.emitError( new InvalidEntry<>(302, String.format("Field '%s' is an array, which is not supported.", conf.unionField), record)); return; } else { emitter.emitError( new InvalidEntry<>(303, String.format("Could not determine type for field '%s' with value of class '%s'.", conf.unionField, val.getClass().getName()), record)); return; } boolean foundSchema = false; for (Schema unionSchema : fieldSchema.getUnionSchemas()) { if (unionSchema.equals(valSchema)) { foundSchema = true; } } if (!foundSchema) { emitter.emitError( new InvalidEntry<>(400, String.format("Field '%s' has schema '%s', which is not in its union schema.", conf.unionField, valSchema), record)); return; } Schema inputSchema = record.getSchema(); List<Schema.Field> fields = new ArrayList<>(inputSchema.getFields().size()); for (Schema.Field inputSchemaField : inputSchema.getFields()) { String fieldName = inputSchemaField.getName(); if (fieldName.equals(conf.unionField)) { fields.add(Schema.Field.of(fieldName, valSchema)); } else { fields.add(inputSchemaField); } } Schema.Type valType = valSchema.getType(); String port = valType == Schema.Type.RECORD ? valSchema.getRecordName() : valType.name().toLowerCase(); Schema outputSchema = conf.modifySchema ? Schema.recordOf(inputSchema.getRecordName() + "." + port, fields) : inputSchema; StructuredRecord.Builder builder = StructuredRecord.builder(outputSchema); for (Schema.Field inputSchemaField : inputSchema.getFields()) { String fieldName = inputSchemaField.getName(); builder.set(fieldName, record.get(fieldName)); } emitter.emit(port, builder.build()); } UnionSplitter(Conf conf); @Override void configurePipeline(MultiOutputPipelineConfigurer multiOutputPipelineConfigurer); @Override void transform(StructuredRecord record, MultiOutputEmitter<StructuredRecord> emitter); }
UnionSplitter extends SplitterTransform<StructuredRecord, StructuredRecord> { @Override public void transform(StructuredRecord record, MultiOutputEmitter<StructuredRecord> emitter) { if (conf.unionField == null) { emitter.emit(record.getSchema().getRecordName(), record); return; } Schema.Field schemaField = record.getSchema().getField(conf.unionField); if (schemaField == null) { emitter.emitError(new InvalidEntry<>(100, String.format("Field '%s' does not exist.", conf.unionField), record)); return; } Schema fieldSchema = schemaField.getSchema(); if (fieldSchema.getType() != Schema.Type.UNION) { emitter.emitError(new InvalidEntry<>(200, String.format("Field '%s' is not of type union, but is of type '%s'.", conf.unionField, fieldSchema.getType()), record)); return; } Object val = record.get(conf.unionField); Schema valSchema; if (val == null) { valSchema = Schema.of(Schema.Type.NULL); } else if (val instanceof Boolean) { valSchema = Schema.of(Schema.Type.BOOLEAN); } else if (val instanceof ByteBuffer || val instanceof byte[] || val instanceof Byte[]) { valSchema = Schema.of(Schema.Type.BYTES); } else if (val instanceof Integer) { valSchema = Schema.of(Schema.Type.INT); } else if (val instanceof Long) { valSchema = Schema.of(Schema.Type.LONG); } else if (val instanceof Float) { valSchema = Schema.of(Schema.Type.FLOAT); } else if (val instanceof Double) { valSchema = Schema.of(Schema.Type.DOUBLE); } else if (val instanceof String) { valSchema = Schema.of(Schema.Type.STRING); } else if (val instanceof StructuredRecord) { valSchema = ((StructuredRecord) val).getSchema(); } else if (val.getClass().isEnum()) { emitter.emitError( new InvalidEntry<>(300, String.format("Field '%s' is an Enum, which is not supported.", conf.unionField), record)); return; } else if (val instanceof Map) { emitter.emitError( new InvalidEntry<>(301, String.format("Field '%s' is a Map, which is not supported.", conf.unionField), record)); return; } else if (val instanceof Collection) { emitter.emitError( new InvalidEntry<>(302, String.format("Field '%s' is an array, which is not supported.", conf.unionField), record)); return; } else { emitter.emitError( new InvalidEntry<>(303, String.format("Could not determine type for field '%s' with value of class '%s'.", conf.unionField, val.getClass().getName()), record)); return; } boolean foundSchema = false; for (Schema unionSchema : fieldSchema.getUnionSchemas()) { if (unionSchema.equals(valSchema)) { foundSchema = true; } } if (!foundSchema) { emitter.emitError( new InvalidEntry<>(400, String.format("Field '%s' has schema '%s', which is not in its union schema.", conf.unionField, valSchema), record)); return; } Schema inputSchema = record.getSchema(); List<Schema.Field> fields = new ArrayList<>(inputSchema.getFields().size()); for (Schema.Field inputSchemaField : inputSchema.getFields()) { String fieldName = inputSchemaField.getName(); if (fieldName.equals(conf.unionField)) { fields.add(Schema.Field.of(fieldName, valSchema)); } else { fields.add(inputSchemaField); } } Schema.Type valType = valSchema.getType(); String port = valType == Schema.Type.RECORD ? valSchema.getRecordName() : valType.name().toLowerCase(); Schema outputSchema = conf.modifySchema ? Schema.recordOf(inputSchema.getRecordName() + "." + port, fields) : inputSchema; StructuredRecord.Builder builder = StructuredRecord.builder(outputSchema); for (Schema.Field inputSchemaField : inputSchema.getFields()) { String fieldName = inputSchemaField.getName(); builder.set(fieldName, record.get(fieldName)); } emitter.emit(port, builder.build()); } UnionSplitter(Conf conf); @Override void configurePipeline(MultiOutputPipelineConfigurer multiOutputPipelineConfigurer); @Override void transform(StructuredRecord record, MultiOutputEmitter<StructuredRecord> emitter);  }