org.apache.avro.SchemaBuilder.FieldAssembler Java Examples

@SuppressWarnings("fallthrough")
private Schema getGranularityBasedSchema() {
  FieldAssembler<Schema> assembler =
      SchemaBuilder.record("GenericRecordTimePartition").namespace("gobblin.writer.partitioner").fields();

  // Construct the fields in reverse order
  if (!Strings.isNullOrEmpty(this.writerPartitionSuffix)) {
    assembler = assembler.name(SUFFIX).type(Schema.create(Schema.Type.STRING)).noDefault();
  }
  assembler = assembler.name(this.granularity.toString()).type(Schema.create(Schema.Type.STRING)).noDefault();

  if (!Strings.isNullOrEmpty(this.writerPartitionPrefix)) {
    assembler = assembler.name(PREFIX).type(Schema.create(Schema.Type.STRING)).noDefault();
  }

  Schema schema = assembler.endRecord();
  Collections.reverse(schema.getFields());
  return schema;
}
 

private Schema initDefaultSchema(Schema designSchema) {
    AvroRegistry avroReg = new AvroRegistry();
    FieldAssembler<Schema> record = SchemaBuilder.record("Main").fields();
    for (SplunkJSONEventField metadataField : SplunkJSONEventField.getMetadataFields()) {
        Schema base = avroReg.getConverter(metadataField.getDataType()).getSchema();
        FieldBuilder<Schema> fieldBuilder = record.name(metadataField.getName());
        if (metadataField.getName().equals(SplunkJSONEventField.TIME.getName())) {
            String datePattern;
            Field designField = designSchema.getField(metadataField.getName());
            if (designField != null) {
                datePattern = designField.getProp(SchemaConstants.TALEND_COLUMN_PATTERN);
            } else {
                datePattern = designSchema.getProp(SchemaConstants.TALEND_COLUMN_PATTERN);
            }
            if (datePattern == null || datePattern.isEmpty()) {
                datePattern = "dd-MM-yyyy";
            }
            fieldBuilder.prop(SchemaConstants.TALEND_COLUMN_PATTERN, datePattern);
        }
        fieldBuilder.type(AvroUtils.wrapAsNullable(base)).noDefault();
    }
    Schema defaultSchema = record.endRecord();
    return defaultSchema;
}
 

public static Schema createTestSchema5(String tablename) {
    FieldAssembler<Schema> builder = SchemaBuilder.builder().record(tablename).fields();

    Schema schema = AvroUtils._int();
    schema = wrap(schema);
    builder = builder.name("ID").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "ID").type(schema).noDefault();

    schema = AvroUtils._string();
    schema = wrap(schema);
    builder = builder.name("NAME").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "NAME").type(schema).noDefault();

    schema = AvroUtils._string();// TODO : fix it as should be object type
    schema = wrap(schema);
    builder = builder.name("RESULTSET").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "RESULTSET").type(schema)
            .noDefault();

    return builder.endRecord();
}
 

public static Schema createTestSchema(String tablename) {
    FieldAssembler<Schema> builder = SchemaBuilder.builder().record(tablename).fields();

    Schema schema = AvroUtils._int();
    schema = wrap(schema);
    builder = builder.name("ID").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "ID").type(schema).noDefault();

    schema = AvroUtils._string();
    schema = wrap(schema);
    builder = builder.name("NAME").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "NAME").type(schema).noDefault();

    return builder.endRecord();
}
 

/**
 * Constructs an avro schema from a pinot schema
 * @param schema
 * @return
 */
public static Schema constructAvroSchemaFromPinotSchema(org.apache.pinot.common.data.Schema schema) {
  Schema avroSchema = null;

  RecordBuilder<Schema> recordBuilder = SchemaBuilder.record("record");
  FieldAssembler<Schema> fieldAssembler = recordBuilder.fields();

  for (FieldSpec fieldSpec : schema.getAllFieldSpecs()) {
    String fieldName = fieldSpec.getName();
    DataType dataType = fieldSpec.getDataType();
    BaseFieldTypeBuilder<Schema> baseFieldTypeBuilder = fieldAssembler.name(fieldName).type().nullable();
    switch (dataType) {
      case BOOLEAN:
        fieldAssembler = baseFieldTypeBuilder.booleanType().noDefault();
        break;
      case DOUBLE:
        fieldAssembler = baseFieldTypeBuilder.doubleType().noDefault();
        break;
      case FLOAT:
        fieldAssembler = baseFieldTypeBuilder.floatType().noDefault();
        break;
      case INT:
        fieldAssembler = baseFieldTypeBuilder.intType().noDefault();
        break;
      case LONG:
        fieldAssembler = baseFieldTypeBuilder.longType().noDefault();
        break;
      case STRING:
        fieldAssembler = baseFieldTypeBuilder.stringType().noDefault();
        break;
      default:
        break;
    }
  }

  avroSchema = fieldAssembler.endRecord();
  LOGGER.info("Avro Schema {}", avroSchema.toString(true));

  return avroSchema;
}
 

