Intelligent sampling for large datasets

[jtalmi]

Lilac currently does not work well with large datasets living in cloud storage. We often need to inspect samples of multi-GB/TB, partitioned parquet datasets but we aren't able to load the entires dataset locally. Ideally this would mean sampling a few rows from each shard, but a simpler version would be taking a single shard.

If Lilac adds better support for sampling, it could be used to share/inspect/evaluate large pre-training datasets

[nsthorat]

We've definitely been focusing on the single-node solution but we'll start thinking about larger datasets soon.

How big is one of your shards and in what format? We can currently read just a single shard (e.g. if it's parquet) if you use the parquet source. We have sampling for the HuggingFace source, but we can also add a sampling mechanism for arbitrary data.

[jtalmi]

We try to keep shards small...maybe <200 MB. The main issue with taking a single shard is skewness, if the dataset isn't properly shuffled. An ideal solution would sample across all shards of a parquet dataset, but I don't know off-hand how feasible that is.

[dsmilkov]

Is your data in HF datasets format, or just raw parquet files? Also we are in slack room with HuggingFace so I just asked for their advice on how to best load a single shard, and ideally a random sample of many shards. Will follow up shortly.

[dsmilkov]

Also I'm assuming the s3 buckets are private? We are using duckdb for really fast parquet reads and will use it for sampling. They support S3 reads directly: https://duckdb.org/docs/guides/import/s3_import.html

You will just need to set some env variables to authorize with S3:

S3_REGION=
S3_ENDPOINT=
S3_ACCESS_KEY=
S3_SECRET_KEY=  

Let me know if this would work for you

[jtalmi]

We are using GCP actually! And we have our own custom datasets in parquet format usually (some .jsonl, .csv, but parquet the main one).

Does duckdb enable smart sampling parquet files within s3/gcs buckets, so that the sample is taken from many shards?

[dsmilkov]

I just sent #778 for review. In summary, when you have many shards, and you don't want to fully shuffle a dataset, we fetch a few rows from each shard. This is done by creating a batched reader for each shard and flipping a coin to select the next reader.

while ...:
  index = random.randint(0, len(readers) - 1)
  reader = readers[index]
  batch = reader.read_next_batch() # Fetch a batch of rows.

I tested this on 4 shards. How many shards does your large dataset have ? 100s, 1,000s, 10,000s? Because the current implementation will open a reader for each shard, and not sure how the OS will handle having many readers open at once.