Java Code Examples for org.apache.beam.sdk.values.Row#withSchema()

public static Row castRow(Row input, Schema inputSchema, Schema outputSchema) {
  if (input == null) {
    return null;
  }

  Row.Builder output = Row.withSchema(outputSchema);
  for (int i = 0; i < outputSchema.getFieldCount(); i++) {
    Schema.Field outputField = outputSchema.getField(i);

    int fromFieldIdx = inputSchema.indexOf(outputField.getName());
    Schema.Field inputField = inputSchema.getField(fromFieldIdx);

    Object inputValue = input.getValue(fromFieldIdx);
    Object outputValue = castValue(inputValue, inputField.getType(), outputField.getType());

    output.addValue(outputValue);
  }

  return output.build();
}
 

/**
 * 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();
}
 

@Override
public Row apply(T input) {
  Schema schema = context.getSchema();
  Row.Builder builder = Row.withSchema(schema);
  for (Convert convert : converters) {
    builder.addValue(convert.getFromProtoMessage((Message) input));
  }
  return builder.build();
}
 

/**
 * 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();
}
 

@ProcessElement
public void processElement(ProcessContext c) {
  Row row = c.element();
  Row.Builder builder = Row.withSchema(keySchema);
  for (Integer i : keyIndex) {
    builder.addValue(row.getValue(i));
  }
  Row keyRow = builder.build();
  c.output(KV.of(keyRow, row));
}
 

@ProcessElement
public void processElement(ProcessContext c) {
  Row row = c.element();
  IntervalWindow window = windowFn.assignWindow(row.getDateTime(windowFieldIndex).toInstant());
  Row.Builder builder = Row.withSchema(outputSchema);
  builder.addValues(row.getValues());
  builder.addValue(window.start());
  builder.addValue(window.end());
  c.output(builder.build());
}
 

@ProcessElement
public void processElement(ProcessContext c) {
  Row row = c.element();
  Collection<IntervalWindow> windows =
      windowFn.assignWindows(row.getDateTime(windowFieldIndex).toInstant());
  for (IntervalWindow window : windows) {
    Row.Builder builder = Row.withSchema(outputSchema);
    builder.addValues(row.getValues());
    builder.addValue(window.start());
    builder.addValue(window.end());
    c.output(builder.build());
  }
}
 

@Override
public Row extractOutput(Object[] accumulator) {
  // Build a row containing a field for every aggregate that was registered.
  CoCombineResult coCombineResult = getComposedCombineFn().extractOutput(accumulator);
  Row.Builder output = Row.withSchema(getOutputSchema());
  for (FieldAggregation fieldAggregation : getFieldAggregations()) {
    Object aggregate = coCombineResult.get(fieldAggregation.combineTag);
    output.addValue(aggregate);
  }
  return output.build();
}
 

/** Select a sub Row from an input Row. */
private static Row selectRow(
    Row input,
    FieldAccessDescriptor fieldAccessDescriptor,
    Schema inputSchema,
    Schema outputSchema) {
  if (fieldAccessDescriptor.getAllFields()) {
    return input;
  }

  Row.Builder output = Row.withSchema(outputSchema);
  selectIntoRow(inputSchema, input, output, fieldAccessDescriptor);
  return output.build();
}
 

@Override
public Row toBaseType(Value input) {
  EnumerationType.Value caseType = input.getCaseType();
  int setFieldIndex = oneOfSchema.indexOf(enumerationType.toString(caseType));
  Row.Builder builder = Row.withSchema(oneOfSchema);
  for (int i = 0; i < oneOfSchema.getFieldCount(); ++i) {
    Object value = (i == setFieldIndex) ? input.getValue() : null;
    builder = builder.addValue(value);
  }
  return builder.build();
}
 

public static Object rowFromProto(SchemaApi.Row row, FieldType fieldType) {
  Row.Builder builder = Row.withSchema(fieldType.getRowSchema());
  for (int i = 0; i < row.getValuesCount(); ++i) {
    builder.addValue(
        fieldValueFromProto(fieldType.getRowSchema().getField(i).getType(), row.getValues(i)));
  }
  return builder.build();
}
 

@Override
public Row mapRow(ResultSet rs) throws Exception {
  Row.Builder rowBuilder = Row.withSchema(schema);
  for (int i = 0; i < schema.getFieldCount(); i++) {
    rowBuilder.addValue(fieldExtractors.get(i).extract(rs, i + 1));
  }
  return rowBuilder.build();
}
 

private static Row kafkaSourceRecordToBeamRow(Struct value, Row.Builder rowBuilder) {
  org.apache.beam.sdk.schemas.Schema beamSchema = rowBuilder.getSchema();
  for (org.apache.beam.sdk.schemas.Schema.Field f : beamSchema.getFields()) {
    switch (f.getType().getTypeName()) {
      case INT16:
        rowBuilder.addValue(value.getInt16(f.getName()));
        break;
      case INT32:
        rowBuilder.addValue(value.getInt32(f.getName()));
        break;
      case INT64:
        rowBuilder.addValue(value.getInt64(f.getName()));
        break;
      case FLOAT:
        rowBuilder.addValue(value.getFloat32(f.getName()));
        break;
      case DOUBLE:
        rowBuilder.addValue(value.getFloat64(f.getName()));
        break;
      case BOOLEAN:
        rowBuilder.addValue(value.getBoolean(f.getName()));
        break;
      case STRING:
        rowBuilder.addValue(value.getString(f.getName()));
        break;
      case BYTES:
        rowBuilder.addValue(value.getBytes(f.getName()));
        break;
      case ROW:
        Row.Builder nestedRowBuilder = Row.withSchema(f.getType().getRowSchema());
        rowBuilder.addValue(
            kafkaSourceRecordToBeamRow(value.getStruct(f.getName()), nestedRowBuilder));
        break;
      case MAP:
        throw new DataException("Map types are not supported.");
      case ARRAY:
        throw new DataException("Array types are not supported.");
      default:
        throw new DataException(
            String.format("Unsupported data type: {}", f.getType().getTypeName()));
    }
  }
  return rowBuilder.build();
}
 

private static Row handleValue(org.apache.kafka.connect.data.Schema schema, Object value) {
  org.apache.beam.sdk.schemas.Schema beamSchema =  kafkaSchemaToBeamRowSchema(schema);
  Row.Builder rowBuilder = Row.withSchema(beamSchema);
  return kafkaSourceRecordToBeamRow((Struct) value, rowBuilder);
}