private Schema getDateTimeFormatBasedSchema() {
  FieldAssembler<Schema> assembler =
      SchemaBuilder.record("GenericRecordTimePartition").namespace("gobblin.writer.partitioner").fields();

  if (!Strings.isNullOrEmpty(this.writerPartitionPrefix)) {
    assembler = assembler.name(PREFIX).type(Schema.create(Schema.Type.STRING)).noDefault();
  }
  assembler = assembler.name(PARTITIONED_PATH).type(Schema.create(Schema.Type.STRING)).noDefault();
  if (!Strings.isNullOrEmpty(this.writerPartitionSuffix)) {
    assembler = assembler.name(SUFFIX).type(Schema.create(Schema.Type.STRING)).noDefault();
  }

  return assembler.endRecord();
}
 

public static Schema createSPSchema3(String tablename) {
    FieldAssembler<Schema> builder = SchemaBuilder.builder().record(tablename).fields();

    Schema schema = AvroUtils._int();
    schema = wrap(schema);
    builder = builder.name("PARAMETER1").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "PARAMETER1").type(schema)
            .noDefault();

    schema = AvroUtils._string();
    schema = wrap(schema);
    builder = builder.name("PARAMETER2").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "PARAMETER2").type(schema)
            .noDefault();

    return builder.endRecord();
}
 

public static Schema createSPSchema2(String tablename) {
    FieldAssembler<Schema> builder = SchemaBuilder.builder().record(tablename).fields();

    Schema schema = AvroUtils._string();
    schema = wrap(schema);
    builder = builder.name("PARAMETER").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "PARAMETER").type(schema)
            .noDefault();

    return builder.endRecord();
}
 

public static Schema createSPSchema1(String tablename) {
    FieldAssembler<Schema> builder = SchemaBuilder.builder().record(tablename).fields();

    Schema schema = AvroUtils._int();
    schema = wrap(schema);
    builder = builder.name("PARAMETER").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "PARAMETER").type(schema)
            .noDefault();

    return builder.endRecord();
}
 

public static Schema createTestSchema4(String tablename) {
    FieldAssembler<Schema> builder = SchemaBuilder.builder().record(tablename).fields();

    Schema schema = AvroUtils._string();// TODO : fix it as should be object type
    schema = wrap(schema);
    builder = builder.name("RESULTSET").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "RESULTSET").type(schema)
            .noDefault();

    return builder.endRecord();
}
 

public static void createEmptyAvroStream(final OutputStream outStream) throws IOException {
    final FieldAssembler<Schema> builder = SchemaBuilder.record("NiFi_ExecuteSQL_Record").namespace("any.data").fields();
    final Schema schema = builder.endRecord();

    final DatumWriter<GenericRecord> datumWriter = new GenericDatumWriter<>(schema);
    try (final DataFileWriter<GenericRecord> dataFileWriter = new DataFileWriter<>(datumWriter)) {
        dataFileWriter.create(schema, outStream);
    }
}
 

public static Schema createTestSchema2(String tablename) {
    FieldAssembler<Schema> builder = SchemaBuilder.builder().record(tablename).fields();

    Schema schema = AvroUtils._int();
    schema = wrap(schema);
    builder = builder.name("ID").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "ID")
            .prop(SchemaConstants.TALEND_COLUMN_IS_KEY, "true").type(schema).noDefault();

    schema = AvroUtils._string();
    schema = wrap(schema);
    builder = builder.name("NAME").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "NAME").type(schema).noDefault();

    return builder.endRecord();
}
 

