@Test public void testCreateMapWithNullValue() { Map<Integer, String> data = new HashMap(); data.put(1, "value1"); data.put(2, "value2"); data.put(3, null); data.put(4, null); Schema type = Stream.of(Schema.Field.of("map", FieldType.map(FieldType.INT32, FieldType.STRING, true))) .collect(toSchema()); Row row = Row.withSchema(type).addValue(data).build(); assertEquals(data, row.getMap("map")); }
@Test public void testCreateMapWithArrayValue() { Map<Integer, List<String>> data = ImmutableMap.<Integer, List<String>>builder() .put(1, Arrays.asList("value1")) .put(2, Arrays.asList("value2")) .build(); Schema type = Stream.of( Schema.Field.of( "map", FieldType.map(FieldType.INT32, FieldType.array(FieldType.STRING)))) .collect(toSchema()); Row row = Row.withSchema(type).addValue(data).build(); assertEquals(data, row.getMap("map")); }
@Test public void testByteBufferEquality() { byte[] a0 = new byte[] {1, 2, 3, 4}; byte[] b0 = new byte[] {1, 2, 3, 4}; Schema schema = Schema.of(Schema.Field.of("bytes", Schema.FieldType.BYTES)); Row a = Row.withSchema(schema).addValue(ByteBuffer.wrap(a0)).build(); Row b = Row.withSchema(schema).addValue(ByteBuffer.wrap(b0)).build(); Assert.assertEquals(a, b); } }
@Test public void testCreateMapWithPrimitiveValue() { Map<Integer, String> data = ImmutableMap.<Integer, String>builder() .put(1, "value1") .put(2, "value2") .put(3, "value3") .put(4, "value4") .build(); Schema type = Stream.of(Schema.Field.of("map", FieldType.map(FieldType.INT32, FieldType.STRING))) .collect(toSchema()); Row row = Row.withSchema(type).addValue(data).build(); assertEquals(data, row.getMap("map")); }
@Test public void testAccessMapElement() { PCollection<Row> input = pCollectionOf2Elements(); Schema resultType = Schema.builder().addNullableField("f_mapElem", Schema.FieldType.INT32).build(); PCollection<Row> result = input.apply( "sqlQuery", SqlTransform.query("SELECT f_intStringMap['key11'] FROM PCOLLECTION")); PAssert.that(result) .containsInAnyOrder( Row.withSchema(resultType).addValues(11).build(), Row.withSchema(resultType).addValue(null).build()); pipeline.run(); }
@Test public void testToHex() throws Exception { Schema resultType = Schema.builder().addStringField("field").build(); Row resultRow = Row.withSchema(resultType).addValue("666f6f626172").build(); Row resultRow2 = Row.withSchema(resultType).addValue("20").build(); Row resultRow3 = Row.withSchema(resultType).addValue("616263414243").build(); Row resultRow4 = Row.withSchema(resultType).addValue("616263414243d0b6d189d184d096d0a9d0a4").build(); String sql = "SELECT TO_HEX(f_bytes) FROM PCOLLECTION WHERE f_func = 'TO_HEX'"; PCollection<Row> result = boundedInputBytes.apply("testUdf", SqlTransform.query(sql)); PAssert.that(result).containsInAnyOrder(resultRow, resultRow2, resultRow3, resultRow4); pipeline.run().waitUntilFinish(); }
@Test public void testRowWithArray() { BeamSqlEnv sqlEnv = BeamSqlEnv.inMemory(readOnlyTableProvider); PCollection<Row> stream = BeamSqlRelUtils.toPCollection( pipeline, sqlEnv.parseQuery( "SELECT rowWithArrayTestTable.col.field3[2] FROM rowWithArrayTestTable")); PAssert.that(stream) .containsInAnyOrder( Row.withSchema(Schema.builder().addInt64Field("int64").build()).addValue(6L).build()); pipeline.run().waitUntilFinish(Duration.standardMinutes(2)); }
@Test @Ignore public void testConsistentWithEqualsMapWithBytesKeyField() throws Exception { FieldType fieldType = FieldType.map(FieldType.BYTES, FieldType.INT32); Schema schema = Schema.of(Schema.Field.of("f1", fieldType)); RowCoder coder = RowCoder.of(schema); Map<byte[], Integer> map1 = Collections.singletonMap(new byte[] {1, 2, 3, 4}, 1); Row row1 = Row.withSchema(schema).addValue(map1).build(); Map<byte[], Integer> map2 = Collections.singletonMap(new byte[] {1, 2, 3, 4}, 1); Row row2 = Row.withSchema(schema).addValue(map2).build(); Assume.assumeTrue(coder.consistentWithEquals()); CoderProperties.coderConsistentWithEquals(coder, row1, row2); }
@Override public PCollection<Row> expand(PCollection<String> input) { return input .apply( "linesToRows", MapElements.into(TypeDescriptors.rows()) .via(s -> Row.withSchema(SCHEMA).addValue(s).build())) .setRowSchema(SCHEMA); } }
@Override public <T> SerializableFunction<T, Row> toRowFunction(TypeDescriptor<T> typeDescriptor) { return v -> Row.withSchema(schemaFor(typeDescriptor)).addValue(v).build(); }
@Test(expected = IllegalArgumentException.class) public void testIgnoreNullFail() { Schema inputSchema = Schema.of(Schema.Field.nullable("f0", Schema.FieldType.INT32)); Schema outputSchema = Schema.of(Schema.Field.of("f0", Schema.FieldType.INT32)); Cast.castRow(Row.withSchema(inputSchema).addValue(null).build(), inputSchema, outputSchema); }
@Test public void testNNestedullValuesSetters() throws NoSuchSchemaException { SchemaRegistry registry = SchemaRegistry.createDefault(); Row row = Row.withSchema(NESTED_NULLABLE_SCHEMA).addValue(null).build(); POJOWithNestedNullable pojo = registry.getFromRowFunction(POJOWithNestedNullable.class).apply(row); assertNull(pojo.nested); } }
@Test public void testConsistentWithEqualsBytesField() throws Exception { Schema schema = Schema.of(Schema.Field.of("f1", FieldType.BYTES)); Row row1 = Row.withSchema(schema).addValue(new byte[] {1, 2, 3, 4}).build(); Row row2 = Row.withSchema(schema).addValue(new byte[] {1, 2, 3, 4}).build(); RowCoder coder = RowCoder.of(schema); Assume.assumeTrue(coder.consistentWithEquals()); CoderProperties.coderConsistentWithEquals(coder, row1, row2); }
@Test public void testByteArrayEquality() { byte[] a0 = new byte[] {1, 2, 3, 4}; byte[] b0 = new byte[] {1, 2, 3, 4}; Schema schema = Schema.of(Schema.Field.of("bytes", Schema.FieldType.BYTES)); Row a = Row.withSchema(schema).addValue(a0).build(); Row b = Row.withSchema(schema).addValue(b0).build(); Assert.assertEquals(a, b); }
private static void unnestRow(Row input, Row.Builder output) { for (int i = 0; i < input.getSchema().getFieldCount(); ++i) { Field field = input.getSchema().getField(i); if (field.getType().getTypeName().isCompositeType()) { unnestRow(input.getRow(i), output); } else { output.addValue(input.getValue(i)); } } } /** A {@link PTransform} that unnests nested row. */
@Test public void testConsistentWithEqualsArrayOfBytes() throws Exception { FieldType fieldType = FieldType.array(FieldType.BYTES); Schema schema = Schema.of(Schema.Field.of("f1", fieldType)); RowCoder coder = RowCoder.of(schema); List<byte[]> list1 = Collections.singletonList(new byte[] {1, 2, 3, 4}); Row row1 = Row.withSchema(schema).addValue(list1).build(); List<byte[]> list2 = Collections.singletonList(new byte[] {1, 2, 3, 4}); Row row2 = Row.withSchema(schema).addValue(list2).build(); Assume.assumeTrue(coder.consistentWithEquals()); CoderProperties.coderConsistentWithEquals(coder, row1, row2); }
@Test public void testCreateRegisteredSchema() { p.getSchemaRegistry() .registerSchemaForClass( String.class, STRING_SCHEMA, s -> Row.withSchema(STRING_SCHEMA).addValue(s).build(), r -> r.getString("field")); PCollection<String> out = p.apply(Create.of("a", "b", "c", "d")); assertThat(out.getCoder(), instanceOf(SchemaCoder.class)); }
@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(); } }
@ProcessElement public void process(@Element Row inputRow, OutputReceiver<Row> output) { for (Object element : inputRow.getArray(0)) { output.output(Row.withSchema(schema).addValue(element).build()); } } }
@Test public void testCreateExplicitSchema() { PCollection<String> out = p.apply( Create.of("a", "b", "c", "d") .withSchema( STRING_SCHEMA, s -> Row.withSchema(STRING_SCHEMA).addValue(s).build(), r -> r.getString("field"))); assertThat(out.getCoder(), instanceOf(SchemaCoder.class)); }