wehub-resource-sync
103 lines
3.0 KiB
Python
103 lines
3.0 KiB
Python
import pyarrow as pa
|
|
import uuid
|
|
import random
|
|
import string
|
|
import ray
|
|
import pyarrow.parquet as pq
|
|
from tqdm import tqdm
|
|
|
|
STRING_PLACEHOLDER = ""
|
|
UUID_PLACEHOLDER = uuid.UUID(int=0)
|
|
INT_PLACEHOLDER = 0
|
|
|
|
TARGET_SIZE_BYTES = 4096
|
|
NUM_FILES = 50
|
|
|
|
SCHEMA = pa.schema(
|
|
[
|
|
("metadata00", pa.string()),
|
|
("metadata01", pa.list_(pa.binary(16))),
|
|
("metadata02", pa.string()),
|
|
("metadata03", pa.uint64()),
|
|
("metadata04", pa.list_(pa.binary(16))),
|
|
("metadata05", pa.list_(pa.binary(16))),
|
|
("metadata06", pa.binary(16)),
|
|
("metadata07", pa.string()),
|
|
("metadata08", pa.binary(16)),
|
|
("metadata09", pa.uint64()),
|
|
("metadata10", pa.binary(16)),
|
|
("metadata11", pa.list_(pa.binary(16))),
|
|
("metadata12", pa.uint64()),
|
|
("metadata13", pa.uint64()),
|
|
("metadata14", pa.list_(pa.binary(16))),
|
|
("span_text", pa.string()),
|
|
("metadata15", pa.binary(16)),
|
|
("metadata16", pa.string()),
|
|
("metadata17", pa.list_(pa.binary(16))),
|
|
("metadata18", pa.list_(pa.binary(16))),
|
|
]
|
|
)
|
|
|
|
|
|
def random_word(min_len=3, max_len=8):
|
|
length = random.randint(min_len, max_len)
|
|
return "".join(random.choices(string.ascii_lowercase, k=length))
|
|
|
|
|
|
def create_random_sentence():
|
|
sentence = ""
|
|
while len(sentence.encode("utf-8")) < TARGET_SIZE_BYTES:
|
|
word = random_word()
|
|
sentence += word + " " # space between words
|
|
|
|
# Trim to exact size
|
|
sentence_bytes = sentence.encode("utf-8")[:TARGET_SIZE_BYTES]
|
|
return sentence_bytes.decode("utf-8", errors="ignore")
|
|
|
|
|
|
def create_row():
|
|
return {
|
|
"metadata00": STRING_PLACEHOLDER,
|
|
"metadata01": [UUID_PLACEHOLDER.bytes],
|
|
"metadata02": STRING_PLACEHOLDER,
|
|
"metadata03": INT_PLACEHOLDER,
|
|
"metadata04": [UUID_PLACEHOLDER.bytes],
|
|
"metadata05": [UUID_PLACEHOLDER.bytes],
|
|
"metadata06": UUID_PLACEHOLDER.bytes,
|
|
"metadata07": STRING_PLACEHOLDER,
|
|
"metadata08": UUID_PLACEHOLDER.bytes,
|
|
"metadata09": INT_PLACEHOLDER,
|
|
"metadata10": UUID_PLACEHOLDER.bytes,
|
|
"metadata11": [UUID_PLACEHOLDER.bytes],
|
|
"metadata12": INT_PLACEHOLDER,
|
|
"metadata13": None if random.random() < 0.01 else INT_PLACEHOLDER,
|
|
"metadata14": [UUID_PLACEHOLDER.bytes],
|
|
"span_text": create_random_sentence(),
|
|
"metadata15": UUID_PLACEHOLDER.bytes,
|
|
"metadata16": STRING_PLACEHOLDER,
|
|
"metadata17": [UUID_PLACEHOLDER.bytes],
|
|
"metadata18": [UUID_PLACEHOLDER.bytes],
|
|
}
|
|
|
|
|
|
@ray.remote
|
|
def write_table(i: int):
|
|
rows = []
|
|
for _ in range(20_000):
|
|
rows.append(create_row())
|
|
|
|
table = pa.Table.from_pylist(rows, schema=SCHEMA)
|
|
pq.write_table(
|
|
table, f"s3://ray-benchmark-data-internal-us-west-2/text-spans/{i}.parquet"
|
|
)
|
|
|
|
|
|
refs = [write_table.remote(i) for i in range(NUM_FILES)]
|
|
|
|
pbar = tqdm(total=len(refs))
|
|
while refs:
|
|
ready, refs = ray.wait(refs, num_returns=1)
|
|
pbar.update(len(ready))
|
|
|
|
pbar.close()
|