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}})