private static Schema getAvroSchema(Map<String, String> fields)
{
    Schema.Parser parser = new Schema.Parser();
    FieldAssembler<Schema> fieldAssembler = getFieldBuilder();
    for (Map.Entry<String, String> field : fields.entrySet()) {
        FieldBuilder<Schema> fieldBuilder = fieldAssembler.name(field.getKey());
        Schema fieldSchema = parser.parse(field.getValue());
        GenericDefault<Schema> genericDefault = fieldBuilder.type(fieldSchema);
        switch (fieldSchema.getType()) {
            case ARRAY:
                genericDefault.withDefault(ImmutableList.of());
                break;
            case MAP:
                genericDefault.withDefault(ImmutableMap.of());
                break;
            case UNION:
                if (fieldSchema.getTypes().stream()
                        .map(Schema::getType)
                        .anyMatch(Schema.Type.NULL::equals)) {
                    genericDefault.withDefault(null);
                }
                else {
                    genericDefault.noDefault();
                }
                break;
            case NULL:
                genericDefault.withDefault(null);
            default:
                genericDefault.noDefault();
        }
    }
    return fieldAssembler.endRecord();
}
 

private Schema createTestSchema() {
    FieldAssembler<Schema> builder = SchemaBuilder.builder().record("TEST").fields();

    Schema schema = AvroUtils._int();
    schema = wrap(schema);
    builder = builder.name("ID").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "ID1")
            .prop(SchemaConstants.TALEND_COLUMN_IS_KEY, "true").type(schema).noDefault();

    schema = AvroUtils._string();
    schema = wrap(schema);
    builder = builder.name("NAME").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "NAME1").type(schema).noDefault();

    return builder.endRecord();
}
 

public Schema getMakeRowSchema(boolean isDynamic) {
    FieldAssembler<Schema> fa = SchemaBuilder.builder().record("MakeRowRecord").fields() //
            .name("Name").type().nullable().stringType().noDefault() //
            .name("ShippingStreet").type().nullable().stringType().noDefault() //
            .name("ShippingPostalCode").type().nullable().intType().noDefault() //
            .name("BillingStreet").type().nullable().stringType().noDefault() //
            .name("BillingState").type().nullable().stringType().noDefault() //
            .name("BillingPostalCode").type().nullable().stringType().noDefault();
    if (isDynamic) {
        fa = fa.name("ShippingState").type().nullable().stringType().noDefault();
    }

    return fa.endRecord();
}
 

/**
 * Create dynamic schema with date pattern in dynamic column. Required to test parsing of Date presented as String.
 */
private Schema createDynamicSchemaWithDatePattern() {
    FieldAssembler<Schema> record = SchemaBuilder.record("Main").fields();
    Schema defaultSchema = record.endRecord();
    AvroUtils.setIncludeAllFields(defaultSchema, true);
    AvroUtils.setProperty(defaultSchema, SchemaConstants.TALEND_COLUMN_PATTERN, DATE_PATTERN);
    return defaultSchema;
}
 

/**
 * add field to Schema.
 */
private FieldAssembler<Schema> addField(FieldAssembler<Schema> record, String name, Class<?> type, AvroRegistry avroReg) {
    Schema base = avroReg.getConverter(type).getSchema();
    FieldBuilder<Schema> fieldBuilder = record.name(name);
    fieldBuilder.type(AvroUtils.wrapAsNullable(base)).noDefault();
    return record;
}
 

/**
 * create schema for indexed record with Date presented as String.
 */
private Schema createSchemaWithTimeAsString() {
    AvroRegistry avroReg = new AvroRegistry();
    FieldAssembler<Schema> record = SchemaBuilder.record("Main").fields();
    addField(record, "FieldString", String.class, avroReg);
    addField(record, "FieldInt", Integer.class, avroReg);
    addField(record, "Description", String.class, avroReg);
    addField(record, "time", String.class, avroReg);
    Schema defaultSchema = record.endRecord();
    return defaultSchema;
}
 

/**
 * Create default static schema for indexed record.
 */
private Schema createSchema() {
    AvroRegistry avroReg = new AvroRegistry();
    FieldAssembler<Schema> record = SchemaBuilder.record("Main").fields();
    addField(record, "FieldString", String.class, avroReg);
    addField(record, "FieldInt", Integer.class, avroReg);
    Schema defaultSchema = record.endRecord();
    return defaultSchema;
}
 

