def test_row_coder_nested_struct(self): Pair = typing.NamedTuple('Pair', [('left', Person), ('right', Person)]) value = Pair(self.PEOPLE[0], self.PEOPLE[1]) coder = RowCoder(typing_to_runner_api(Pair).row_type.schema) self.assertEqual(value, coder.decode(coder.encode(value)))
def test_create_row_coder_from_schema(self): schema = schema_pb2.Schema( id="person", fields=[ schema_pb2.Field( name="name", type=schema_pb2.FieldType(atomic_type=schema_pb2.STRING)), schema_pb2.Field( name="age", type=schema_pb2.FieldType(atomic_type=schema_pb2.INT32)), schema_pb2.Field(name="address", type=schema_pb2.FieldType( atomic_type=schema_pb2.STRING, nullable=True)), schema_pb2.Field( name="aliases", type=schema_pb2.FieldType(array_type=schema_pb2.ArrayType( element_type=schema_pb2.FieldType( atomic_type=schema_pb2.STRING)))), schema_pb2.Field( name="knows_javascript", type=schema_pb2.FieldType(atomic_type=schema_pb2.BOOLEAN)), schema_pb2.Field(name="payload", type=schema_pb2.FieldType( atomic_type=schema_pb2.BYTES, nullable=True)), ]) coder = RowCoder(schema) for test_case in self.PEOPLE: self.assertEqual(test_case, coder.decode(coder.encode(test_case)))
def test_encoding_position_reorder_fields(self): fields = [("field1", str), ("field2", int), ("field3", int)] expected = typing.NamedTuple('expected', fields) reorder = schema_pb2.Schema( id="new_order", fields=[ schema_pb2.Field( name="field3", type=schema_pb2.FieldType(atomic_type=schema_pb2.STRING), encoding_position=2), schema_pb2.Field( name="field2", type=schema_pb2.FieldType(atomic_type=schema_pb2.INT32), encoding_position=1), schema_pb2.Field( name="field1", type=schema_pb2.FieldType(atomic_type=schema_pb2.INT32), encoding_position=0) ]) old_coder = RowCoder.from_type_hint(expected, None) new_coder = RowCoder(reorder) encode_expected = old_coder.encode(expected("foo", 7, 12)) encode_reorder = new_coder.encode(expected(12, 7, "foo")) self.assertEqual(encode_expected, encode_reorder)
def test_create_row_coder_from_schema(self): schema = schema_pb2.Schema( id="person", fields=[ schema_pb2.Field( name="name", type=schema_pb2.FieldType(atomic_type=schema_pb2.STRING)), schema_pb2.Field( name="age", type=schema_pb2.FieldType(atomic_type=schema_pb2.INT32)), schema_pb2.Field( name="address", type=schema_pb2.FieldType( atomic_type=schema_pb2.STRING, nullable=True)), schema_pb2.Field( name="aliases", type=schema_pb2.FieldType( array_type=schema_pb2.ArrayType( element_type=schema_pb2.FieldType( atomic_type=schema_pb2.STRING)))), schema_pb2.Field( name="knows_javascript", type=schema_pb2.FieldType(atomic_type=schema_pb2.BOOLEAN)), schema_pb2.Field( name="payload", type=schema_pb2.FieldType( atomic_type=schema_pb2.BYTES, nullable=True)), schema_pb2.Field( name="custom_metadata", type=schema_pb2.FieldType( map_type=schema_pb2.MapType( key_type=schema_pb2.FieldType( atomic_type=schema_pb2.STRING), value_type=schema_pb2.FieldType( atomic_type=schema_pb2.INT64), ))), schema_pb2.Field( name="favorite_time", type=schema_pb2.FieldType( logical_type=schema_pb2.LogicalType( urn="beam:logical_type:micros_instant:v1", representation=schema_pb2.FieldType( row_type=schema_pb2.RowType( schema=schema_pb2.Schema( id="micros_instant", fields=[ schema_pb2.Field( name="seconds", type=schema_pb2.FieldType( atomic_type=schema_pb2.INT64)), schema_pb2.Field( name="micros", type=schema_pb2.FieldType( atomic_type=schema_pb2.INT64)), ])))))), ]) coder = RowCoder(schema) for test_case in self.PEOPLE: self.assertEqual(test_case, coder.decode(coder.encode(test_case)))
def test_create_row_coder_from_named_tuple(self): expected_coder = RowCoder(typing_to_runner_api(Person).row_type.schema) real_coder = coders_registry.get_coder(Person) for test_case in self.PEOPLE: self.assertEqual(expected_coder.encode(test_case), real_coder.encode(test_case)) self.assertEqual(test_case, real_coder.decode(real_coder.encode(test_case)))
def test_row_accepts_trailing_zeros_truncated(self): expected_coder = RowCoder( typing_to_runner_api(NullablePerson).row_type.schema) person = NullablePerson(None, np.int32(25), "Westeros", ["Mother of Dragons"], False, None, {"dragons": 3}, None, "NotNull") out = expected_coder.encode(person) # 9 fields, 1 null byte, field 0, 5, 7 are null new_payload = bytes([9, 1, 1 | 1 << 5 | 1 << 7]) + out[4:] new_value = expected_coder.decode(new_payload) self.assertEqual(person, new_value)