89 lines
2.8 KiB
Python
89 lines
2.8 KiB
Python
from io import BytesIO
|
|
|
|
try:
|
|
from clickhouse_cityhash.cityhash import CityHash128
|
|
except ImportError:
|
|
raise RuntimeError(
|
|
'Package clickhouse-cityhash is required to use compression'
|
|
)
|
|
|
|
from .native import BlockOutputStream, BlockInputStream
|
|
from ..bufferedreader import CompressedBufferedReader
|
|
from ..bufferedwriter import CompressedBufferedWriter
|
|
from ..compression import get_decompressor_cls
|
|
from ..defines import BUFFER_SIZE
|
|
from ..reader import read_binary_uint8, read_binary_uint128
|
|
from ..writer import write_binary_uint8, write_binary_uint128
|
|
|
|
|
|
class CompressedBlockOutputStream(BlockOutputStream):
|
|
def __init__(self, compressor_cls, compress_block_size, fout, context):
|
|
self.compressor_cls = compressor_cls
|
|
self.compress_block_size = compress_block_size
|
|
self.raw_fout = fout
|
|
|
|
self.compressor = self.compressor_cls()
|
|
self.fout = CompressedBufferedWriter(self.compressor, BUFFER_SIZE)
|
|
super(CompressedBlockOutputStream, self).__init__(self.fout, context)
|
|
|
|
def get_compressed_hash(self, data):
|
|
return CityHash128(data)
|
|
|
|
def finalize(self):
|
|
self.fout.flush()
|
|
|
|
compressed = self.get_compressed()
|
|
compressed_size = len(compressed)
|
|
|
|
compressed_hash = self.get_compressed_hash(compressed)
|
|
write_binary_uint128(compressed_hash, self.raw_fout)
|
|
|
|
block_size = self.compress_block_size
|
|
|
|
i = 0
|
|
while i < compressed_size:
|
|
self.raw_fout.write(compressed[i:i + block_size])
|
|
i += block_size
|
|
|
|
self.raw_fout.flush()
|
|
|
|
def get_compressed(self):
|
|
compressed = BytesIO()
|
|
|
|
if self.compressor.method_byte is not None:
|
|
write_binary_uint8(self.compressor.method_byte, compressed)
|
|
extra_header_size = 1 # method
|
|
else:
|
|
extra_header_size = 0
|
|
|
|
data = self.compressor.get_compressed_data(extra_header_size)
|
|
compressed.write(data)
|
|
|
|
return compressed.getvalue()
|
|
|
|
|
|
class CompressedBlockInputStream(BlockInputStream):
|
|
def __init__(self, fin, context):
|
|
self.raw_fin = fin
|
|
fin = CompressedBufferedReader(self.read_block, BUFFER_SIZE)
|
|
super(CompressedBlockInputStream, self).__init__(fin, context)
|
|
|
|
def get_compressed_hash(self, data):
|
|
return CityHash128(data)
|
|
|
|
def read_block(self):
|
|
compressed_hash = read_binary_uint128(self.raw_fin)
|
|
method_byte = read_binary_uint8(self.raw_fin)
|
|
|
|
decompressor_cls = get_decompressor_cls(method_byte)
|
|
decompressor = decompressor_cls(self.raw_fin)
|
|
|
|
if decompressor.method_byte is not None:
|
|
extra_header_size = 1 # method
|
|
else:
|
|
extra_header_size = 0
|
|
|
|
return decompressor.get_decompressed_data(
|
|
method_byte, compressed_hash, extra_header_size
|
|
)
|