Preparing a Dataset

This document walks through the full process of preparing a new non-synthetic dataset for benchmarking. The high-level steps are:

  1. Load source data into S3
  2. Create a dataset config
  3. Sample the data at each size you need
  4. Convert the sampled parquet data to CSV
  5. Load the heap into Postgres
  6. Snapshot or restore the heap with pgBackRest
  7. Run the benchmark

Unless otherwise noted, run commands from the benchmarks/ directory.

Step 1: Load Source Data into S3

Upload your source data as partitioned parquet files to S3. Each table should be in its own subdirectory under a source/parquet/ path (filename doesn’t matter):

s3://<dataset-bucket>/datasets/{dataset-name}/source/parquet/
├── {table_a}/
│   ├── part-0001.parquet
│   ├── part-0002.parquet
│   └── ...
├── {table_b}/
│   └── ...
└── {table_c}/
    └── ...

For example, for the Stack Overflow dataset:

s3://<dataset-bucket>/datasets/stackoverflow/source/parquet/
├── stackoverflow_posts/
├── comments/
└── users/

Step 2: Create a Config

Writing the Config

Create a TOML config file at datasets/{dataset-name}/config.toml that describes your table relationships. The config specifies which table is the root (the one that gets sampled directly) and how child tables relate to it via joins.

sampling_seed = 723
s3_base_path = "s3://<dataset-bucket>/datasets/root_dataset"

[root_table]
name = "root_table_name"
primary_key = "id"

[[tables]]
name = "child_table"
parent = "root_table_name"
parent_join_col = "id"
join_col = "parent_id"

Fields:

  • sampling_seed: Seed for deterministic, reproducible sampling.
  • s3_base_path: Optional default source for load-heap. It can be overridden with --data-source.
  • [root_table]: The primary table. The --rows argument controls how many rows are sampled from this table.
  • primary_key: Unique, non-null key used for deterministic root-table sampling.
  • [[tables]]: One entry per non-root table. Child tables specify parent, parent_join_col, and join_col to define the relationship.

Child tables are not sampled independently. Instead, they are filtered via an inner join with their parent table, so only rows that reference a sampled parent row are kept. This preserves referential integrity across the dataset.

Tables can form a hierarchy (a child can be a parent of another table). They are processed in topological order.

See datasets/stackoverflow/config.toml for a real example.

Step 3: Sample the Source Data

Running the Sampling Tool

Run the sample command once for each dataset size you need. The --rows argument sets the target row count for the root table. Output goes to the sampled/{size}/parquet/ path.

# Sample to 10k rows
cargo run --release -- sample \
  --input s3://<dataset-bucket>/datasets/stackoverflow/source/parquet/ \
  --output s3://<dataset-bucket>/datasets/stackoverflow/sampled/10k/parquet/ \
  --config ./datasets/stackoverflow/config.toml \
  --rows 10000

# Sample to 100k rows
cargo run --release -- sample \
  --input s3://<dataset-bucket>/datasets/stackoverflow/source/parquet/ \
  --output s3://<dataset-bucket>/datasets/stackoverflow/sampled/100k/parquet/ \
  --config ./datasets/stackoverflow/config.toml \
  --rows 100000

# Sample to 1m rows
cargo run --release -- sample \
  --input s3://<dataset-bucket>/datasets/stackoverflow/source/parquet/ \
  --output s3://<dataset-bucket>/datasets/stackoverflow/sampled/1m/parquet/ \
  --config ./datasets/stackoverflow/config.toml \
  --rows 1000000

Notes:

  • The output path must be empty (no pre-existing data).
  • For small targets (<=100k rows), sampling is exact using reservoir sampling. For larger targets, it uses system sampling and the result will be approximate (within ~3-5%).
  • Use --dry-run to validate inputs and see planned row counts without writing anything.

Step 4: Convert Sampled Data to CSV

Run the convert command for each sampled size to produce CSV versions. The --tables flag takes a comma-separated list of all tables to convert.

# Convert 10k sampled data
cargo run --release -- convert \
  --input s3://<dataset-bucket>/datasets/stackoverflow/sampled/10k/parquet/ \
  --output s3://<dataset-bucket>/datasets/stackoverflow/sampled/10k/csv/ \
  --tables stackoverflow_posts,comments,users

