def test_rows(self):
s1 = SampleBatch(
{
"a": np.array([[1, 1], [2, 2], [3, 3]]),
"b": {"c": np.array([[4, 4], [5, 5], [6, 6]])},
SampleBatch.SEQ_LENS: np.array([1, 2]),
}
)
check(
next(s1.rows()),
{"a": [1, 1], "b": {"c": [4, 4]}, SampleBatch.SEQ_LENS: 1},
)
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}})
