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