add support for a codec compression level (used only by deflate currently)
MODIFY
fastavro/_write.pyx
fastavro/_write.pyx
@@ -346,7 +346,7 @@ cpdef write_header(bytearray fo, dict metadata, bytes sync_marker):
write_data(fo, header, HEADER_SCHEMA)
-cpdef null_write_block(object fo, bytes block_bytes):
+cpdef null_write_block(object fo, bytes block_bytes, compression_level):
"""Write block in "null" codec."""
cdef bytearray tmp = bytearray()
write_long(tmp, len(block_bytes))
@@ -354,11 +354,14 @@ cpdef null_write_block(object fo, bytes block_bytes):
fo.write(block_bytes)
-cpdef deflate_write_block(object fo, bytes block_bytes):
+cpdef deflate_write_block(object fo, bytes block_bytes, compression_level):
"""Write block in "deflate" codec."""
cdef bytearray tmp = bytearray()
# The first two characters and last character are zlib
# wrappers around deflate data.
+ if compression_level is not None:
+ data = zlib.compress(block_bytes, compression_level)[2:-1]
+ else:
data = zlib.compress(block_bytes)[2:-1]
write_long(tmp, len(data))
@@ -366,7 +369,7 @@ cpdef deflate_write_block(object fo, bytes block_bytes):
fo.write(data)
-cpdef bzip2_write_block(object fo, bytes block_bytes):
+cpdef bzip2_write_block(object fo, bytes block_bytes, compression_level):
"""Write block in "bzip2" codec."""
cdef bytearray tmp = bytearray()
data = bz2.compress(block_bytes)
@@ -383,7 +386,7 @@ BLOCK_WRITERS = {
def _missing_dependency(codec, library):
- def missing(fo, block_bytes):
+ def missing(fo, block_bytes, compression_level):
raise ValueError(
"{} codec is supported but you ".format(codec)
+ "need to install {}".format(library)
@@ -399,7 +402,7 @@ else:
BLOCK_WRITERS['snappy'] = snappy_write_block
-cpdef snappy_write_block(object fo, bytes block_bytes):
+cpdef snappy_write_block(object fo, bytes block_bytes, compression_level):
"""Write block in "snappy" codec."""
cdef bytearray tmp = bytearray()
data = snappy.compress(block_bytes)
@@ -420,7 +423,7 @@ else:
BLOCK_WRITERS["zstandard"] = zstandard_write_block
-cpdef zstandard_write_block(object fo, bytes block_bytes):
+cpdef zstandard_write_block(object fo, bytes block_bytes, compression_level):
"""Write block in "zstandard" codec."""
cdef bytearray tmp = bytearray()
data = zstd.ZstdCompressor().compress(block_bytes)
@@ -437,7 +440,7 @@ else:
BLOCK_WRITERS["lz4"] = lz4_write_block
-cpdef lz4_write_block(object fo, bytes block_bytes):
+cpdef lz4_write_block(object fo, bytes block_bytes, compression_level):
"""Write block in "lz4" codec."""
cdef bytearray tmp = bytearray()
data = lz4.block.compress(block_bytes)
@@ -472,6 +475,7 @@ cdef class Writer(object):
cdef public object metadata
cdef public long64 sync_interval
cdef public object block_writer
+ cdef public object compression_level
def __init__(self,
fo,
@@ -480,7 +484,8 @@ cdef class Writer(object):
sync_interval=1000 * SYNC_SIZE,
metadata=None,
validator=None,
- sync_marker=None):
+ sync_marker=None,
+ compression_level=None):
cdef bytearray tmp = bytearray()
self.fo = fo
@@ -489,6 +494,7 @@ cdef class Writer(object):
self.io = MemoryIO()
self.block_count = 0
self.sync_interval = sync_interval
+ self.compression_level = compression_level
if appendable(self.fo):
# Seed to the beginning to read the header
@@ -536,7 +542,7 @@ cdef class Writer(object):
cdef bytearray tmp = bytearray()
write_long(tmp, self.block_count)
self.fo.write(tmp)
- self.block_writer(self.fo, self.io.getvalue())
+ self.block_writer(self.fo, self.io.getvalue(), self.compression_level)
self.fo.write(self.sync_marker)
self.io.clear()
self.block_count = 0
@@ -556,7 +562,7 @@ cdef class Writer(object):
cdef bytearray tmp = bytearray()
write_long(tmp, block.num_records)
self.fo.write(tmp)
- self.block_writer(self.fo, block.bytes_.getvalue())
+ self.block_writer(self.fo, block.bytes_.getvalue(), self.compression_level)
self.fo.write(self.sync_marker)
def flush(self):
@@ -572,7 +578,8 @@ def writer(fo,
sync_interval=1000 * SYNC_SIZE,
metadata=None,
validator=None,
- sync_marker=None):
+ sync_marker=None,
+ codec_compression_level=None):
# Sanity check that records is not a single dictionary (as that is a common
# mistake and the exception that gets raised is not helpful)
if isinstance(records, dict):
@@ -586,6 +593,7 @@ def writer(fo,
metadata,
validator,
sync_marker,
+ codec_compression_level,
)
for record in records:

