Examples with ReadRowsResponseToInternalRowIteratorConverter com.google.cloud.spark.bigquery.ReadRowsResponseToInternalRowIteratorConverter used on opensource projects

Example 1 with ReadRowsResponseToInternalRowIteratorConverter

use of com.google.cloud.spark.bigquery.ReadRowsResponseToInternalRowIteratorConverter in project spark-bigquery-connector by GoogleCloudDataproc.

the class BigQueryDataSourceReaderContext method createConverter.

private ReadRowsResponseToInternalRowIteratorConverter createConverter(ImmutableList<String> selectedFields, ReadSessionResponse readSessionResponse, Optional<StructType> userProvidedSchema) {
    ReadRowsResponseToInternalRowIteratorConverter converter;
    DataFormat format = readSessionCreatorConfig.getReadDataFormat();
    if (format == DataFormat.AVRO) {
        Schema schema = SchemaConverters.getSchemaWithPseudoColumns(readSessionResponse.getReadTableInfo());
        if (selectedFields.isEmpty()) {
            // means select *
            selectedFields = schema.getFields().stream().map(Field::getName).collect(ImmutableList.toImmutableList());
        } else {
            Set<String> requiredColumnSet = ImmutableSet.copyOf(selectedFields);
            schema = Schema.of(schema.getFields().stream().filter(field -> requiredColumnSet.contains(field.getName())).collect(Collectors.toList()));
        }
        return ReadRowsResponseToInternalRowIteratorConverter.avro(schema, selectedFields, readSessionResponse.getReadSession().getAvroSchema().getSchema(), userProvidedSchema);
    }
    throw new IllegalArgumentException("No known converted for " + readSessionCreatorConfig.getReadDataFormat());
}

Example 2 with ReadRowsResponseToInternalRowIteratorConverter

use of com.google.cloud.spark.bigquery.ReadRowsResponseToInternalRowIteratorConverter in project spark-bigquery-connector by GoogleCloudDataproc.

the class BigQueryInputPartitionReaderContextTest method testReadAvro.

@Test
public void testReadAvro() throws Exception {
    TableInfo allTypesTableInfo = allTypesTableInfo();
    ReadRowsResponse.Builder readRowsResponse = ReadRowsResponse.newBuilder();
    TextFormat.merge(ALL_TYPES_TABLE_READ_ROWS_RESPONSE_STR, readRowsResponse);
    Iterator<ReadRowsResponse> readRowsResponses = ImmutableList.of(readRowsResponse.build()).iterator();
    ReadRowsResponseToInternalRowIteratorConverter converter = ReadRowsResponseToInternalRowIteratorConverter.avro(ALL_TYPES_TABLE_BIGQUERY_SCHEMA, ALL_TYPES_TABLE_FIELDS, ALL_TYPES_TABLE_AVRO_RAW_SCHEMA, Optional.empty());
    BigQueryInputPartitionReaderContext reader = new BigQueryInputPartitionReaderContext(readRowsResponses, converter, null);
    assertThat(reader.next()).isTrue();
    InternalRow row = reader.get();
    assertThat(reader.next()).isFalse();
    assertThat(row.numFields()).isEqualTo(15);
    assertThat(row.getString(0)).isEqualTo("hello");
}