astronomy-commons / astronomy-commons/lsdb

Excessive I/O in LSDB/Dask compared to pyarrow.dataset

Open
#1,010 2 comments 0 reactions 1 assignee Assigned to @smcguire-cmu View on GitHub
performance
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.

Image

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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.