def test_roundtrip_rt_pa_rt(self, rt_dset: rt.Dataset) -> None:
"""Test round-tripping from rt.Dataset to pyarrow.Table and back."""
result_pa_tbl = rt_dset.to_arrow()
result_rt_dset = rt.Dataset.from_arrow(result_pa_tbl, zero_copy_only=False)
assert rt_dset.keys() == result_rt_dset.keys()
for col_name in rt_dset.keys():
# relaxed_cat_check=True, because we're not trying to test specific details of Categorical conversion
# here, we're more interested in the dataset-level stuff.
assert_array_or_cat_equal(rt_dset[col_name], result_rt_dset[col_name], relaxed_cat_check=True)
def test_append_dataset_dict(self):
ds = Dataset({
_k: list(range(_i * 10, (_i + 1) * 10))
for _i, _k in enumerate('abcdefghijklmnop')
})
ds0 = ds[:3].copy()
ds1 = ds[6:9].copy()
ds2 = ds[11:15].copy()
dd = {'one': ds0, 'two': ds1, 'three': ds2}
ds = append_dataset_dict(dd, 'keyfield')
ucols = set()
for _d in dd.values():
ucols.update(_d)
self.assertEqual(set(ds.keys()), ucols)
self.assertEqual(ds.get_nrows(),
sum(_d.get_nrows() for _d in dd.values()))
keyfield = []
for _d in dd.values():
keyfield.extend(_d.keyfield)
self.assertTrue((ds.keyfield == keyfield).all())
self.assertTrue((ds.a[:10] == range(10)).all())
self.assertTrue((ds.g[10:20] == range(60, 70)).all())
self.assertTrue((ds.l[20:30] == range(110, 120)).all())