private void setupDefaultSchema() {
    AvroRegistry avroReg = new AvroRegistry();
    FieldAssembler<Schema> record = SchemaBuilder.record("Main").fields();
    for (SplunkJSONEventField metadataField : SplunkJSONEventField.getMetadataFields()) {
        Schema base = avroReg.getConverter(metadataField.getDataType()).getSchema();
        FieldBuilder<Schema> fieldBuilder = record.name(metadataField.getName());
        if (metadataField.getName().equals(SplunkJSONEventField.TIME.getName())) {
            fieldBuilder.prop(SchemaConstants.TALEND_COLUMN_PATTERN, "dd-MM-yyyy");
        }
        fieldBuilder.type(AvroUtils.wrapAsNullable(base)).noDefault();
    }
    Schema defaultSchema = record.endRecord();
    schema.schema.setValue(defaultSchema);
}
 

public static Schema createTestSchema() {
    FieldAssembler<Schema> builder = SchemaBuilder.builder().record("TEST").fields();

    Schema schema = AvroUtils._int();
    schema = wrap(schema);
    builder = builder.name("ID").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "ID").type(schema).noDefault();

    schema = AvroUtils._string();
    schema = wrap(schema);
    builder = builder.name("NAME").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "NAME").type(schema).noDefault();

    return builder.endRecord();
}
 

private Schema inferSchemaRecord(IndexedRecord in) {
    FieldAssembler<Schema> builder = SchemaBuilder.builder().record(in.getSchema().getName()).fields();
    for (Schema.Field field : in.getSchema().getFields()) {
        Schema fieldSchema = inferSchema(field);
        Object fieldDefault = field.defaultVal();
        if (null == fieldDefault) {
            builder = builder.name(field.name()).type(fieldSchema).noDefault();
        } else {
            builder = builder.name(field.name()).type(fieldSchema).withDefault(fieldDefault);
        }
    }
    return builder.endRecord();
}
 

private static void addNullableField(
        FieldAssembler<Schema> builder,
        String columnName,
        Function<BaseTypeBuilder<UnionAccumulator<NullDefault<Schema>>>, UnionAccumulator<NullDefault<Schema>>> func
) {
    final BaseTypeBuilder<UnionAccumulator<NullDefault<Schema>>> and = builder.name(columnName).type().unionOf().nullBuilder().endNull().and();
    func.apply(and).endUnion().noDefault();
}
 

public static Schema createAllTypesSchema(String tablename) {
    FieldAssembler<Schema> builder = SchemaBuilder.builder().record(tablename).fields();

    // sql (smallint)short -> avro short
    Schema schema = AvroUtils._short();
    schema = wrap(schema);
    builder = builder.name("SMALL_INT_COL").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "SMALL_INT_COL").type(schema)
            .noDefault();

    schema = AvroUtils._int();
    schema = wrap(schema);
    builder = builder.name("INT_COL").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "INT_COL").type(schema).noDefault();

    // sql bigint -> avro long
    schema = AvroUtils._long();
    schema = wrap(schema);
    builder = builder.name("BIG_INT_COL").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "BIG_INT_COL").type(schema)
            .noDefault();

    schema = AvroUtils._float();
    schema = wrap(schema);
    builder = builder.name("REAL_COL").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "REAL_COL").type(schema)
            .noDefault();

    schema = AvroUtils._double();
    schema = wrap(schema);
    builder = builder.name("DOUBLE_COL").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "DOUBLE_COL").type(schema)
            .noDefault();

    // We don't use avro logical type, but use our own implementation of BigDecimal
    schema = AvroUtils._decimal();
    schema = wrap(schema);
    builder = builder.name("DECIMAL_COL").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "DECIMAL_COL").type(schema)
            .noDefault();

    // sql char -> avro string
    schema = AvroUtils._string();
    schema = wrap(schema);
    builder = builder.name("CHAR_COL").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "CHAR_COL").type(schema)
            .noDefault();

    // sql varchar -> avro string
    schema = AvroUtils._string();
    schema = wrap(schema);
    builder = builder.name("VARCHAR_COL").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "VARCHAR_COL").type(schema)
            .noDefault();

    // sql blob -> avro bytes
    schema = AvroUtils._bytes();
    schema = wrap(schema);
    builder = builder.name("BLOB_COL").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "BLOB_COL").type(schema)
            .noDefault();

    // sql clob -> avro string
    // TBD it could also be bytes
    schema = AvroUtils._string();
    schema = wrap(schema);
    builder = builder.name("CLOB_COL").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "CLOB_COL").type(schema)
            .noDefault();

    // sql date -> avro date
    schema = AvroUtils._date();
    schema = wrap(schema);
    builder = builder.name("DATE_COL").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "DATE_COL").type(schema)
            .noDefault();

    // sql time -> avro date
    schema = AvroUtils._date();
    schema = wrap(schema);
    builder = builder.name("TIME_COL").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "TIME_COL").type(schema)
            .noDefault();

    // sql timestamp -> avro date
    schema = AvroUtils._date();
    schema = wrap(schema);
    builder = builder.name("TIMESTAMP_COL").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "TIMESTAMP_COL").type(schema)
            .noDefault();

    schema = AvroUtils._boolean();
    schema = wrap(schema);
    builder = builder.name("BOOLEAN_COL").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "BOOLEAN_COL").type(schema)
            .noDefault();

    return builder.endRecord();
}
 

