Skip to content

Apply the indices mapping in Dataset.unique - #8383

Open
LeSingh1 wants to merge 1 commit into
huggingface:mainfrom
LeSingh1:fix/unique-indices-mapping
Open

Apply the indices mapping in Dataset.unique#8383
LeSingh1 wants to merge 1 commit into
huggingface:mainfrom
LeSingh1:fix/unique-indices-mapping

Conversation

@LeSingh1

@LeSingh1 LeSingh1 commented Aug 4, 2026

Copy link
Copy Markdown

Dataset.unique only takes the indices mapping into account when it has a different number of rows than the underlying table:

if self._indices is not None and self._indices.num_rows != self._data.num_rows:

An indices mapping can also repeat some rows and drop others while keeping the same length. In that case unique reads the raw table and returns values that are not in the dataset:

ds = Dataset.from_dict({"col_1": ["a", "b", "c"]}).select([0, 0, 1])
ds["col_1"]         # ['a', 'a', 'b']
ds.unique("col_1")  # ['a', 'b', 'c']  <- 'c' is not in the dataset

class_encode_column builds its ClassLabel names from unique, so it inherits the bug and creates a class for rows that were dropped:

ds.class_encode_column("col_1").features["col_1"].names  # ['a', 'b', 'c']

This is reachable from ordinary use — select with repeats (bootstrap resampling), and any indices mapping whose length happens to match the table.

This applies the indices mapping whenever there is one. Only the requested column is materialized, via the existing query_table(..., key=column, indices=self._indices) path, so it deliberately avoids falling back to flatten_indices() — that would make shuffle().unique() rewrite the whole dataset.

Added test_unique_with_indices_mapping_of_same_length and test_class_encode_column_with_indices_mapping_of_same_length; both fail on main with assert ['a', 'b', 'c'] == ['a', 'b'] and pass with the change. tests/test_arrow_dataset.py, tests/test_dataset_dict.py and tests/test_search.py pass (470 passed, 38 skipped).

…th as the table

Dataset.unique only took the indices mapping into account when it had a
different number of rows than the underlying table. An indices mapping can
also repeat some rows and drop others while keeping the same length, e.g.
after ds.select([0, 0, 1]) or a bootstrap resampling, and in that case unique
was reading the raw table and reported values that are not in the dataset.

class_encode_column builds its ClassLabel names from unique, so it inherited
the same problem and created classes for rows that were filtered out.

Apply the indices mapping whenever there is one. Only the requested column is
materialized, so this doesn't fall back to flattening the whole dataset.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant