Examples with Row org.apache.beam.sdk.values.Row used on opensource projects

Example 21 with Row

use of org.apache.beam.sdk.values.Row in project beam by apache.

the class ZetaSQLPushDownTest method initializeBeamTableProvider.

private static void initializeBeamTableProvider() {
    Table projectTable = getTable("InMemoryTableProject", PushDownOptions.PROJECT);
    Table bothTable = getTable("InMemoryTableBoth", PushDownOptions.BOTH);
    Row[] rows = { row(BASIC_SCHEMA, 100L, 1L, "one", 100L), row(BASIC_SCHEMA, 200L, 2L, "two", 200L) };
    tableProvider = new TestTableProvider();
    tableProvider.createTable(projectTable);
    tableProvider.createTable(bothTable);
    tableProvider.addRows(projectTable.getName(), rows);
    tableProvider.addRows(bothTable.getName(), rows);
}

Example 22 with Row

use of org.apache.beam.sdk.values.Row in project beam by apache.

the class TextTableProviderTest method testInvalidJson.

@Test
public void testInvalidJson() throws Exception {
    File deadLetterFile = new File(tempFolder.getRoot(), "dead-letter-file");
    Files.write(tempFolder.newFile("test.json").toPath(), INVALID_JSON_TEXT.getBytes(Charsets.UTF_8));
    String query = "SELECT * FROM test";
    String ddl = String.format("CREATE EXTERNAL TABLE test %s TYPE text LOCATION '%s/*' " + "TBLPROPERTIES '{\"format\":\"json\", \"deadLetterFile\": \"%s\"}'", SQL_JSON_SCHEMA, tempFolder.getRoot(), deadLetterFile.getAbsoluteFile());
    PCollection<Row> rows = pipeline.apply(SqlTransform.query(query).withDdlString(ddl));
    PAssert.that(rows).empty();
    pipeline.run();
    assertThat(new NumberedShardedFile(deadLetterFile.getAbsoluteFile() + "*").readFilesWithRetries(Sleeper.DEFAULT, BackOff.STOP_BACKOFF), containsInAnyOrder(INVALID_JSON_TEXT));
}

Example 23 with Row

use of org.apache.beam.sdk.values.Row in project beam by apache.

the class BeamRowToStorageApiProto method toProtoValue.

private static Object toProtoValue(FieldDescriptor fieldDescriptor, FieldType beamFieldType, Object value) {
    switch(beamFieldType.getTypeName()) {
        case ROW:
            return messageFromBeamRow(fieldDescriptor.getMessageType(), (Row) value);
        case ARRAY:
            List<Object> list = (List<Object>) value;
            @Nullable FieldType arrayElementType = beamFieldType.getCollectionElementType();
            if (arrayElementType == null) {
                throw new RuntimeException("Unexpected null element type!");
            }
            return list.stream().map(v -> toProtoValue(fieldDescriptor, arrayElementType, v)).collect(Collectors.toList());
        case ITERABLE:
            Iterable<Object> iterable = (Iterable<Object>) value;
            @Nullable FieldType iterableElementType = beamFieldType.getCollectionElementType();
            if (iterableElementType == null) {
                throw new RuntimeException("Unexpected null element type!");
            }
            return StreamSupport.stream(iterable.spliterator(), false).map(v -> toProtoValue(fieldDescriptor, iterableElementType, v)).collect(Collectors.toList());
        case MAP:
            throw new RuntimeException("Map types not supported by BigQuery.");
        default:
            return scalarToProtoValue(beamFieldType, value);
    }
}

Example 24 with Row

use of org.apache.beam.sdk.values.Row in project beam by apache.

the class JdbcIOTest method testReadRowsWithNumericFieldsWithExcessPrecision.

@Test
public void testReadRowsWithNumericFieldsWithExcessPrecision() {
    PCollection<Row> rows = pipeline.apply(JdbcIO.readRows().withDataSourceConfiguration(DATA_SOURCE_CONFIGURATION).withQuery(String.format("SELECT CAST(1 AS NUMERIC(10, 2)) AS T1 FROM %s WHERE name = ?", READ_TABLE_NAME)).withStatementPreparator(preparedStatement -> preparedStatement.setString(1, TestRow.getNameForSeed(1))));
    Schema expectedSchema = Schema.of(Schema.Field.of("T1", FieldType.logicalType(FixedPrecisionNumeric.of(NUMERIC.getName(), 10, 2)).withNullable(false)));
    assertEquals(expectedSchema, rows.getSchema());
    PCollection<Row> output = rows.apply(Select.fieldNames("T1"));
    PAssert.that(output).containsInAnyOrder(ImmutableList.of(Row.withSchema(expectedSchema).addValues(BigDecimal.valueOf(1).setScale(2, RoundingMode.HALF_UP)).build()));
    pipeline.run();
}

