#!/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))