GeezerVison is showing github.com — feast-benchmarks/data_generator.py at main · feast-dev/feast-benchmarks

Expand file tree

data_generator.py

Copy path

More file actions

19 lines (17 loc) · 712 Bytes

Copy path

19 lines (17 loc) · 712 Bytes

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

import datetime

import pyarrow as pa

import pyarrow.parquet as pq

import pandas as pd

import numpy as np

from pathlib import Path

def generate_data(num_rows: int, num_features: int, key_space: int) -> pd.DataFrame:

features = [f"feature_{i}" for i in range(num_features)]

columns = ["entity", "event_timestamp"] + features

df = pd.DataFrame(0, index=np.arange(num_rows), columns=columns)

df["event_timestamp"] = datetime.datetime.utcnow()

for column in ["entity"] + features:

df[column] = np.random.randint(1, key_space, num_rows)

return df

if __name__ == "__main__":

df = generate_data(10**4, 250, 10**4)

df.to_parquet(Path(__file__).parent / "generated_data.parquet")