Java Code Examples for org.apache.beam.sdk.schemas.Schema#getFields()

/**
 * Strict conversion from AVRO to Beam, strict because it doesn't do widening or narrowing during
 * conversion. If Schema is not provided, one is inferred from the AVRO schema.
 */
public static Row toBeamRowStrict(GenericRecord record, @Nullable Schema schema) {
  if (schema == null) {
    schema = toBeamSchema(record.getSchema());
  }

  Row.Builder builder = Row.withSchema(schema);
  org.apache.avro.Schema avroSchema = record.getSchema();

  for (Schema.Field field : schema.getFields()) {
    Object value = record.get(field.getName());
    org.apache.avro.Schema fieldAvroSchema = avroSchema.getField(field.getName()).schema();
    builder.addValue(convertAvroFieldStrict(value, fieldAvroSchema, field.getType()));
  }

  return builder.build();
}
 

private static void allLeafFields(
    Schema schema,
    List<String> nameComponents,
    SerializableFunction<List<String>, String> nameFn,
    Map<String, String> fieldsSelected) {
  for (Field field : schema.getFields()) {
    nameComponents.add(field.getName());
    if (field.getType().getTypeName().isCompositeType()) {
      allLeafFields(field.getType().getRowSchema(), nameComponents, nameFn, fieldsSelected);
    } else {
      String newName = nameFn.apply(nameComponents);
      fieldsSelected.put(String.join(".", nameComponents), newName);
    }
    nameComponents.remove(nameComponents.size() - 1);
  }
}
 

private List<Convert> createConverters(Schema schema) {
  List<Convert> fieldOverlays = new ArrayList<>();
  for (Schema.Field field : schema.getFields()) {
    fieldOverlays.add(createConverter(field));
  }
  return fieldOverlays;
}
 

/**
 * Converts all properties of an {@code Entity} to Beam {@code Row}.
 *
 * @param schema Target row {@code Schema}.
 * @param values A map of property names and values.
 * @return resulting Beam {@code Row}.
 */
private Row extractRowFromProperties(Schema schema, Map<String, Value> values) {
  Row.Builder builder = Row.withSchema(schema);
  // It is not a guarantee that the values will be in the same order as the schema.
  // Maybe metadata:
  // https://cloud.google.com/appengine/docs/standard/python/datastore/metadataqueries
  // TODO: figure out in what order the elements are in (without relying on Beam schema).
  for (Schema.Field field : schema.getFields()) {
    Value val = values.get(field.getName());
    builder.addValue(convertValueToObject(field.getType(), val));
  }
  return builder.build();
}
 

/**
 * Converts an entire {@code Row} to an appropriate DataStore {@code Entity.Builder}.
 *
 * @param row {@code Row} to convert.
 * @return resulting {@code Entity.Builder}.
 */
private Entity.Builder constructEntityFromRow(Schema schema, Row row) {
  Entity.Builder entityBuilder = Entity.newBuilder();
  for (Schema.Field field : schema.getFields()) {
    Value val = mapObjectToValue(row.getValue(field.getName()));
    entityBuilder.putProperties(field.getName(), val);
  }
  return entityBuilder;
}
 

private static Object[] rowToAvatica(Row row) {
  Schema schema = row.getSchema();
  Object[] convertedColumns = new Object[schema.getFields().size()];
  int i = 0;
  for (Schema.Field field : schema.getFields()) {
    convertedColumns[i] = fieldToAvatica(field.getType(), row.getBaseValue(i, Object.class));
    ++i;
  }
  return convertedColumns;
}
 

private static Schema uniquifyNames(Schema schema) {
  Schema.Builder builder = new Schema.Builder();
  for (Field field : schema.getFields()) {
    builder.addField(UUID.randomUUID().toString(), uniquifyNames(field.getType()));
  }
  return builder.build();
}
 

/** Converts a Beam Schema into an AVRO schema. */
public static org.apache.avro.Schema toAvroSchema(
    Schema beamSchema, @Nullable String name, @Nullable String namespace) {
  final String schemaName = Strings.isNullOrEmpty(name) ? "topLevelRecord" : name;
  final String schemaNamespace = namespace == null ? "" : namespace;
  String childNamespace =
      !"".equals(schemaNamespace) ? schemaNamespace + "." + schemaName : schemaName;
  List<org.apache.avro.Schema.Field> fields = Lists.newArrayList();
  for (Schema.Field field : beamSchema.getFields()) {
    org.apache.avro.Schema.Field recordField = toAvroField(field, childNamespace);
    fields.add(recordField);
  }
  return org.apache.avro.Schema.createRecord(schemaName, null, schemaNamespace, false, fields);
}
 

private Map<String, String> getMapping(Schema schema) {
  Map<String, String> mapping = Maps.newHashMap();
  for (Field field : schema.getFields()) {
    String underscore = CaseFormat.LOWER_UNDERSCORE.to(CaseFormat.LOWER_CAMEL, field.getName());
    mapping.put(underscore, field.getName());
    // The Avro compiler might add a $ at the end of a getter to disambiguate.
    mapping.put(underscore + "$", field.getName());
    // If the field is in camel case already, then it's the identity mapping.
    mapping.put(field.getName(), field.getName());
  }
  return mapping;
}
 

private static List<TableFieldSchema> toTableFieldSchema(Schema schema) {
  List<TableFieldSchema> fields = new ArrayList<>(schema.getFieldCount());
  for (Field schemaField : schema.getFields()) {
    FieldType type = schemaField.getType();

    TableFieldSchema field = new TableFieldSchema().setName(schemaField.getName());
    if (schemaField.getDescription() != null && !"".equals(schemaField.getDescription())) {
      field.setDescription(schemaField.getDescription());
    }

    if (!schemaField.getType().getNullable()) {
      field.setMode(Mode.REQUIRED.toString());
    }
    if (type.getTypeName().isCollectionType()) {
      type = type.getCollectionElementType();
      if (type.getTypeName().isCollectionType() || type.getTypeName().isMapType()) {
        throw new IllegalArgumentException("Array of collection is not supported in BigQuery.");
      }
      field.setMode(Mode.REPEATED.toString());
    }
    if (TypeName.ROW == type.getTypeName()) {
      Schema subType = type.getRowSchema();
      field.setFields(toTableFieldSchema(subType));
    }
    if (TypeName.MAP == type.getTypeName()) {
      throw new IllegalArgumentException("Maps are not supported in BigQuery.");
    }
    field.setType(toStandardSQLTypeName(type).toString());

    fields.add(field);
  }
  return fields;
}