MODIFY
fastavro/_write_py.py
fastavro/_write_py.py
@@ -245,22 +245,26 @@ def write_header(encoder, metadata, sync_marker):
write_data(encoder, header, HEADER_SCHEMA)
-def null_write_block(encoder, block_bytes):
+def null_write_block(encoder, block_bytes, compression_level):
"""Write block in "null" codec."""
encoder.write_long(len(block_bytes))
encoder._fo.write(block_bytes)
-def deflate_write_block(encoder, block_bytes):
+def deflate_write_block(encoder, block_bytes, compression_level):
"""Write block in "deflate" codec."""
# The first two characters and last character are zlib
# wrappers around deflate data.
+ if compression_level is not None:
+ data = zlib.compress(block_bytes, compression_level)[2:-1]
+ else:
+ data = zlib.compress(block_bytes)[2:-1]
data = zlib.compress(block_bytes)[2:-1]
encoder.write_long(len(data))
encoder._fo.write(data)
-def bzip2_write_block(encoder, block_bytes):
+def bzip2_write_block(encoder, block_bytes, compression_level):
"""Write block in "bzip2" codec."""
data = bz2.compress(block_bytes)
encoder.write_long(len(data))
@@ -275,7 +279,7 @@ BLOCK_WRITERS = {
def _missing_codec_lib(codec, library):
- def missing(encoder, block_bytes):
+ def missing(encoder, block_bytes, compression_level):
raise ValueError(
"{} codec is supported but you ".format(codec)
+ "need to install {}".format(library)
@@ -283,7 +287,7 @@ def _missing_codec_lib(codec, library):
return missing
-def snappy_write_block(encoder, block_bytes):
+def snappy_write_block(encoder, block_bytes, compression_level):
"""Write block in "snappy" codec."""
data = snappy.compress(block_bytes)
encoder.write_long(len(data) + 4) # for CRC
@@ -299,7 +303,7 @@ else:
BLOCK_WRITERS['snappy'] = snappy_write_block
-def zstandard_write_block(encoder, block_bytes):
+def zstandard_write_block(encoder, block_bytes, compression_level):
"""Write block in "zstandard" codec."""
data = zstd.ZstdCompressor().compress(block_bytes)
encoder.write_long(len(data))
@@ -314,7 +318,7 @@ else:
BLOCK_WRITERS["zstandard"] = zstandard_write_block
-def lz4_write_block(encoder, block_bytes):
+def lz4_write_block(encoder, block_bytes, compression_level):
"""Write block in "lz4" codec."""
data = lz4.block.compress(block_bytes)
encoder.write_long(len(data))
@@ -358,7 +362,8 @@ class Writer(GenericWriter):
sync_interval=1000 * SYNC_SIZE,
metadata=None,
validator=None,
- sync_marker=None):
+ sync_marker=None,
+ compression_level=None):
GenericWriter.__init__(self, schema, metadata, validator)
self.metadata['avro.codec'] = codec
@@ -369,6 +374,7 @@ class Writer(GenericWriter):
self.io = BinaryEncoder(MemoryIO())
self.block_count = 0
self.sync_interval = sync_interval
+ self.compression_level = compression_level
if appendable(self.encoder._fo):
# Seed to the beginning to read the header
@@ -401,7 +407,9 @@ class Writer(GenericWriter):
def dump(self):
self.encoder.write_long(self.block_count)
- self.block_writer(self.encoder, self.io._fo.getvalue())
+ self.block_writer(
+ self.encoder, self.io._fo.getvalue(), self.compression_level
+ )
self.encoder._fo.write(self.sync_marker)
self.io._fo.truncate(0)
self.io._fo.seek(0, SEEK_SET)
@@ -420,7 +428,9 @@ class Writer(GenericWriter):
if self.io._fo.tell() or self.block_count > 0:
self.dump()
self.encoder.write_long(block.num_records)
- self.block_writer(self.encoder, block.bytes_.getvalue())
+ self.block_writer(
+ self.encoder, block.bytes_.getvalue(), self.compression_level
+ )
self.encoder._fo.write(self.sync_marker)
def flush(self):
@@ -438,7 +448,8 @@ class JSONWriter(GenericWriter):
sync_interval=1000 * SYNC_SIZE,
metadata=None,
validator=None,
- sync_marker=None):
+ sync_marker=None,
+ codec_compression_level=None):
GenericWriter.__init__(self, schema, metadata, validator)
self.encoder = fo
@@ -460,7 +471,8 @@ def writer(fo,
sync_interval=1000 * SYNC_SIZE,
metadata=None,
validator=None,
- sync_marker=None):
+ sync_marker=None,
+ codec_compression_level=None):
"""Write records to fo (stream) according to schema
Parameters
@@ -486,6 +498,9 @@ def writer(fo,
sync_marker: bytes, optional
A byte string used as the avro sync marker. If not provided, a random
byte string will be used.
+ codec_compression_level: int, optional
+ Compression level to use with the specified codec (if the codec
+ supports it)
Example::
@@ -548,6 +563,7 @@ def writer(fo,
metadata,
validator,
sync_marker,
+ codec_compression_level,
)
for record in records:

MODIFY
tests/test_compression.py
tests/test_compression.py
@@ -129,3 +129,34 @@ def test_unsupported_codec():
with pytest.raises(ValueError, match="Unrecognized codec"):
list(fastavro.reader(modified_file))
+
+
+def test_compression_level():
+ """https://github.com/fastavro/fastavro/issues/377"""
+ schema = {
+ "doc": "A weather reading.",
+ "name": "Weather",
+ "namespace": "test",
+ "type": "record",
+ "fields": [
+ {"name": "station", "type": "string"},
+ {"name": "time", "type": "long"},
+ {"name": "temp", "type": "int"},
+ ],
+ }
+
+ records = [
+ {"station": "011990-99999", "temp": 0, "time": 1433269388},
+ {"station": "011990-99999", "temp": 22, "time": 1433270389},
+ {"station": "011990-99999", "temp": -11, "time": 1433273379},
+ {"station": "012650-99999", "temp": 111, "time": 1433275478},
+ ]
+
+ file = MemoryIO()
+ fastavro.writer(
+ file, schema, records, codec="deflate", codec_compression_level=9
+ )
+
+ file.seek(0)
+ out_records = list(fastavro.reader(file))
+ assert records == out_records

add support for a codec compression level (used only by deflate currently)