fix zstd zrle
1 changed file+28−6
benchcompress/src/benchcompress/algorithms/zstd/__init__.pymodified+28−6View file
@@ -104,6 +104,16 @@ def zstd_markov_zrle_encode(x: np.ndarray, level: int) -> bytes:
104104 # Get run lengths for zero/non-zero sequences
105105 run_lengths = get_run_lengths(x)
106106
107+ # Determine run length dtype code
108+ if run_lengths.dtype == np.uint8:
109+ run_length_dtype_code = 0
110+ elif run_lengths.dtype == np.uint16:
111+ run_length_dtype_code = 1
112+ elif run_lengths.dtype == np.uint32:
113+ run_length_dtype_code = 2
114+ else:
115+ raise ValueError(f"Unsupported run length dtype: {run_lengths.dtype}")
116+
107117 # Extract non-zero data
108118 non_zero_arrays = []
109119 array_pos = 0
@@ -130,13 +140,14 @@ def zstd_markov_zrle_encode(x: np.ndarray, level: int) -> bytes:
130140 initial_bytes = initial.tobytes()
131141 run_lengths_bytes = run_lengths.tobytes()
132142
133- # Create header with lengths
143+ # Create header with lengths and dtype code
134144 header = struct.pack(
135- "QQQQ",
145+ "QQQQB",
136146 len(coeffs_bytes),
137147 len(initial_bytes),
138148 len(run_lengths_bytes),
139149 len(run_lengths),
150+ run_length_dtype_code,
140151 )
141152
142153 # Compress residuals
@@ -153,9 +164,9 @@ def zstd_markov_zrle_decode(x: bytes, dtype: str) -> np.ndarray:
153164 import struct
154165
155166 # Extract header
156- header_size = struct.calcsize("QQQQ")
157- coeffs_len, initial_len, run_lengths_len, num_run_lengths = struct.unpack(
158- "QQQQ", x[:header_size]
167+ header_size = struct.calcsize("QQQQB")
168+ coeffs_len, initial_len, run_lengths_len, num_run_lengths, run_length_dtype_code = (
169+ struct.unpack("QQQQB", x[:header_size])
159170 )
160171
161172 # Extract components
@@ -164,9 +175,20 @@ def zstd_markov_zrle_decode(x: bytes, dtype: str) -> np.ndarray:
164175 pos += coeffs_len
165176 initial = np.frombuffer(x[pos : pos + initial_len], dtype=dtype)
166177 pos += initial_len
167- run_lengths = np.frombuffer(x[pos : pos + run_lengths_len], dtype=np.uint32)
178+
179+ # Get run lengths with proper dtype
180+ if run_length_dtype_code == 0:
181+ run_lengths = np.frombuffer(x[pos : pos + run_lengths_len], dtype=np.uint8)
182+ elif run_length_dtype_code == 1:
183+ run_lengths = np.frombuffer(x[pos : pos + run_lengths_len], dtype=np.uint16)
184+ elif run_length_dtype_code == 2:
185+ run_lengths = np.frombuffer(x[pos : pos + run_lengths_len], dtype=np.uint32)
186+ else:
187+ raise ValueError(f"Unsupported run length dtype code: {run_length_dtype_code}")
168188 pos += run_lengths_len
169189
190+ assert len(run_lengths) == num_run_lengths
191+
170192 # Decompress residuals
171193 decompressor = zstd.ZstdDecompressor()
172194 resid_buf = decompressor.decompress(x[pos:])