Function generate_sample_data_csv

python/ray/data/tests/test_checkpoint.py:67–85

(tmp_path)

@pytest.fixture
def generate_sample_data_csv(tmp_path):
    def _generate(id_type: Literal["int", "str"] = "int"):
        # Generate a dummy dataset with SAMPLE_DATA_NUM_ROWS rows and columns [ID_COL, "col1"]
        ids = (
            [f"id_{i}" for i in range(SAMPLE_DATA_NUM_ROWS)]
            if id_type == "str"
            else list(range(SAMPLE_DATA_NUM_ROWS))
        )

        data = [{ID_COL: id_val, "col1": random.random()} for id_val in ids]

        f_path = os.path.join(tmp_path, "sample_data.csv")
        with open(f_path, mode="w", newline="") as file:
            writer = csv.DictWriter(file, fieldnames=data[0].keys())
            writer.writeheader()
            writer.writerows(data)
        return f_path

    return _generate


@pytest.fixture