# Convert 100k sampled data
cargo run --release -- convert \
  --input s3://<dataset-bucket>/datasets/stackoverflow/sampled/100k/parquet/ \
  --output s3://<dataset-bucket>/datasets/stackoverflow/sampled/100k/csv/ \
  --tables stackoverflow_posts,comments,users

# Convert 1m sampled data
cargo run --release -- convert \
  --input s3://<dataset-bucket>/datasets/stackoverflow/sampled/1m/parquet/ \
  --output s3://<dataset-bucket>/datasets/stackoverflow/sampled/1m/csv/ \
  --tables stackoverflow_posts,comments,users

You can also convert the full source data:

cargo run --release -- convert \
  --input s3://<dataset-bucket>/datasets/stackoverflow/source/parquet/ \
  --output s3://<dataset-bucket>/datasets/stackoverflow/source/csv/ \
  --tables stackoverflow_posts,comments,users

Notes:

  • The output path must be empty.
  • Row counts are verified after conversion to ensure no data is lost.
  • Use --dry-run to validate without writing.
  • AWS credentials must be accessible via the standard credential chain (env vars, ~/.aws/credentials, or instance metadata).

Step 5: Load the Heap

Load the sampled CSV data into Postgres with load-heap:

POSTGRES_URL="postgresql://localhost:28818/postgres"

cargo run --release -- load-heap \
  --url "${POSTGRES_URL}" \
  --dataset stackoverflow \
  --size 100k

load-heap reads from {s3_base_path}/sampled/{size}/csv/{table}/, creates tables from datasets/{dataset}/create_tables.sql, and bulk-loads CSV files. Use --data-source to override s3_base_path from the dataset config.

Step 6: Snapshot or Restore the Heap

The snapshot commands shell out to pgbackrest, so install and configure pgBackRest before using them. Use --config to provide your own pgBackRest config, or pass repository settings to generate one. The examples below use ParadeDB’s CI snapshot repository: s3://paradedb-ci-benchmarks/snapshots/{dataset}/{size}/.

Stop Postgres before snapshotting a loaded heap:

PGDATA="/home/runner/.pgrx/data-18"
BACKREST_ARGS=(
  --stanza bench
  --repo-bucket paradedb-ci-benchmarks
  --repo-path-prefix /snapshots
  --repo-region us-east-1
  --repo-endpoint s3.us-east-1.amazonaws.com
  --repo-s3-key-type shared
)

(cd ../pg_search && cargo pgrx stop pg18)
cargo run --release -- snapshot-heap \
  --dataset stackoverflow \
  --size 100k \
  --pgdata "${PGDATA}" \
  "${BACKREST_ARGS[@]}"

To restore that heap later, keep Postgres stopped, restore the snapshot, and start Postgres:

cargo run --release -- restore-heap \
  --dataset stackoverflow \
  --size 100k \
  --pgdata "${PGDATA}" \
  "${BACKREST_ARGS[@]}"
(cd ../pg_search && cargo pgrx start pg18)

Useful snapshot options:

  • --repo-bucket: S3 bucket for the pgBackRest repository.
  • --repo-path-prefix: Prefix inside the bucket.
  • --repo-s3-key-type: Use shared, auto, or web-id.
  • --process-max: Max pgBackRest worker processes. Defaults to available CPUs.

Step 7: Run Benchmarks

After loading or restoring the heap, run benchmark:

cargo run --release -- benchmark \
  --url "${POSTGRES_URL}" \
  --dataset stackoverflow \
  --index bm25

Final S3 Layout

After preparing data and publishing snapshots, your artifacts will look like this:

s3://<dataset-bucket>/datasets/{dataset-name}/
├── source/
│   ├── parquet/
│   │   ├── {table_a}/
│   │   ├── {table_b}/
│   │   └── {table_c}/
│   └── csv/
│       ├── {table_a}/
│       ├── {table_b}/
│       └── {table_c}/
└── sampled/
    ├── 10k/
    │   ├── parquet/
    │   │   ├── {table_a}/
    │   │   ├── {table_b}/
    │   │   └── {table_c}/
    │   └── csv/
    │       ├── {table_a}/
    │       ├── {table_b}/
    │       └── {table_c}/
    ├── 100k/
    │   └── ...
    └── 1m/
        └── ...

s3://<snapshot-bucket>/snapshots/{dataset-name}/
├── 100k/
├── 1m/
└── 20m/