astronomy-commons / astronomy-commons/lsdb
Excessive I/O in LSDB/Dask compared to pyarrow.dataset
- Dominant language
- Python
- Stars
- 55
- Forks
- 26
- Avg merge
- 4d 1h
- Merged PRs (30d)
- 8
Description
While testing the high-speed link between `shire` (storage) and `arnor` compute at UW, we found a dramatic difference in I/O between LSDB/Dask and pyarrow.dataset. We wanted to fully load the storage link, so we made tests with low compute and high I/O.
The operation we chose was the calculation of the mean of a single column (`phot_g_mean_mag`) from Gaia DR3. @mtauraso used `strace` to measure the storage usage, and collected and plotted the stats below.
The graphs below show IOPS (left) and bytes read (right), as histograms with logarithmic X-axes. Legend:
- The green line represents the ideal or minimum amount of operations/bytes required to read the column out of the catalog.
- The blue line represents the operations/bytes used by `pyarrow.dataset`.
- the orange line represents the operations/bytes used by LSDB/Dask.
The code for LSDB/Dask looked like:
```python
def catalog_mean(df, pixel, target_column=''):
mean_val = df[target_column].mean()
result = pd.DataFrame([{
"pixel": pixel,
f"{target_column}_mean": mean_val,
}])
return result
# ...
lsdb_gaia = lsdb.open_catalog(gaia_root, columns=['phot_g_mean_mag'])
unrealized = lsdb_gaia.map_partitions(
catalog_mean,
include_pixel=True,
# Keyword arguments after 'include_pixel=' are passed to your
# function
target_column="phot_g_mean_mag",
)
result = unrealized.compute()
# We need to do a final reduction step to get the true mean
float(result['phot_g_mean_mag_mean'].mean())
```
The code for pyarrow.dataset looked like:
```python
parquet_root = f"{gaia_root}/dataset"
pyarrow_ds = ds.parquet_dataset(f"{parquet_root}/_metadata")
def dataset_mean(dataset, field: str, *, use_threads: bool = True):
total_sum = None # Arrow Scalar
total_count = 0 # Python int
for batch in dataset.to_batches(columns=[field], use_threads=use_threads):
col = batch.column(0)
b_sum = pc.sum(col) # Scalar (or null if all-null)
b_count = pc.count(col, mode="only_valid") # Int64 Scalar
if not pc.is_null(b_sum).as_py() and b_count.as_py() > 0:
total_sum = b_sum if total_sum is None else pc.add(total_sum, b_sum)
total_count += b_count.as_py()
if total_sum is None or total_count == 0:
return None
# Avoid pc.divide to steer clear of Expression mixing; just unwrap to Python
return pc.cast(total_sum, pa.float64()).as_py() / float(total_count)
dataset_mean(pyarrow_ds, "phot_g_mean_mag")
```
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.