Skip to content

[Python Schema] Python schema support custom Avro configurations for Enum type #12642

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
codelipenghui merged 3 commits into from
Nov 7, 2021
Merged

[Python Schema] Python schema support custom Avro configurations for Enum type #12642

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
386a1b3
Python client support set `required`, `default`, `required_default` f…
gaoran10 Nov 5, 2021
a40943f
fix test
gaoran10 Nov 6, 2021
2cdcd49
fix test
gaoran10 Nov 7, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion pulsar-client-cpp/python/examples/company.avsc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@
{"name": "age", "type": ["null", "int"]}
]
}}]},
{"name": "labels", "type": ["null", {"type": "map", "values": "string"}]}
{"name": "labels", "type": ["null", {"type": "map", "values": "string"}]},
{"name": "companyType", "type": ["null", {"type": "enum", "name": "CompanyType", "symbols":
["companyType1", "companyType2", "companyType3"]}]}
]
}
2 changes: 1 addition & 1 deletion pulsar-client-cpp/python/pulsar/schema/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
#

from .definition import Record, Field, Null, Boolean, Integer, Long, \
Float, Double, Bytes, String, Array, Map
Float, Double, Bytes, String, Array, Map, CustomEnum

from .schema import Schema, BytesSchema, StringSchema, JsonSchema
from .schema_avro import AvroSchema
18 changes: 12 additions & 6 deletions pulsar-client-cpp/python/pulsar/schema/definition.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,8 +44,7 @@ def _get_fields(cls, dct):
fields = OrderedDict()
for name, value in dct.items():
if issubclass(type(value), EnumMeta):
# Wrap Python enums
value = _Enum(value)
value = CustomEnum(value)
elif type(value) == RecordMeta:
# We expect an instance of a record rather than the class itself
value = value()
Expand Down Expand Up @@ -125,6 +124,12 @@ def schema_info(cls, defined_names):
schema['namespace'] = cls._avro_namespace
schema['fields'] = []

def get_filed_default_value(value):
if isinstance(value, Enum):
return value.name
else:
return value

if cls._sorted_fields:
fields = sorted(cls._fields.keys())
else:
Expand All @@ -135,7 +140,7 @@ def schema_info(cls, defined_names):
if field._required else ['null', field.schema_info(defined_names)]
schema['fields'].append({
'name': name,
'default': field.default(),
'default': get_filed_default_value(field.default()),
'type': field_type
}) if field.required_default() else schema['fields'].append({
'name': name,
Expand Down Expand Up @@ -360,15 +365,16 @@ def default(self):

# Complex types

class _Enum(Field):
def __init__(self, enum_type):

class CustomEnum(Field):
def __init__(self, enum_type, default=None, required=False, required_default=False):
if not issubclass(enum_type, Enum):
raise Exception(enum_type + " is not a valid Enum type")
self.enum_type = enum_type
self.values = {}
for x in enum_type.__members__.values():
self.values[x.value] = x
super(_Enum, self).__init__()
super(CustomEnum, self).__init__(default, required, required_default)

def type(self):
return 'enum'
Expand Down
118 changes: 58 additions & 60 deletions pulsar-client-cpp/python/schema_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@ class Example(Record):
g = Double()
h = Bytes()
i = Map(String())
j = CustomEnum(Color)

fastavro.parse_schema(Example.schema())
self.assertEqual(Example.schema(), {
Expand All @@ -74,16 +75,23 @@ class Example(Record):
{"name": "i", "type": ["null", {
"type": "map",
"values": "string"}]
},
},
{"name": "j", "type": ["null", "Color"]}
]
})

def test_complex(self):
class Color(Enum):
red = 1
green = 2
blue = 3

class MySubRecord(Record):
_sorted_fields = True
x = Integer()
y = Long()
z = String()
color = CustomEnum(Color)

