def _encode_sample(self, idxes: List[int]) -> SampleBatchType:
"""Fetches concatenated samples at given indeces from the storage."""
samples = [self._storage[i] for i in idxes]
if samples:
# Assume all samples are of same type
sample_type = type(samples[0])
out = sample_type.concat_samples(samples)
else:
out = SampleBatch()
out.decompress_if_needed()
return out
def test_compression(self):
"""Tests, whether compression and decompression work properly."""
s1 = SampleBatch({
"a": np.array([1, 2, 3, 2, 3, 4]),
"b": {
"c": np.array([4, 5, 6, 5, 6, 7])
},
})
# Test, whether compressing happens in-place.
s1.compress(columns={"a", "b"}, bulk=True)
self.assertTrue(is_compressed(s1["a"]))
self.assertTrue(is_compressed(s1["b"]["c"]))
self.assertTrue(isinstance(s1["b"], dict))
# Test, whether de-compressing happens in-place.
s1.decompress_if_needed(columns={"a", "b"})
check(s1["a"], [1, 2, 3, 2, 3, 4])
check(s1["b"]["c"], [4, 5, 6, 5, 6, 7])
it = s1.rows()
next(it)
check(next(it), {"a": 2, "b": {"c": 5}})
