Coverage for src/loman/serialization/compression.py: 100%

1"""Blob compression. 

2 

3Compression is a choice the caller makes, named on the profile as a codec and 

4level: ``"zstd:1"``, ``"zlib:6"``, or ``"none"``. Nothing here estimates, 

5samples or guesses. 

6 

7There used to be an ``"auto"`` mode that compressed the first 256 KiB of a blob, 

8extrapolated, and skipped compression when the projection looked poor. It was 

9deleted rather than tuned, because it was wrong in both directions on data whose 

10character changes part way through --- which market data routinely does. On a 

11payload with a random head and a compressible tail it projected a 3.6% saving 

12against an actual 36.8%, and silently stored the blob raw. 

13 

14The only reason to guess was that zlib rejects incompressible data at about 

1543 MB/s, so compressing to find out was expensive. zstd does the same at roughly 

161 GB/s, and compresses real data better and faster besides, which is why it is a 

17required dependency rather than an extra. Compressing to find out now costs 

18about a second per gigabyte of incompressible data, so there is nothing left for 

19a heuristic to save. 

20 

21One rule remains, and it is a measurement rather than a prediction: whichever of 

22the compressed and raw payloads is smaller is the one stored. A blob that did 

23not shrink is written as-is and recorded as ``"none"``, so no future read pays a 

24decompression step for nothing. 

25""" 

26 

27from __future__ import annotations 

28 

29import zlib 

30from collections.abc import Callable 

31 

32from loman.exception import SerializationError 

33 

34Codec = tuple[Callable[[bytes], bytes], Callable[[bytes], bytes]] 

35 

36 

37def _zlib_codec(level: int) -> Codec: 

38 """Return the compress/decompress pair for zlib at *level*.""" 

39 return (lambda data: zlib.compress(data, level), zlib.decompress) 

40 

41 

42def _zstd_codec(level: int) -> Codec: 

43 """Return the compress/decompress pair for zstd at *level*.""" 

44 import zstandard 

45 

46 def compress(data: bytes) -> bytes: 

47 return zstandard.ZstdCompressor(level=level).compress(data) 

48 

49 def decompress(data: bytes) -> bytes: 

50 return zstandard.ZstdDecompressor().decompress(data) 

51 

52 return compress, decompress 

53 

54 

55_FAMILIES: dict[str, Callable[[int], Codec]] = { 

56 "zlib": _zlib_codec, 

57 "zstd": _zstd_codec, 

58} 

59 

60# Level 1 for zstd: on measured data it gives 9.3x on a realistic price series 

61# at 568 MB/s, and rejects incompressible data at 1067 MB/s. Higher levels cost 

62# materially more for little further gain on numeric payloads. 

63_DEFAULT_LEVELS = {"zlib": 6, "zstd": 1} 

64 

65#: What the efficient profile uses when the caller names nothing. 

66DEFAULT_COMPRESSION = "zstd:1" 

67 

68#: Accepted when compression should not happen at all. 

69NO_COMPRESSION = "none" 

70 

71 

72def parse_spec(spec: str) -> tuple[str, int | None]: 

73 """Split a compression spec such as ``"zlib:1"`` into family and level.""" 

74 family, _, level = spec.partition(":") 

75 if not level: 

76 return family, None 

77 try: 

78 return family, int(level) 

79 except ValueError: 

80 msg = f"Invalid compression level in {spec!r}: expected an integer after ':'" 

81 raise ValueError(msg) from None 

82 

83 

84def get_codec(spec: str) -> Codec: 

85 """Return the compress/decompress pair named by *spec*.""" 

86 family, level = parse_spec(spec) 

87 factory = _FAMILIES.get(family) 

88 if factory is None: 

89 known = sorted([NO_COMPRESSION, *_FAMILIES]) 

90 msg = f"Unknown compression {spec!r}; expected one of {known}" 

91 raise ValueError(msg) 

92 return factory(_DEFAULT_LEVELS[family] if level is None else level) 

93 

94 

95def register_codec(family: str, factory: Callable[[int], Codec], default_level: int) -> None: 

96 """Register a compression family, so a user can bring their own. 

97 

98 :param family: Name used in a compression spec, before the ``:``. 

99 :param factory: Called with a level, returning ``(compress, decompress)``. 

100 :param default_level: Level used when a spec names no level. 

101 """ 

102 _FAMILIES[family] = factory 

103 _DEFAULT_LEVELS[family] = default_level 

104 

105 

106def compress_blob(data: bytes, spec: str, *, compressible: bool = True) -> tuple[bytes, str]: 

107 """Return *data* compressed according to *spec*, and the spec actually used. 

108 

109 The returned spec is ``"none"`` whenever the raw bytes are what got stored, 

110 whether because compression was not asked for or because it did not shrink 

111 them. That is what the manifest records, so a reader never has to know what 

112 was originally requested. 

113 

114 :param data: The payload. 

115 :param spec: ``"none"``, or a family with an optional level. 

116 :param compressible: False for payloads that already compress themselves, 

117 such as parquet. Skips compression entirely, which is how 

118 double-compression is prevented. 

119 """ 

120 if spec == NO_COMPRESSION or not compressible or not data: 

121 return data, NO_COMPRESSION 

122 

123 compress, _ = get_codec(spec) 

124 compressed = compress(data) 

125 

126 # Store whichever is smaller. Incompressible data comes back slightly larger 

127 # than it went in, since a codec adds framing, and keeping that would mean a 

128 # bigger file *and* a decompression step on every future read. 

129 if len(compressed) >= len(data): 

130 return data, NO_COMPRESSION 

131 return compressed, spec 

132 

133 

134def decompress_blob(data: bytes, spec: str) -> bytes: 

135 """Return *data* decompressed according to *spec*.""" 

136 if spec in (NO_COMPRESSION, "", None): 

137 return data 

138 try: 

139 _, decompress = get_codec(spec) 

140 except ValueError as exc: 

141 msg = f"Cannot read a blob compressed with {spec!r}: {exc}" 

142 raise SerializationError(msg) from exc 

143 return decompress(data)