Function test_concatenate_datasets

tests/test_arrow_dataset.py:3644–3658

(dataset_type, axis, expected_shape, dataset_dict, arrow_path)

@pytest.mark.parametrize("dataset_type", ["in_memory", "memory_mapped", "mixed"])
@pytest.mark.parametrize("axis, expected_shape", [(0, (4, 3)), (1, (2, 6))])
def test_concatenate_datasets(dataset_type, axis, expected_shape, dataset_dict, arrow_path):
    table = {
        "in_memory": InMemoryTable.from_pydict(dataset_dict),
        "memory_mapped": MemoryMappedTable.from_file(arrow_path),
    }
    tables = [
        table[dataset_type if dataset_type != "mixed" else "memory_mapped"].slice(0, 2),  # shape = (2, 3)
        table[dataset_type if dataset_type != "mixed" else "in_memory"].slice(2, 4),  # shape = (2, 3)
    ]
    if axis == 1:  # don't duplicate columns
        tables[1] = tables[1].rename_columns([col + "_bis" for col in tables[1].column_names])
    datasets = [Dataset(table) for table in tables]
    dataset = concatenate_datasets(datasets, axis=axis)
    assert dataset.shape == expected_shape
    assert_arrow_metadata_are_synced_with_dataset_features(dataset)


def test_concatenate_datasets_new_columns():