@Test public void testNestedFieldByName() { FieldAccessDescriptor fieldAccessDescriptor = FieldAccessDescriptor.withFieldNames("field1") .withNestedField("field1", FieldAccessDescriptor.withAllFields()); fieldAccessDescriptor = fieldAccessDescriptor.resolve(NESTED_SCHEMA2); assertTrue(fieldAccessDescriptor.fieldIdsAccessed().isEmpty()); assertEquals(1, fieldAccessDescriptor.nestedFields().size()); FieldAccessDescriptor nestedAccess = fieldAccessDescriptor.nestedFields().get(1); assertTrue(nestedAccess.allFields()); }
@Test public void testArrayNestedField() { FieldAccessDescriptor level1 = FieldAccessDescriptor.withFieldNames("field2"); FieldAccessDescriptor level2 = FieldAccessDescriptor.withFieldNames("field1").withNestedField("field1", level1); FieldAccessDescriptor resolved = level2.resolve(NESTED_ARRAY_SCHEMA); assertTrue(resolved.fieldIdsAccessed().isEmpty()); assertEquals(1, resolved.nestedFields().size()); resolved = resolved.nestedFields().get(1); assertEquals(Sets.newHashSet(2), resolved.fieldIdsAccessed()); }
@Test public void testPartialAccessNestedField() { FieldAccessDescriptor level1 = FieldAccessDescriptor.withFieldNames("field2"); FieldAccessDescriptor level2 = FieldAccessDescriptor.withFieldNames("field1").withNestedField("field1", level1); FieldAccessDescriptor level3 = FieldAccessDescriptor.withFieldNames("field1").withNestedField("field1", level2); FieldAccessDescriptor resolved = level3.resolve(NESTED_SCHEMA2); assertTrue(resolved.fieldIdsAccessed().isEmpty()); assertEquals(1, resolved.nestedFields().size()); resolved = resolved.nestedFields().get(1); assertTrue(resolved.fieldIdsAccessed().isEmpty()); assertEquals(1, resolved.nestedFields().size()); resolved = resolved.nestedFields().get(1); assertEquals(Sets.newHashSet(2), resolved.fieldIdsAccessed()); }
@Test public void testNestedFieldById() { FieldAccessDescriptor fieldAccessDescriptor = FieldAccessDescriptor.withFieldNames("field1") .withNestedField(1, FieldAccessDescriptor.withAllFields()); fieldAccessDescriptor = fieldAccessDescriptor.resolve(NESTED_SCHEMA2); assertTrue(fieldAccessDescriptor.fieldIdsAccessed().isEmpty()); assertEquals(1, fieldAccessDescriptor.nestedFields().size()); FieldAccessDescriptor nestedAccess = fieldAccessDescriptor.nestedFields().get(1); assertTrue(nestedAccess.allFields()); }
@Test public void testMapNestedField() { FieldAccessDescriptor level1 = FieldAccessDescriptor.withFieldNames("field2"); FieldAccessDescriptor level2 = FieldAccessDescriptor.withFieldNames("field1").withNestedField("field1", level1); FieldAccessDescriptor resolved = level2.resolve(NESTED_MAP_SCHEMA); assertTrue(resolved.fieldIdsAccessed().isEmpty()); assertEquals(1, resolved.nestedFields().size()); resolved = resolved.nestedFields().get(1); assertEquals(Sets.newHashSet(2), resolved.fieldIdsAccessed()); } }
static Schema getOutputSchema(Schema inputSchema, FieldAccessDescriptor fieldAccessDescriptor) { if (fieldAccessDescriptor.allFields()) { return inputSchema; } Schema.Builder builder = new Schema.Builder(); for (int fieldId : fieldAccessDescriptor.fieldIdsAccessed()) { builder.addField(inputSchema.getField(fieldId)); } for (Map.Entry<Integer, FieldAccessDescriptor> nested : fieldAccessDescriptor.nestedFields().entrySet()) { Field field = inputSchema.getField(nested.getKey()); FieldAccessDescriptor nestedDescriptor = nested.getValue(); FieldType nestedType = FieldType.row(getOutputSchema(field.getType().getRowSchema(), nestedDescriptor)); if (field.getNullable()) { builder.addNullableField(field.getName(), nestedType); } else { builder.addField(field.getName(), nestedType); } } return builder.build(); }
static Row selectRow( Row input, FieldAccessDescriptor fieldAccessDescriptor, Schema inputSchema, Schema outputSchema) { if (fieldAccessDescriptor.allFields()) { return input; } else { Row.Builder output = Row.withSchema(outputSchema); for (int fieldId : fieldAccessDescriptor.fieldIdsAccessed()) { output.addValue(input.getValue(fieldId)); } for (Map.Entry<Integer, FieldAccessDescriptor> nested : fieldAccessDescriptor.nestedFields().entrySet()) { String fieldName = inputSchema.nameOf(nested.getKey()); Schema nestedInputSchema = inputSchema.getField(nested.getKey()).getType().getRowSchema(); Schema nestedOutputSchema = outputSchema.getField(fieldName).getType().getRowSchema(); output.addValue( selectRow( input.getValue(fieldName), nested.getValue(), nestedInputSchema, nestedOutputSchema)); } return output.build(); } } }