Method test_select

tests/test_arrow_dataset.py:2220–2312

(self, in_memory)

                        assert_arrow_metadata_are_synced_with_dataset_features(inverted_dset)

    def test_select(self, in_memory):
        with tempfile.TemporaryDirectory() as tmp_dir:
            with self._create_dummy_dataset(in_memory, tmp_dir) as dset:
                # select every two example
                indices = list(range(0, len(dset), 2))
                tmp_file = os.path.join(tmp_dir, "test.arrow")
                fingerprint = dset._fingerprint
                with dset.select(indices, indices_cache_file_name=tmp_file) as dset_select_even:
                    self.assertIsNotNone(dset_select_even._indices)  # an indices mapping is created
                    self.assertTrue(os.path.exists(tmp_file))
                    self.assertEqual(len(dset_select_even), 15)
                    for row in dset_select_even:
                        self.assertEqual(int(row["filename"][-1]) % 2, 0)
                    self.assertDictEqual(dset.features, Features({"filename": Value("string")}))
                    self.assertDictEqual(dset_select_even.features, Features({"filename": Value("string")}))
                    self.assertNotEqual(dset_select_even._fingerprint, fingerprint)

        with tempfile.TemporaryDirectory() as tmp_dir:
            with self._create_dummy_dataset(in_memory, tmp_dir) as dset:
                indices = list(range(0, len(dset)))
                with dset.select(indices) as dset_select_all:
                    # no indices mapping, since the indices are contiguous
                    # (in this case the arrow table is simply sliced, which is more efficient)
                    self.assertIsNone(dset_select_all._indices)
                    self.assertEqual(len(dset_select_all), len(dset))
                    self.assertListEqual(list(dset_select_all), list(dset))
                    self.assertDictEqual(dset.features, Features({"filename": Value("string")}))
                    self.assertDictEqual(dset_select_all.features, Features({"filename": Value("string")}))
                    self.assertNotEqual(dset_select_all._fingerprint, fingerprint)
                indices = range(0, len(dset))
                with dset.select(indices) as dset_select_all:
                    # same but with range
                    self.assertIsNone(dset_select_all._indices)
                    self.assertEqual(len(dset_select_all), len(dset))
                    self.assertListEqual(list(dset_select_all), list(dset))
                    self.assertDictEqual(dset.features, Features({"filename": Value("string")}))
                    self.assertDictEqual(dset_select_all.features, Features({"filename": Value("string")}))
                    self.assertNotEqual(dset_select_all._fingerprint, fingerprint)

        with tempfile.TemporaryDirectory() as tmp_dir:
            with self._create_dummy_dataset(in_memory, tmp_dir) as dset:
                bad_indices = list(range(5))
                bad_indices[-1] = len(dset) + 10  # out of bounds
                tmp_file = os.path.join(tmp_dir, "test.arrow")
                self.assertRaises(
                    Exception,
                    dset.select,
                    indices=bad_indices,
                    indices_cache_file_name=tmp_file,
                    writer_batch_size=2,
                )
                self.assertFalse(os.path.exists(tmp_file))

        with tempfile.TemporaryDirectory() as tmp_dir:
            with self._create_dummy_dataset(in_memory, tmp_dir) as dset:
                indices = iter(range(len(dset)))  # iterator of contiguous indices
                with dset.select(indices) as dset_select_all:
                    # no indices mapping, since the indices are contiguous