Function _check_data

tests/dataset/test_local_datasets.py:181–225

(ms_dataset: Dataset | DatasetDict,
                oc_dataset: Dataset | DatasetDict,
                sample_size)

def _check_data(ms_dataset: Dataset | DatasetDict,
                oc_dataset: Dataset | DatasetDict,
                sample_size):
    assert type(ms_dataset) == type(
        oc_dataset
    ), f'Dataset type not match: {type(ms_dataset)} != {type(oc_dataset)}'

    # match DatasetDict
    if isinstance(oc_dataset, DatasetDict):
        assert ms_dataset.keys() == oc_dataset.keys(
        ), f'DatasetDict not match: {ms_dataset.keys()} != {oc_dataset.keys()}'

        for key in ms_dataset.keys():
            _check_data(ms_dataset[key], oc_dataset[key], sample_size=sample_size)

    elif isinstance(oc_dataset, Dataset):
        # match by cols
        assert set(ms_dataset.column_names) == set(
            oc_dataset.column_names
        ), f'Column names do not match: {ms_dataset.column_names} != {oc_dataset.column_names}'

        # Check that the number of rows is the same
        assert len(ms_dataset) == len(
            oc_dataset
        ), f'Number of rows do not match: {len(ms_dataset)} != {len(oc_dataset)}'

        # Randomly sample indices
        sample_indices = random.sample(range(len(ms_dataset)),
                                       min(sample_size, len(ms_dataset)))

        for i, idx in enumerate(sample_indices):
            for col in ms_dataset.column_names:
                ms_value = clean_string(str(ms_dataset[col][idx]))
                oc_value = clean_string(str(oc_dataset[col][idx]))
                try:
                    assert ms_value == oc_value, f"Value mismatch in column '{col}', index {idx}: {ms_value} != {oc_value}"
                except AssertionError as e:
                    print(f"Assertion failed for column '{col}', index {idx}")
                    print(f"ms_data: {ms_dataset[idx]}")
                    print(f'oc_data: {oc_dataset[idx]}')
                    print(f'ms_value: {ms_value} ({type(ms_value)})')
                    print(f'oc_value: {oc_value} ({type(oc_value)})')
                    raise e
    else:
        raise ValueError(f'Datasets type not supported {type(ms_dataset)}')


if __name__ == '__main__':