From 9d6660bea18d2e7e7b3a75e58ddc78b4e8e46ec8 Mon Sep 17 00:00:00 2001 From: Sandro Campos Date: Fri, 16 Oct 2026 11:45:43 -0400 Subject: [PATCH] Fix benchmarks --- benchmarks/test_column_mean.py | 8 ++++---- benchmarks/test_crossmatch.py | 4 ++-- benchmarks/test_local_io.py | 20 ++++++++++---------- benchmarks/test_protocols.py | 8 ++++++-- 4 files changed, 22 insertions(+), 18 deletions(-) diff --git a/benchmarks/test_column_mean.py b/benchmarks/test_column_mean.py index 58dd6bf..6e00288 100644 --- a/benchmarks/test_column_mean.py +++ b/benchmarks/test_column_mean.py @@ -5,8 +5,8 @@ import nested_pandas as npd -def test_pyarrow_mean(gaia_collection_path, lbench): - gaia_root = gaia_collection_path / "gaia" +def test_pyarrow_mean(gaia_local_collection_path, lbench): + gaia_root = gaia_local_collection_path / "gaia" parquet_root = f"{gaia_root}/dataset" pyarrow_ds = ds.parquet_dataset(f"{parquet_root}/_metadata") @@ -31,7 +31,7 @@ def dataset_mean(dataset, field: str, *, use_threads: bool = True): lbench(dataset_mean, pyarrow_ds, "phot_g_mean_mag") -def test_lsdb_mean(gaia_collection_path, lbench_dask): +def test_lsdb_mean(gaia_local_collection_path, lbench_dask): def catalog_mean(df, target_column=""): result = npd.NestedFrame( { @@ -41,7 +41,7 @@ def catalog_mean(df, target_column=""): ) return result - lsdb_gaia = lsdb.open_catalog(gaia_collection_path, columns=["phot_g_mean_mag"]) + lsdb_gaia = lsdb.open_catalog(gaia_local_collection_path, columns=["phot_g_mean_mag"]) unrealized = lsdb_gaia.map_partitions( catalog_mean, target_column="phot_g_mean_mag", diff --git a/benchmarks/test_crossmatch.py b/benchmarks/test_crossmatch.py index f656422..0b009e2 100644 --- a/benchmarks/test_crossmatch.py +++ b/benchmarks/test_crossmatch.py @@ -3,10 +3,10 @@ import nested_pandas as npd from hats.pixel_math import HealpixPixel -from lsdb.dask.crossmatch_catalog_data import perform_crossmatch +from lsdb.operations.functions.crossmatch_catalog_data import perform_crossmatch +from lsdb.operations.functions.merge_catalog_functions import generate_meta_df_for_joined_tables from hats.io import pixel_catalog_file from lsdb.core.crossmatch.kdtree_match import KdTreeCrossmatch -from lsdb.dask.merge_catalog_functions import generate_meta_df_for_joined_tables @pytest.mark.lbench_memory diff --git a/benchmarks/test_local_io.py b/benchmarks/test_local_io.py index 09423bc..064358b 100644 --- a/benchmarks/test_local_io.py +++ b/benchmarks/test_local_io.py @@ -5,8 +5,8 @@ import pytest -def test_local_catalog_partition_read(gaia_collection_path, lbench_dask): - gaia = lsdb.read_hats(gaia_collection_path) +def test_local_catalog_partition_read(gaia_local_collection_path, lbench_dask): + gaia = lsdb.read_hats(gaia_local_collection_path) cat = gaia.partitions[0] def load_partition(): @@ -15,8 +15,8 @@ def load_partition(): lbench_dask(load_partition) -def test_local_catalog_npd_read(gaia_collection_path, lbench): - gaia = lsdb.read_hats(gaia_collection_path) +def test_local_catalog_npd_read(gaia_local_collection_path, lbench): + gaia = lsdb.read_hats(gaia_local_collection_path) partition_0_pixel = gaia.partitions[0].get_healpix_pixels()[0] partition_0_path = hats.io.paths.pixel_catalog_file(gaia.hc_structure.catalog_base_dir, partition_0_pixel) @@ -26,8 +26,8 @@ def load_partition_npd(): lbench(load_partition_npd) -def test_local_catalog_pd_read(gaia_collection_path, lbench): - gaia = lsdb.read_hats(gaia_collection_path) +def test_local_catalog_pd_read(gaia_local_collection_path, lbench): + gaia = lsdb.read_hats(gaia_local_collection_path) partition_0_pixel = gaia.partitions[0].get_healpix_pixels()[0] partition_0_path = hats.io.paths.pixel_catalog_file(gaia.hc_structure.catalog_base_dir, partition_0_pixel) @@ -38,8 +38,8 @@ def load_partition_pd(): @pytest.mark.benchmark(min_rounds=1) -def test_local_catalog_multi_partition_read(gaia_collection_path, lbench_dask): - gaia = lsdb.read_hats(gaia_collection_path) +def test_local_catalog_multi_partition_read(gaia_local_collection_path, lbench_dask): + gaia = lsdb.read_hats(gaia_local_collection_path) n_partitions = 10 cat = gaia.partitions[:n_partitions] @@ -50,8 +50,8 @@ def load_partitions(): @pytest.mark.benchmark(min_rounds=1) -def test_local_catalog_multi_partition_npd_read(gaia_collection_path, lbench): - gaia = lsdb.read_hats(gaia_collection_path) +def test_local_catalog_multi_partition_npd_read(gaia_local_collection_path, lbench): + gaia = lsdb.read_hats(gaia_local_collection_path) n_partitions = 10 partition_paths = [] for i in range(n_partitions): diff --git a/benchmarks/test_protocols.py b/benchmarks/test_protocols.py index b14dd2e..32f3709 100644 --- a/benchmarks/test_protocols.py +++ b/benchmarks/test_protocols.py @@ -83,7 +83,9 @@ def crossmatch(): cone = lsdb.ConeSearch(ra=0.0, dec=0.0, radius_arcsec=3600) gaia_int = lsdb.open_catalog(GAIA_URLS[internal_protocol], search_filter=cone) gaia_ext = lsdb.open_catalog(GAIA_URLS[external_protocol], search_filter=cone) - xmatch = gaia_int.crossmatch(gaia_ext, radius_arcsec=1, suffixes=("_1", "_2")) + xmatch = gaia_int.crossmatch( + gaia_ext, radius_arcsec=1, suffixes=("_1", "_2"), suffix_method="all_columns" + ) xmatch.compute() lbench_dask(crossmatch) @@ -98,7 +100,9 @@ def crossmatch(): cone = lsdb.ConeSearch(ra=0.0, dec=0.0, radius_arcsec=3600) gaia_int = lsdb.open_catalog(GAIA_URLS[internal_protocol], columns=["ra", "dec"], search_filter=cone) gaia_ext = lsdb.open_catalog(GAIA_URLS[external_protocol], columns=["ra", "dec"], search_filter=cone) - xmatch = gaia_int.crossmatch(gaia_ext, radius_arcsec=1, suffixes=("_1", "_2")) + xmatch = gaia_int.crossmatch( + gaia_ext, radius_arcsec=1, suffixes=("_1", "_2"), suffix_method="all_columns" + ) xmatch.compute() lbench_dask(crossmatch)