Example 25 with Row

use of org.apache.beam.sdk.values.Row in project beam by apache.

the class JdbcIOTest method testGetPreparedStatementSetCallerForLogicalTypes.

@Test
public void testGetPreparedStatementSetCallerForLogicalTypes() throws Exception {
    FieldType fixedLengthStringType = LogicalTypes.fixedLengthString(JDBCType.VARCHAR, 4);
    Schema schema = Schema.builder().addField("logical_date_col", LogicalTypes.JDBC_DATE_TYPE).addField("logical_time_col", LogicalTypes.JDBC_TIME_TYPE).addField("logical_time_with_tz_col", LogicalTypes.JDBC_TIMESTAMP_WITH_TIMEZONE_TYPE).addField("logical_fixed_length_string_col", fixedLengthStringType).addField("logical_fixed_length_string_nullable_col", fixedLengthStringType.withNullable(true)).addField("logical_uuid_col", LogicalTypes.JDBC_UUID_TYPE).addField("logical_other_col", LogicalTypes.OTHER_AS_STRING_TYPE).build();
    long epochMilli = 1558719710000L;
    DateTime dateTime = new DateTime(epochMilli, ISOChronology.getInstanceUTC());
    DateTime time = new DateTime(34567000L, /* value must be less than num millis in one day */
    ISOChronology.getInstanceUTC());
    Row row = Row.withSchema(schema).addValues(dateTime.withTimeAtStartOfDay(), time, dateTime, "Test", null, UUID.randomUUID(), "{}").build();
    PreparedStatement psMocked = mock(PreparedStatement.class);
    JdbcUtil.getPreparedStatementSetCaller(LogicalTypes.JDBC_DATE_TYPE).set(row, psMocked, 0, SchemaUtil.FieldWithIndex.of(schema.getField(0), 0));
    JdbcUtil.getPreparedStatementSetCaller(LogicalTypes.JDBC_TIME_TYPE).set(row, psMocked, 1, SchemaUtil.FieldWithIndex.of(schema.getField(1), 1));
    JdbcUtil.getPreparedStatementSetCaller(LogicalTypes.JDBC_TIMESTAMP_WITH_TIMEZONE_TYPE).set(row, psMocked, 2, SchemaUtil.FieldWithIndex.of(schema.getField(2), 2));
    JdbcUtil.getPreparedStatementSetCaller(fixedLengthStringType).set(row, psMocked, 3, SchemaUtil.FieldWithIndex.of(schema.getField(3), 3));
    JdbcUtil.getPreparedStatementSetCaller(fixedLengthStringType.withNullable(true)).set(row, psMocked, 4, SchemaUtil.FieldWithIndex.of(schema.getField(4), 4));
    JdbcUtil.getPreparedStatementSetCaller(LogicalTypes.JDBC_UUID_TYPE).set(row, psMocked, 5, SchemaUtil.FieldWithIndex.of(schema.getField(5), 5));
    JdbcUtil.getPreparedStatementSetCaller(LogicalTypes.OTHER_AS_STRING_TYPE).set(row, psMocked, 6, SchemaUtil.FieldWithIndex.of(schema.getField(6), 6));
    verify(psMocked, times(1)).setDate(1, new Date(row.getDateTime(0).getMillis()));
    verify(psMocked, times(1)).setTime(2, new Time(row.getDateTime(1).getMillis()));
    Calendar cal = Calendar.getInstance(TimeZone.getTimeZone("UTC"));
    cal.setTimeInMillis(epochMilli);
    verify(psMocked, times(1)).setTimestamp(3, new Timestamp(cal.getTime().getTime()), cal);
    verify(psMocked, times(1)).setString(4, row.getString(3));
    verify(psMocked, times(1)).setString(5, row.getString(4));
    verify(psMocked, times(1)).setObject(6, row.getLogicalTypeValue(5, UUID.class));
    verify(psMocked, times(1)).setObject(7, row.getString(6), java.sql.Types.OTHER);
}