public static Schema createTestSchema3(boolean nullableForAnyColumn, String tablename) {
    FieldAssembler<Schema> builder = SchemaBuilder.builder().record(tablename).fields();

    Schema schema = AvroUtils._int();
    schema = wrap(schema, nullableForAnyColumn);
    builder = builder.name("C1").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "C1")
            .prop(SchemaConstants.TALEND_COLUMN_IS_KEY, "true").type(schema).noDefault();

    schema = AvroUtils._short();
    schema = wrap(schema, nullableForAnyColumn);
    builder = builder.name("C2").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "C2").type(schema).noDefault();

    schema = AvroUtils._long();
    schema = wrap(schema, nullableForAnyColumn);
    builder = builder.name("C3").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "C3").type(schema).noDefault();

    schema = AvroUtils._float();
    schema = wrap(schema, nullableForAnyColumn);
    builder = builder.name("C4").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "C4").type(schema).noDefault();

    schema = AvroUtils._double();
    schema = wrap(schema, nullableForAnyColumn);
    builder = builder.name("C5").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "C5").type(schema).noDefault();

    schema = AvroUtils._float();
    schema = wrap(schema, nullableForAnyColumn);
    builder = builder.name("C6").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "C6").type(schema).noDefault();

    schema = AvroUtils._decimal();
    schema = wrap(schema, nullableForAnyColumn);
    builder = builder.name("C7").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "C7").type(schema).noDefault();

    schema = AvroUtils._decimal();
    schema = wrap(schema, nullableForAnyColumn);
    builder = builder.name("C8").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "C8").type(schema).noDefault();

    schema = AvroUtils._boolean();
    schema = wrap(schema, nullableForAnyColumn);
    builder = builder.name("C9").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "C9").type(schema).noDefault();

    schema = AvroUtils._string();
    schema = wrap(schema, nullableForAnyColumn);
    builder = builder.name("C10").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "C10").type(schema).noDefault();

    schema = AvroUtils._date();
    schema = wrap(schema, nullableForAnyColumn);
    builder = builder.name("C11").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "C11").type(schema).noDefault();

    schema = AvroUtils._date();
    schema = wrap(schema, nullableForAnyColumn);
    builder = builder.name("C12").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "C12").type(schema).noDefault();

    schema = AvroUtils._date();
    schema = wrap(schema, nullableForAnyColumn);
    builder = builder.name("C13").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "C13").type(schema).noDefault();

    schema = AvroUtils._string();
    schema = wrap(schema, nullableForAnyColumn);
    builder = builder.name("C14").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "C14").type(schema).noDefault();

    schema = AvroUtils._string();
    schema = wrap(schema, nullableForAnyColumn);
    builder = builder.name("C15").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "C15").type(schema).noDefault();
    
    schema = AvroUtils._bytes();
    schema = wrap(schema, nullableForAnyColumn);
    builder = builder.name("C16").prop(SchemaConstants.TALEND_COLUMN_DB_COLUMN_NAME, "C16").type(schema).noDefault();

    return builder.endRecord();
}
 

private static FieldAssembler<Schema> getFieldBuilder()
{
    return SchemaBuilder.record("test_schema")
            .namespace("io.prestosql.decoder.avro")
            .fields();
}