| #!/usr/bin/env python |
| |
| ## |
| # 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 |
| # |
| # https://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. |
| |
| """Read/Write Avro File Object Containers.""" |
| |
| from __future__ import absolute_import, division, print_function |
| |
| import io |
| import os |
| import random |
| import zlib |
| |
| import avro.io |
| import avro.schema |
| from avro.codecs import Codecs |
| |
| # |
| # Constants |
| # |
| VERSION = 1 |
| MAGIC = bytes(b'Obj' + bytearray([VERSION])) |
| MAGIC_SIZE = len(MAGIC) |
| SYNC_SIZE = 16 |
| SYNC_INTERVAL = 4000 * SYNC_SIZE # TODO(hammer): make configurable |
| META_SCHEMA = avro.schema.parse("""\ |
| {"type": "record", "name": "org.apache.avro.file.Header", |
| "fields" : [ |
| {"name": "magic", "type": {"type": "fixed", "name": "magic", "size": %d}}, |
| {"name": "meta", "type": {"type": "map", "values": "bytes"}}, |
| {"name": "sync", "type": {"type": "fixed", "name": "sync", "size": %d}}]} |
| """ % (MAGIC_SIZE, SYNC_SIZE)) |
| |
| NULL_CODEC = 'null' |
| VALID_CODECS = Codecs.supported_codec_names() |
| VALID_ENCODINGS = ['binary'] # not used yet |
| |
| CODEC_KEY = "avro.codec" |
| SCHEMA_KEY = "avro.schema" |
| |
| # |
| # Exceptions |
| # |
| |
| |
| class DataFileException(avro.schema.AvroException): |
| """ |
| Raised when there's a problem reading or writing file object containers. |
| """ |
| |
| def __init__(self, fail_msg): |
| avro.schema.AvroException.__init__(self, fail_msg) |
| |
| # |
| # Write Path |
| # |
| |
| |
| class _DataFile(object): |
| """Mixin for methods common to both reading and writing.""" |
| |
| block_count = 0 |
| _meta = None |
| _sync_marker = None |
| |
| def __enter__(self): |
| return self |
| |
| def __exit__(self, type, value, traceback): |
| # Perform a close if there's no exception |
| if type is None: |
| self.close() |
| |
| def get_meta(self, key): |
| return self.meta.get(key) |
| |
| def set_meta(self, key, val): |
| self.meta[key] = val |
| |
| @property |
| def sync_marker(self): |
| return self._sync_marker |
| |
| @property |
| def meta(self): |
| """Read-only dictionary of metadata for this datafile.""" |
| if self._meta is None: |
| self._meta = {} |
| return self._meta |
| |
| @property |
| def codec(self): |
| """Meta are stored as bytes, but codec is returned as a string.""" |
| try: |
| return self.get_meta(CODEC_KEY).decode() |
| except AttributeError: |
| return "null" |
| |
| @codec.setter |
| def codec(self, value): |
| """Meta are stored as bytes, but codec is set as a string.""" |
| if value not in VALID_CODECS: |
| raise DataFileException("Unknown codec: {!r}".format(value)) |
| self.set_meta(CODEC_KEY, value.encode()) |
| |
| @property |
| def schema(self): |
| """Meta are stored as bytes, but schema is returned as a string.""" |
| return self.get_meta(SCHEMA_KEY).decode() |
| |
| @schema.setter |
| def schema(self, value): |
| """Meta are stored as bytes, but schema is set as a string.""" |
| self.set_meta(SCHEMA_KEY, value.encode()) |
| |
| |
| class DataFileWriter(_DataFile): |
| |
| # TODO(hammer): make 'encoder' a metadata property |
| def __init__(self, writer, datum_writer, writers_schema=None, codec=NULL_CODEC): |
| """ |
| If the schema is not present, presume we're appending. |
| |
| @param writer: File-like object to write into. |
| """ |
| self._writer = writer |
| self._encoder = avro.io.BinaryEncoder(writer) |
| self._datum_writer = datum_writer |
| self._buffer_writer = io.BytesIO() |
| self._buffer_encoder = avro.io.BinaryEncoder(self._buffer_writer) |
| self.block_count = 0 |
| self._header_written = False |
| |
| if writers_schema is not None: |
| self._sync_marker = generate_sixteen_random_bytes() |
| self.codec = codec |
| self.schema = str(writers_schema) |
| self.datum_writer.writers_schema = writers_schema |
| else: |
| # open writer for reading to collect metadata |
| dfr = DataFileReader(writer, avro.io.DatumReader()) |
| |
| # TODO(hammer): collect arbitrary metadata |
| # collect metadata |
| self._sync_marker = dfr.sync_marker |
| self.codec = dfr.codec |
| |
| # get schema used to write existing file |
| self.schema = schema_from_file = dfr.schema |
| self.datum_writer.writers_schema = avro.schema.parse(schema_from_file) |
| |
| # seek to the end of the file and prepare for writing |
| writer.seek(0, 2) |
| self._header_written = True |
| |
| # read-only properties |
| writer = property(lambda self: self._writer) |
| encoder = property(lambda self: self._encoder) |
| datum_writer = property(lambda self: self._datum_writer) |
| buffer_writer = property(lambda self: self._buffer_writer) |
| buffer_encoder = property(lambda self: self._buffer_encoder) |
| |
| def _write_header(self): |
| header = {'magic': MAGIC, |
| 'meta': self.meta, |
| 'sync': self.sync_marker} |
| self.datum_writer.write_data(META_SCHEMA, header, self.encoder) |
| self._header_written = True |
| |
| @property |
| def codec(self): |
| """Meta are stored as bytes, but codec is returned as a string.""" |
| return self.get_meta(CODEC_KEY).decode() |
| |
| @codec.setter |
| def codec(self, value): |
| """Meta are stored as bytes, but codec is set as a string.""" |
| if value not in VALID_CODECS: |
| raise DataFileException("Unknown codec: {!r}".format(value)) |
| self.set_meta(CODEC_KEY, value.encode()) |
| |
| # TODO(hammer): make a schema for blocks and use datum_writer |
| def _write_block(self): |
| if not self._header_written: |
| self._write_header() |
| |
| if self.block_count > 0: |
| # write number of items in block |
| self.encoder.write_long(self.block_count) |
| |
| # write block contents |
| uncompressed_data = self.buffer_writer.getvalue() |
| codec = Codecs.get_codec(self.codec) |
| compressed_data, compressed_data_length = codec.compress(uncompressed_data) |
| |
| # Write length of block |
| self.encoder.write_long(compressed_data_length) |
| |
| # Write block |
| self.writer.write(compressed_data) |
| |
| # write sync marker |
| self.writer.write(self.sync_marker) |
| |
| # reset buffer |
| self.buffer_writer.truncate(0) |
| self.buffer_writer.seek(0) |
| self.block_count = 0 |
| |
| def append(self, datum): |
| """Append a datum to the file.""" |
| self.datum_writer.write(datum, self.buffer_encoder) |
| self.block_count += 1 |
| |
| # if the data to write is larger than the sync interval, write the block |
| if self.buffer_writer.tell() >= SYNC_INTERVAL: |
| self._write_block() |
| |
| def sync(self): |
| """ |
| Return the current position as a value that may be passed to |
| DataFileReader.seek(long). Forces the end of the current block, |
| emitting a synchronization marker. |
| """ |
| self._write_block() |
| return self.writer.tell() |
| |
| def flush(self): |
| """Flush the current state of the file, including metadata.""" |
| self._write_block() |
| self.writer.flush() |
| |
| def close(self): |
| """Close the file.""" |
| self.flush() |
| self.writer.close() |
| |
| |
| class DataFileReader(_DataFile): |
| """Read files written by DataFileWriter.""" |
| # TODO(hammer): allow user to specify expected schema? |
| # TODO(hammer): allow user to specify the encoder |
| |
| def __init__(self, reader, datum_reader): |
| self._reader = reader |
| self._raw_decoder = avro.io.BinaryDecoder(reader) |
| self._datum_decoder = None # Maybe reset at every block. |
| self._datum_reader = datum_reader |
| |
| # read the header: magic, meta, sync |
| self._read_header() |
| |
| # get file length |
| self._file_length = self.determine_file_length() |
| |
| # get ready to read |
| self.block_count = 0 |
| self.datum_reader.writers_schema = avro.schema.parse(self.schema) |
| |
| def __iter__(self): |
| return self |
| |
| # read-only properties |
| reader = property(lambda self: self._reader) |
| raw_decoder = property(lambda self: self._raw_decoder) |
| datum_decoder = property(lambda self: self._datum_decoder) |
| datum_reader = property(lambda self: self._datum_reader) |
| file_length = property(lambda self: self._file_length) |
| |
| def determine_file_length(self): |
| """ |
| Get file length and leave file cursor where we found it. |
| """ |
| remember_pos = self.reader.tell() |
| self.reader.seek(0, 2) |
| file_length = self.reader.tell() |
| self.reader.seek(remember_pos) |
| return file_length |
| |
| def is_EOF(self): |
| return self.reader.tell() == self.file_length |
| |
| def _read_header(self): |
| # seek to the beginning of the file to get magic block |
| self.reader.seek(0, 0) |
| |
| # read header into a dict |
| header = self.datum_reader.read_data( |
| META_SCHEMA, META_SCHEMA, self.raw_decoder) |
| |
| # check magic number |
| if header.get('magic') != MAGIC: |
| fail_msg = "Not an Avro data file: %s doesn't match %s."\ |
| % (header.get('magic'), MAGIC) |
| raise avro.schema.AvroException(fail_msg) |
| |
| # set metadata |
| self._meta = header['meta'] |
| |
| # set sync marker |
| self._sync_marker = header['sync'] |
| |
| def _read_block_header(self): |
| self.block_count = self.raw_decoder.read_long() |
| codec = Codecs.get_codec(self.codec) |
| self._datum_decoder = codec.decompress(self.raw_decoder) |
| |
| def _skip_sync(self): |
| """ |
| Read the length of the sync marker; if it matches the sync marker, |
| return True. Otherwise, seek back to where we started and return False. |
| """ |
| proposed_sync_marker = self.reader.read(SYNC_SIZE) |
| if proposed_sync_marker != self.sync_marker: |
| self.reader.seek(-SYNC_SIZE, 1) |
| return False |
| return True |
| |
| def __next__(self): |
| """Return the next datum in the file.""" |
| while self.block_count == 0: |
| if self.is_EOF() or (self._skip_sync() and self.is_EOF()): |
| raise StopIteration |
| self._read_block_header() |
| |
| datum = self.datum_reader.read(self.datum_decoder) |
| self.block_count -= 1 |
| return datum |
| next = __next__ |
| |
| def close(self): |
| """Close this reader.""" |
| self.reader.close() |
| |
| |
| def generate_sixteen_random_bytes(): |
| try: |
| return os.urandom(16) |
| except NotImplementedError: |
| return bytes(random.randrange(256) for i in range(16)) |
