Examples with FieldAccessDescriptor org.apache.beam.sdk.schemas.FieldAccessDescriptor used on opensource projects

Example 41 with FieldAccessDescriptor

use of org.apache.beam.sdk.schemas.FieldAccessDescriptor in project beam by apache.

the class RenameFields method renameSchema.

// Apply the user-specified renames to the input schema.
@VisibleForTesting
static void renameSchema(Schema inputSchema, Collection<RenamePair> renames, Map<UUID, Schema> renamedSchemasMap, Map<UUID, BitSet> nestedFieldRenamedMap) {
    // The mapping of renames to apply at this level of the schema.
    Map<Integer, String> topLevelRenames = Maps.newHashMap();
    // For nested schemas, collect all applicable renames here.
    Multimap<Integer, RenamePair> nestedRenames = ArrayListMultimap.create();
    for (RenamePair rename : renames) {
        FieldAccessDescriptor access = rename.getFieldAccessDescriptor();
        if (!access.fieldIdsAccessed().isEmpty()) {
            // This references a field at this level of the schema.
            Integer fieldId = Iterables.getOnlyElement(access.fieldIdsAccessed());
            topLevelRenames.put(fieldId, rename.getNewName());
        } else {
            // This references a nested field.
            Map.Entry<Integer, FieldAccessDescriptor> nestedAccess = Iterables.getOnlyElement(access.nestedFieldsById().entrySet());
            nestedFieldRenamedMap.computeIfAbsent(inputSchema.getUUID(), s -> new BitSet(inputSchema.getFieldCount())).set(nestedAccess.getKey());
            nestedRenames.put(nestedAccess.getKey(), RenamePair.of(nestedAccess.getValue(), rename.getNewName()));
        }
    }
    Schema.Builder builder = Schema.builder();
    for (int i = 0; i < inputSchema.getFieldCount(); ++i) {
        Field field = inputSchema.getField(i);
        FieldType fieldType = field.getType();
        String newName = topLevelRenames.getOrDefault(i, field.getName());
        Collection<RenamePair> nestedFieldRenames = nestedRenames.asMap().getOrDefault(i, Collections.emptyList());
        builder.addField(newName, renameFieldType(fieldType, nestedFieldRenames, renamedSchemasMap, nestedFieldRenamedMap));
    }
    renamedSchemasMap.put(inputSchema.getUUID(), builder.build());
}

Example 42 with FieldAccessDescriptor

use of org.apache.beam.sdk.schemas.FieldAccessDescriptor in project beam by apache.

the class SelectHelpersTest method testSelectIterableOfRowPartial.

@Test
public void testSelectIterableOfRowPartial() {
    FieldAccessDescriptor fieldAccessDescriptor = FieldAccessDescriptor.withFieldNames("rowIter[].field1").resolve(ITERABLE_SCHEMA);
    Schema outputSchema = SelectHelpers.getOutputSchema(ITERABLE_SCHEMA, fieldAccessDescriptor);
    Schema expectedSchema = Schema.builder().addIterableField("field1", FieldType.STRING).build();
    assertEquals(expectedSchema, outputSchema);
    Row row = selectRow(ITERABLE_SCHEMA, fieldAccessDescriptor, ITERABLE_ROW);
    Row expectedRow = Row.withSchema(expectedSchema).addIterable(ImmutableList.of("first", "first")).build();
    assertEquals(expectedRow, row);
}

Example 43 with FieldAccessDescriptor

use of org.apache.beam.sdk.schemas.FieldAccessDescriptor in project beam by apache.

the class SelectHelpersTest method testSelectArrayOfRowArray.

@Test
public void testSelectArrayOfRowArray() {
    FieldAccessDescriptor fieldAccessDescriptor = FieldAccessDescriptor.withFieldNames("arrayOfRowArray[][].field1").resolve(ARRAY_SCHEMA);
    Schema outputSchema = SelectHelpers.getOutputSchema(ARRAY_SCHEMA, fieldAccessDescriptor);
    Schema expectedSchema = Schema.builder().addArrayField("field1", FieldType.array(FieldType.STRING)).build();
    assertEquals(expectedSchema, outputSchema);
    Row row = selectRow(ARRAY_SCHEMA, fieldAccessDescriptor, ARRAY_ROW);
    Row expectedRow = Row.withSchema(expectedSchema).addArray(ImmutableList.of("first"), ImmutableList.of("first")).build();
    assertEquals(expectedRow, row);
}

Example 44 with FieldAccessDescriptor

use of org.apache.beam.sdk.schemas.FieldAccessDescriptor in project beam by apache.

the class SelectHelpersTest method testSelectNullableNestedRowArray.

@Test
public void testSelectNullableNestedRowArray() {
    FieldAccessDescriptor fieldAccessDescriptor1 = FieldAccessDescriptor.withFieldNames("nestedArray.field1").resolve(NESTED_NULLABLE_SCHEMA);
    Row out1 = selectRow(NESTED_NULLABLE_SCHEMA, fieldAccessDescriptor1, Row.nullRow(NESTED_NULLABLE_SCHEMA));
    assertNull(out1.getValue(0));
    FieldAccessDescriptor fieldAccessDescriptor2 = FieldAccessDescriptor.withFieldNames("nestedArray.*").resolve(NESTED_NULLABLE_SCHEMA);
    Row out2 = selectRow(NESTED_NULLABLE_SCHEMA, fieldAccessDescriptor2, Row.nullRow(NESTED_NULLABLE_SCHEMA));
    assertEquals(Collections.nCopies(4, null), out2.getValues());
}

Example 45 with FieldAccessDescriptor

use of org.apache.beam.sdk.schemas.FieldAccessDescriptor in project beam by apache.

the class SelectHelpersTest method testSelectArrayOfRow.

@Test
public void testSelectArrayOfRow() {
    FieldAccessDescriptor fieldAccessDescriptor = FieldAccessDescriptor.withFieldNames("rowArray").resolve(ARRAY_SCHEMA);
    Schema outputSchema = SelectHelpers.getOutputSchema(ARRAY_SCHEMA, fieldAccessDescriptor);
    Schema expectedSchema = Schema.builder().addArrayField("rowArray", FieldType.row(FLAT_SCHEMA)).build();
    assertEquals(expectedSchema, outputSchema);
    Row row = selectRow(ARRAY_SCHEMA, fieldAccessDescriptor, ARRAY_ROW);
    Row expectedRow = Row.withSchema(expectedSchema).addArray(FLAT_ROW, FLAT_ROW).build();
    assertEquals(expectedRow, row);
}