| ################################################################################ |
| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| ################################################################################ |
| import glob |
| import os |
| import tempfile |
| from typing import Tuple, List |
| |
| from avro.datafile import DataFileReader |
| from avro.io import DatumReader |
| from py4j.java_gateway import JavaObject, java_import |
| |
| from pyflink.datastream import MapFunction |
| from pyflink.datastream.connectors.file_system import FileSink |
| from pyflink.datastream.formats.avro import AvroSchema, GenericRecordAvroTypeInfo, \ |
| AvroBulkWriters, AvroInputFormat |
| from pyflink.datastream.tests.test_util import DataStreamTestSinkFunction |
| from pyflink.java_gateway import get_gateway |
| from pyflink.testing.test_case_utils import PyFlinkUTTestCase |
| |
| |
| class FileSourceAvroInputFormatTests(PyFlinkUTTestCase): |
| |
| def setUp(self): |
| super().setUp() |
| self.test_sink = DataStreamTestSinkFunction() |
| self.avro_file_name = tempfile.mktemp(suffix='.avro', dir=self.tempdir) |
| _import_avro_classes() |
| |
| def test_avro_basic_read(self): |
| schema, records = _create_basic_avro_schema_and_records() |
| self._create_avro_file(schema, records) |
| self._build_avro_job(schema) |
| self.env.execute('test_avro_basic_read') |
| results = self.test_sink.get_results(True, False) |
| _check_basic_avro_schema_results(self, results) |
| |
| def test_avro_enum_read(self): |
| schema, records = _create_enum_avro_schema_and_records() |
| self._create_avro_file(schema, records) |
| self._build_avro_job(schema) |
| self.env.execute('test_avro_enum_read') |
| results = self.test_sink.get_results(True, False) |
| _check_enum_avro_schema_results(self, results) |
| |
| def test_avro_union_read(self): |
| schema, records = _create_union_avro_schema_and_records() |
| self._create_avro_file(schema, records) |
| self._build_avro_job(schema) |
| self.env.execute('test_avro_union_read') |
| results = self.test_sink.get_results(True, False) |
| _check_union_avro_schema_results(self, results) |
| |
| def test_avro_array_read(self): |
| schema, records = _create_array_avro_schema_and_records() |
| self._create_avro_file(schema, records) |
| self._build_avro_job(schema) |
| self.env.execute('test_avro_array_read') |
| results = self.test_sink.get_results(True, False) |
| _check_array_avro_schema_results(self, results) |
| |
| def test_avro_map_read(self): |
| schema, records = _create_map_avro_schema_and_records() |
| self._create_avro_file(schema, records) |
| self._build_avro_job(schema) |
| self.env.execute('test_avro_map_read') |
| results = self.test_sink.get_results(True, False) |
| _check_map_avro_schema_results(self, results) |
| |
| def _build_avro_job(self, record_schema): |
| ds = self.env.create_input(AvroInputFormat(self.avro_file_name, record_schema)) |
| ds.map(PassThroughMapFunction()).add_sink(self.test_sink) |
| |
| def _create_avro_file(self, schema: AvroSchema, records: list): |
| jvm = get_gateway().jvm |
| j_file = jvm.java.io.File(self.avro_file_name) |
| j_datum_writer = jvm.org.apache.flink.avro.shaded.org.apache.avro.generic \ |
| .GenericDatumWriter() |
| j_file_writer = jvm.org.apache.flink.avro.shaded.org.apache.avro.file \ |
| .DataFileWriter(j_datum_writer) |
| j_file_writer.create(schema._j_schema, j_file) |
| for r in records: |
| j_file_writer.append(r) |
| j_file_writer.close() |
| |
| |
| class FileSinkAvroWritersTests(PyFlinkUTTestCase): |
| |
| def setUp(self) -> None: |
| super().setUp() |
| # NOTE: parallelism == 1 is required to keep the order of results |
| self.env.set_parallelism(1) |
| self.avro_dir_name = tempfile.mkdtemp(dir=self.tempdir) |
| |
| def test_avro_basic_write(self): |
| schema, objects = _create_basic_avro_schema_and_py_objects() |
| self._build_avro_job(schema, objects) |
| self.env.execute('test_avro_basic_write') |
| results = self._read_avro_file() |
| _check_basic_avro_schema_results(self, results) |
| |
| def test_avro_enum_write(self): |
| schema, objects = _create_enum_avro_schema_and_py_objects() |
| self._build_avro_job(schema, objects) |
| self.env.execute('test_avro_enum_write') |
| results = self._read_avro_file() |
| _check_enum_avro_schema_results(self, results) |
| |
| def test_avro_union_write(self): |
| schema, objects = _create_union_avro_schema_and_py_objects() |
| self._build_avro_job(schema, objects) |
| self.env.execute('test_avro_union_write') |
| results = self._read_avro_file() |
| _check_union_avro_schema_results(self, results) |
| |
| def test_avro_array_write(self): |
| schema, objects = _create_array_avro_schema_and_py_objects() |
| self._build_avro_job(schema, objects) |
| self.env.execute('test_avro_array_write') |
| results = self._read_avro_file() |
| _check_array_avro_schema_results(self, results) |
| |
| def test_avro_map_write(self): |
| schema, objects = _create_map_avro_schema_and_py_objects() |
| self._build_avro_job(schema, objects) |
| self.env.execute('test_avro_map_write') |
| results = self._read_avro_file() |
| _check_map_avro_schema_results(self, results) |
| |
| def _build_avro_job(self, schema, objects): |
| ds = self.env.from_collection(objects) |
| sink = FileSink.for_bulk_format( |
| self.avro_dir_name, AvroBulkWriters.for_generic_record(schema) |
| ).build() |
| ds.map(lambda e: e, output_type=GenericRecordAvroTypeInfo(schema)).sink_to(sink) |
| |
| def _read_avro_file(self) -> List[dict]: |
| records = [] |
| for file in glob.glob(os.path.join(os.path.join(self.avro_dir_name, '**/*'))): |
| for record in DataFileReader(open(file, 'rb'), DatumReader()): |
| records.append(record) |
| return records |
| |
| |
| class PassThroughMapFunction(MapFunction): |
| |
| def map(self, value): |
| return value |
| |
| |
| def _import_avro_classes(): |
| jvm = get_gateway().jvm |
| classes = ['org.apache.avro.generic.GenericData'] |
| prefix = 'org.apache.flink.avro.shaded.' |
| for cls in classes: |
| java_import(jvm, prefix + cls) |
| |
| |
| BASIC_SCHEMA = """ |
| { |
| "type": "record", |
| "name": "test", |
| "fields": [ |
| { "name": "null", "type": "null" }, |
| { "name": "boolean", "type": "boolean" }, |
| { "name": "int", "type": "int" }, |
| { "name": "long", "type": "long" }, |
| { "name": "float", "type": "float" }, |
| { "name": "double", "type": "double" }, |
| { "name": "string", "type": "string" } |
| ] |
| } |
| """ |
| |
| |
| def _create_basic_avro_schema_and_records() -> Tuple[AvroSchema, List[JavaObject]]: |
| schema = AvroSchema.parse_string(BASIC_SCHEMA) |
| records = [_create_basic_avro_record(schema, True, 0, 1, 2, 3, 's1'), |
| _create_basic_avro_record(schema, False, 4, 5, 6, 7, 's2')] |
| return schema, records |
| |
| |
| def _create_basic_avro_schema_and_py_objects() -> Tuple[AvroSchema, List[dict]]: |
| schema = AvroSchema.parse_string(BASIC_SCHEMA) |
| objects = [ |
| {'null': None, 'boolean': True, 'int': 0, 'long': 1, |
| 'float': 2., 'double': 3., 'string': 's1'}, |
| {'null': None, 'boolean': False, 'int': 4, 'long': 5, |
| 'float': 6., 'double': 7., 'string': 's2'}, |
| ] |
| return schema, objects |
| |
| |
| def _check_basic_avro_schema_results(test, results): |
| result1 = results[0] |
| result2 = results[1] |
| test.assertEqual(result1['null'], None) |
| test.assertEqual(result1['boolean'], True) |
| test.assertEqual(result1['int'], 0) |
| test.assertEqual(result1['long'], 1) |
| test.assertAlmostEqual(result1['float'], 2, delta=1e-3) |
| test.assertAlmostEqual(result1['double'], 3, delta=1e-3) |
| test.assertEqual(result1['string'], 's1') |
| test.assertEqual(result2['null'], None) |
| test.assertEqual(result2['boolean'], False) |
| test.assertEqual(result2['int'], 4) |
| test.assertEqual(result2['long'], 5) |
| test.assertAlmostEqual(result2['float'], 6, delta=1e-3) |
| test.assertAlmostEqual(result2['double'], 7, delta=1e-3) |
| test.assertEqual(result2['string'], 's2') |
| |
| |
| ENUM_SCHEMA = """ |
| { |
| "type": "record", |
| "name": "test", |
| "fields": [ |
| { |
| "name": "suit", |
| "type": { |
| "type": "enum", |
| "name": "Suit", |
| "symbols" : ["SPADES", "HEARTS", "DIAMONDS", "CLUBS"] |
| } |
| } |
| ] |
| } |
| """ |
| |
| |
| def _create_enum_avro_schema_and_records() -> Tuple[AvroSchema, List[JavaObject]]: |
| schema = AvroSchema.parse_string(ENUM_SCHEMA) |
| records = [_create_enum_avro_record(schema, 'SPADES'), |
| _create_enum_avro_record(schema, 'DIAMONDS')] |
| return schema, records |
| |
| |
| def _create_enum_avro_schema_and_py_objects() -> Tuple[AvroSchema, List[dict]]: |
| schema = AvroSchema.parse_string(ENUM_SCHEMA) |
| records = [ |
| {'suit': 'SPADES'}, |
| {'suit': 'DIAMONDS'}, |
| ] |
| return schema, records |
| |
| |
| def _check_enum_avro_schema_results(test, results): |
| test.assertEqual(results[0]['suit'], 'SPADES') |
| test.assertEqual(results[1]['suit'], 'DIAMONDS') |
| |
| |
| UNION_SCHEMA = """ |
| { |
| "type": "record", |
| "name": "test", |
| "fields": [ |
| { |
| "name": "union", |
| "type": [ "int", "double", "null" ] |
| } |
| ] |
| } |
| """ |
| |
| |
| def _create_union_avro_schema_and_records() -> Tuple[AvroSchema, List[JavaObject]]: |
| schema = AvroSchema.parse_string(UNION_SCHEMA) |
| records = [_create_union_avro_record(schema, 1), |
| _create_union_avro_record(schema, 2.), |
| _create_union_avro_record(schema, None)] |
| return schema, records |
| |
| |
| def _create_union_avro_schema_and_py_objects() -> Tuple[AvroSchema, List[dict]]: |
| schema = AvroSchema.parse_string(UNION_SCHEMA) |
| records = [ |
| {'union': 1}, |
| {'union': 2.}, |
| {'union': None}, |
| ] |
| return schema, records |
| |
| |
| def _check_union_avro_schema_results(test, results): |
| test.assertEqual(results[0]['union'], 1) |
| test.assertAlmostEqual(results[1]['union'], 2.0, delta=1e-3) |
| test.assertEqual(results[2]['union'], None) |
| |
| |
| # It seems there's bug when array item record contains only one field, which throws |
| # java.lang.ClassCastException: required ... is not a group when reading |
| ARRAY_SCHEMA = """ |
| { |
| "type": "record", |
| "name": "test", |
| "fields": [ |
| { |
| "name": "array", |
| "type": { |
| "type": "array", |
| "items": { |
| "type": "record", |
| "name": "item", |
| "fields": [ |
| { "name": "int", "type": "int" }, |
| { "name": "double", "type": "double" } |
| ] |
| } |
| } |
| } |
| ] |
| } |
| """ |
| |
| |
| def _create_array_avro_schema_and_records() -> Tuple[AvroSchema, List[JavaObject]]: |
| schema = AvroSchema.parse_string(ARRAY_SCHEMA) |
| records = [_create_array_avro_record(schema, [(1, 2.), (3, 4.)]), |
| _create_array_avro_record(schema, [(5, 6.), (7, 8.)])] |
| return schema, records |
| |
| |
| def _create_array_avro_schema_and_py_objects() -> Tuple[AvroSchema, List[dict]]: |
| schema = AvroSchema.parse_string(ARRAY_SCHEMA) |
| records = [ |
| {'array': [{'int': 1, 'double': 2.}, {'int': 3, 'double': 4.}]}, |
| {'array': [{'int': 5, 'double': 6.}, {'int': 7, 'double': 8.}]}, |
| ] |
| return schema, records |
| |
| |
| def _check_array_avro_schema_results(test, results): |
| result1 = results[0] |
| result2 = results[1] |
| test.assertEqual(result1['array'][0]['int'], 1) |
| test.assertAlmostEqual(result1['array'][0]['double'], 2., delta=1e-3) |
| test.assertEqual(result1['array'][1]['int'], 3) |
| test.assertAlmostEqual(result1['array'][1]['double'], 4., delta=1e-3) |
| test.assertEqual(result2['array'][0]['int'], 5) |
| test.assertAlmostEqual(result2['array'][0]['double'], 6., delta=1e-3) |
| test.assertEqual(result2['array'][1]['int'], 7) |
| test.assertAlmostEqual(result2['array'][1]['double'], 8., delta=1e-3) |
| |
| |
| MAP_SCHEMA = """ |
| { |
| "type": "record", |
| "name": "test", |
| "fields": [ |
| { |
| "name": "map", |
| "type": { |
| "type": "map", |
| "values": "long" |
| } |
| } |
| ] |
| } |
| """ |
| |
| |
| def _create_map_avro_schema_and_records() -> Tuple[AvroSchema, List[JavaObject]]: |
| schema = AvroSchema.parse_string(MAP_SCHEMA) |
| records = [_create_map_avro_record(schema, {'a': 1, 'b': 2}), |
| _create_map_avro_record(schema, {'c': 3, 'd': 4})] |
| return schema, records |
| |
| |
| def _create_map_avro_schema_and_py_objects() -> Tuple[AvroSchema, List[dict]]: |
| schema = AvroSchema.parse_string(MAP_SCHEMA) |
| records = [ |
| {'map': {'a': 1, 'b': 2}}, |
| {'map': {'c': 3, 'd': 4}}, |
| ] |
| return schema, records |
| |
| |
| def _check_map_avro_schema_results(test, results): |
| result1 = results[0] |
| result2 = results[1] |
| test.assertEqual(result1['map']['a'], 1) |
| test.assertEqual(result1['map']['b'], 2) |
| test.assertEqual(result2['map']['c'], 3) |
| test.assertEqual(result2['map']['d'], 4) |
| |
| |
| def _create_basic_avro_record(schema: AvroSchema, boolean_value, int_value, long_value, |
| float_value, double_value, string_value): |
| jvm = get_gateway().jvm |
| j_record = jvm.GenericData.Record(schema._j_schema) |
| j_record.put('boolean', boolean_value) |
| j_record.put('int', int_value) |
| j_record.put('long', long_value) |
| j_record.put('float', float_value) |
| j_record.put('double', double_value) |
| j_record.put('string', string_value) |
| return j_record |
| |
| |
| def _create_enum_avro_record(schema: AvroSchema, enum_value): |
| jvm = get_gateway().jvm |
| j_record = jvm.GenericData.Record(schema._j_schema) |
| j_enum = jvm.GenericData.EnumSymbol(schema._j_schema.getField('suit').schema(), enum_value) |
| j_record.put('suit', j_enum) |
| return j_record |
| |
| |
| def _create_union_avro_record(schema, union_value): |
| jvm = get_gateway().jvm |
| j_record = jvm.GenericData.Record(schema._j_schema) |
| j_record.put('union', union_value) |
| return j_record |
| |
| |
| def _create_array_avro_record(schema, item_values: list): |
| jvm = get_gateway().jvm |
| j_record = jvm.GenericData.Record(schema._j_schema) |
| item_schema = AvroSchema(schema._j_schema.getField('array').schema().getElementType()) |
| j_array = jvm.java.util.ArrayList() |
| for idx, item_value in enumerate(item_values): |
| j_item = jvm.GenericData.Record(item_schema._j_schema) |
| j_item.put('int', item_value[0]) |
| j_item.put('double', item_value[1]) |
| j_array.add(j_item) |
| j_record.put('array', j_array) |
| return j_record |
| |
| |
| def _create_map_avro_record(schema, map: dict): |
| jvm = get_gateway().jvm |
| j_record = jvm.GenericData.Record(schema._j_schema) |
| j_map = jvm.java.util.HashMap() |
| for k, v in map.items(): |
| j_map.put(k, v) |
| j_record.put('map', j_map) |
| return j_record |