class Example(Record):
_sorted_fields = True
Expand All @@ -101,9 +109,12 @@ class Example(Record):
"type": ["null", {
"name": "MySubRecord",
"type": "record",
"fields": [{"name": "x", "type": ["null", "int"]},
{"name": "y", "type": ["null", "long"]},
{"name": "z", "type": ["null", "string"]}]
"fields": [
{'name': 'color', 'type': ['null', {'type': 'enum', 'name': 'Color', 'symbols':
['red', 'green', 'blue']}]},
{"name": "x", "type": ["null", "int"]},
{"name": "y", "type": ["null", "long"]},
{"name": "z", "type": ["null", "string"]}]
}]
},
{"name": "sub2",
Expand Down Expand Up @@ -630,6 +641,8 @@ class MyEnum(Enum):
class Example(Record):
name = String()
v = MyEnum
w = CustomEnum(MyEnum)
x = CustomEnum(MyEnum, required=True, default=MyEnum.A, required_default=True)

topic = 'my-json-enum-topic'

Expand All @@ -641,13 +654,15 @@ class Example(Record):
consumer = client.subscribe(topic, 'test',
schema=JsonSchema(Example))

r = Example(name='test', v=MyEnum.C)
r = Example(name='test', v=MyEnum.C, w=MyEnum.B)
producer.send(r)

msg = consumer.receive()

self.assertEqual('test', msg.value().name)
self.assertEqual(MyEnum.C, MyEnum(msg.value().v))
self.assertEqual(MyEnum.B, MyEnum(msg.value().w))
self.assertEqual(MyEnum.A, MyEnum(msg.value().x))
client.close()

def test_avro_enum(self):
Expand All @@ -659,6 +674,8 @@ class MyEnum(Enum):
class Example(Record):
name = String()
v = MyEnum
w = CustomEnum(MyEnum)
x = CustomEnum(MyEnum, required=True, default=MyEnum.B, required_default=True)

topic = 'my-avro-enum-topic'

Expand All @@ -670,12 +687,14 @@ class Example(Record):
consumer = client.subscribe(topic, 'test',
schema=AvroSchema(Example))

r = Example(name='test', v=MyEnum.C)
r = Example(name='test', v=MyEnum.C, w=MyEnum.A)
producer.send(r)

msg = consumer.receive()
msg.value()
self.assertEqual(MyEnum.C, msg.value().v)
self.assertEqual(MyEnum.A, MyEnum(msg.value().w))
self.assertEqual(MyEnum.B, MyEnum(msg.value().x))
client.close()

def test_avro_map_array(self):
Expand Down Expand Up @@ -913,6 +932,11 @@ class MyRecord(Record):
client.close()

def test_serialize_schema_complex(self):
class Color(Enum):
red = 1
green = 2
blue = 3

class NestedObj1(Record):
_sorted_fields = True
na1 = String()
Expand All @@ -925,6 +949,8 @@ class NestedObj2(Record):
nc2 = NestedObj1()

class NestedObj3(Record):
_sorted_fields = True
color = CustomEnum(Color)
na3 = Integer()

class NestedObj4(Record):
Expand All @@ -933,18 +959,14 @@ class NestedObj4(Record):
na4 = String()
nb4 = Integer()

class Color(Enum):
red = 1
green = 2
blue = 3

class ComplexRecord(Record):
_avro_namespace = 'xxx.xxx'
_sorted_fields = True
a = Integer()
b = Integer()
color = Color
color2 = Color
color3 = CustomEnum(Color, required=True, default=Color.red, required_default=True)
nested = NestedObj2()
nested2 = NestedObj2()
mapNested = Map(NestedObj3())
Expand All @@ -970,8 +992,10 @@ class ComplexRecord(Record):
{'name': 'color', 'type': ['null', {'type': 'enum', 'name': 'Color', 'symbols': [
'red', 'green', 'blue']}]},
{'name': 'color2', 'type': ['null', 'Color']},
{'name': 'color3', 'default': 'red', 'type': 'Color'},
{'name': 'mapNested', 'type': ['null', {'type': 'map', 'values':
{'name': 'NestedObj3', 'type': 'record', 'fields': [
{'name': 'color', 'type': ['null', 'Color']},
{'name': 'na3', 'type': ['null', 'int']}
]}}
]},
Expand All @@ -998,12 +1022,12 @@ def encode_and_decode(schema_type):
r = ComplexRecord(a=1, b=2, color=Color.red, color2=Color.blue,
nested=nested_obj2, nested2=nested_obj2,
mapNested={
'a': NestedObj3(na3=1),
'a': NestedObj3(na3=1, color=Color.green),
'b': NestedObj3(na3=2),
'c': NestedObj3(na3=3)
'c': NestedObj3(na3=3, color=Color.red)
}, mapNested2={
'd': NestedObj3(na3=4),
'e': NestedObj3(na3=5),
'd': NestedObj3(na3=4, color=Color.red),
'e': NestedObj3(na3=5, color=Color.blue),
'f': NestedObj3(na3=6)
}, arrayNested=[
NestedObj4(na4='value na4 1', nb4=100),
Expand All @@ -1017,32 +1041,9 @@ def encode_and_decode(schema_type):
data_decode = data_schema.decode(data_encode)
self.assertEqual(data_decode.__class__.__name__, 'ComplexRecord')
self.assertEqual(data_decode, r)
self.assertEqual(data_decode.a, 1)
self.assertEqual(data_decode.b, 2)
self.assertEqual(data_decode.color, Color.red)
self.assertEqual(data_decode.color2, Color.blue)
self.assertEqual(data_decode.nested.na2, 22)
self.assertEqual(data_decode.nested.nb2, True)
self.assertEqual(data_decode.nested.nc2.na1, 'na1 value')
self.assertEqual(data_decode.nested.nc2.nb1, 20.5)
self.assertEqual(data_decode.nested2.na2, 22)
self.assertEqual(data_decode.nested2.nb2, True)
self.assertEqual(data_decode.nested2.nc2.na1, 'na1 value')
self.assertEqual(data_decode.nested2.nc2.nb1, 20.5)
self.assertEqual(data_decode.mapNested['a'].na3, 1)
self.assertEqual(data_decode.mapNested['b'].na3, 2)
self.assertEqual(data_decode.mapNested['c'].na3, 3)
self.assertEqual(data_decode.mapNested2['d'].na3, 4)
self.assertEqual(data_decode.mapNested2['e'].na3, 5)
self.assertEqual(data_decode.mapNested2['f'].na3, 6)
self.assertEqual(data_decode.arrayNested[0].na4, 'value na4 1')
self.assertEqual(data_decode.arrayNested[0].nb4, 100)
self.assertEqual(data_decode.arrayNested[1].na4, 'value na4 2')
self.assertEqual(data_decode.arrayNested[1].nb4, 200)
self.assertEqual(data_decode.arrayNested2[0].na4, 'value na4 3')
self.assertEqual(data_decode.arrayNested2[0].nb4, 300)
self.assertEqual(data_decode.arrayNested2[1].na4, 'value na4 4')
self.assertEqual(data_decode.arrayNested2[1].nb4, 400)
self.assertEqual(r.color3, Color.red)
self.assertEqual(r.mapNested['a'].color, Color.green)
self.assertEqual(r.mapNested['b'].color, None)
print('Encode and decode complex schema finish. schema_type: ', schema_type)

encode_and_decode('avro')
Expand All @@ -1069,8 +1070,12 @@ class NestedObj2(Record):
self.assertEqual(data_decode.na2, 1)
self.assertTrue(data_decode.nb2)


def test_produce_and_consume_complex_schema_data(self):
class Color(Enum):
red = 1
green = 2
blue = 3

class NestedObj1(Record):
na1 = String()
nb1 = Double()
Expand All @@ -1082,6 +1087,7 @@ class NestedObj2(Record):

class NestedObj3(Record):
na3 = Integer()
color = CustomEnum(Color, required=True, required_default=True, default=Color.blue)

class NestedObj4(Record):
na4 = String()
Expand All @@ -1090,6 +1096,7 @@ class NestedObj4(Record):
class ComplexRecord(Record):
a = Integer()
b = Integer()
color = CustomEnum(Color)
nested = NestedObj2()
mapNested = Map(NestedObj3())
arrayNested = Array(NestedObj4())
Expand All @@ -1112,8 +1119,8 @@ def produce_consume_test(schema_type):
nested_obj1 = NestedObj1(na1='na1 value', nb1=20.5)
nested_obj2 = NestedObj2(na2=22, nb2=True, nc2=nested_obj1)
r = ComplexRecord(a=1, b=2, nested=nested_obj2, mapNested={
'a': NestedObj3(na3=1),
'b': NestedObj3(na3=2),
'a': NestedObj3(na3=1, color=Color.red),
'b': NestedObj3(na3=2, color=Color.green),
'c': NestedObj3(na3=3)
}, arrayNested=[
NestedObj4(na4='value na4 1', nb4=100),
Expand All @@ -1125,19 +1132,6 @@ def produce_consume_test(schema_type):
value = msg.value()
self.assertEqual(value.__class__.__name__, 'ComplexRecord')
self.assertEqual(value, r)
self.assertEqual(value.a, 1)
self.assertEqual(value.b, 2)
self.assertEqual(value.nested.na2, 22)
self.assertEqual(value.nested.nb2, True)
self.assertEqual(value.nested.nc2.na1, 'na1 value')
self.assertEqual(value.nested.nc2.nb1, 20.5)
self.assertEqual(value.mapNested['a'].na3, 1)
self.assertEqual(value.mapNested['b'].na3, 2)
self.assertEqual(value.mapNested['c'].na3, 3)
self.assertEqual(value.arrayNested[0].na4, 'value na4 1')
self.assertEqual(value.arrayNested[0].nb4, 100)
self.assertEqual(value.arrayNested[1].na4, 'value na4 2')
self.assertEqual(value.arrayNested[1].nb4, 200)

print('Produce and consume complex schema data finish. schema_type', schema_type)

Expand All @@ -1163,7 +1157,8 @@ def encode_and_decode(schema_definition):
"industry": "software",
"scale": ">100",
"funds": "1000000.0"
}
},
"companyType": "companyType1"
}
data = avro_schema.encode(company)
company_decode = avro_schema.decode(data)
Expand All @@ -1185,7 +1180,9 @@ def encode_and_decode(schema_definition):
{'name': 'age', 'type': ['null', 'int']}
]
}}]},
{'name': 'labels', 'type': ['null', {'type': 'map', 'values': 'string'}]}
{'name': 'labels', 'type': ['null', {'type': 'map', 'values': 'string'}]},
{'name': 'companyType', 'type': ['null', {'type': 'enum', 'name': 'CompanyType', 'symbols':
['companyType1', 'companyType2', 'companyType3']}]}
]
}
encode_and_decode(schema_definition)
Expand Down Expand Up @@ -1218,7 +1215,8 @@ def produce_and_consume(topic, schema_definition):
"industry": "software" + str(i),
"scale": ">100",
"funds": "1000000.0"
}
},
"companyType": "companyType" + str((i % 3) + 1)
}
producer.send(company)

Expand Down