Repository: conda/rattler
Pull request #1347: feat(py): query based on matchspec and package count
Author: login: baszalmstra
Labels: name: python-bindings
Created: 2025-05-22T14:44:23Z | Merged: 2025-05-22T15:57:52Z

Adds the ability to load all repodata records from a SparseRepoData, query packages based on matchspec, and return the total number of packages. 


Diff:
diff --git a/crates/rattler_repodata_gateway/src/gateway/local_subdir.rs b/crates/rattler_repodata_gateway/src/gateway/local_subdir.rs
index 75e23f2bc7..b99671876b 100644
--- a/crates/rattler_repodata_gateway/src/gateway/local_subdir.rs
+++ b/crates/rattler_repodata_gateway/src/gateway/local_subdir.rs
@@ -92,7 +92,7 @@ impl SubdirClient for LocalSubdirClient {
     fn package_names(&self) -> Vec<String> {
         let sparse_repodata: Arc<SparseRepoData> = self.sparse.clone();
         sparse_repodata
-            .package_names()
+            .package_names(PackageFormatSelection::PreferConda)
             .map(std::convert::Into::into)
             .collect()
     }
diff --git a/crates/rattler_repodata_gateway/src/sparse/mod.rs b/crates/rattler_repodata_gateway/src/sparse/mod.rs
index 80f9a05f80..9ba12b5f63 100644
--- a/crates/rattler_repodata_gateway/src/sparse/mod.rs
+++ b/crates/rattler_repodata_gateway/src/sparse/mod.rs
@@ -3,19 +3,12 @@
 
 #![allow(clippy::mem_forget)]
 
-use std::{
-    collections::{HashSet, VecDeque},
-    fmt, io,
-    marker::PhantomData,
-    path::Path,
-};
-
 use bytes::Bytes;
 use fs_err as fs;
 use itertools::Itertools;
 use rattler_conda_types::{
-    compute_package_url, package::ArchiveType, Channel, ChannelInfo, PackageName, PackageRecord,
-    RepoDataRecord,
+    compute_package_url, package::ArchiveType, Channel, ChannelInfo, MatchSpec, Matches,
+    PackageName, PackageRecord, RepoDataRecord,
 };
 use rattler_redaction::Redact;
 use serde::{
@@ -23,6 +16,13 @@ use serde::{
     Deserialize, Deserializer,
 };
 use serde_json::value::RawValue;
+use std::borrow::Borrow;
+use std::{
+    collections::{HashSet, VecDeque},
+    fmt, io,
+    marker::PhantomData,
+    path::Path,
+};
 use superslice::Ext;
 use thiserror::Error;
 
@@ -197,14 +197,81 @@ impl SparseRepoData {
     /// This works by iterating over all elements in the `packages` and
     /// `conda_packages` fields of the repodata and returning the unique
     /// package names.
-    pub fn package_names(&self) -> impl Iterator<Item = &'_ str> {
+    pub fn package_names(
+        &self,
+        package_format_selection: PackageFormatSelection,
+    ) -> impl Iterator<Item = &'_ str> {
         let repo_data = self.inner.borrow_repo_data();
-        let tar_baz2_packages = repo_data.packages.iter().map(|(name, _)| name.package);
-        let conda_packages = repo_data
-            .conda_packages
-            .iter()
-            .map(|(name, _)| name.package);
-        tar_baz2_packages.merge(conda_packages).dedup()
+
+        fn select_package_name<'i>((filename, _): &(PackageFilename<'i>, &'i RawValue)) -> &'i str {
+            filename.package
+        }
+
+        let tar_baz2_packages = repo_data.packages.iter().map(select_package_name);
+        let conda_packages = repo_data.conda_packages.iter().map(select_package_name);
+
+        match package_format_selection {
+            PackageFormatSelection::Both | PackageFormatSelection::PreferConda => {
+                itertools::Either::Left(tar_baz2_packages.merge(conda_packages).dedup())
+            }
+            PackageFormatSelection::OnlyTarBz2 => {
+                itertools::Either::Right(tar_baz2_packages.dedup())
+            }
+            PackageFormatSelection::OnlyConda => itertools::Either::Right(conda_packages.dedup()),
+        }
+    }
+
+    /// Returns the number of records in this instance.
+    pub fn record_count(&self, package_format_selection: PackageFormatSelection) -> usize {
+        match package_format_selection {
+            PackageFormatSelection::PreferConda => {
+                let repo_data = self.inner.borrow_repo_data();
+                let tar_baz2_packages = repo_data.packages.iter().map(|(name, _)| name.package);
+                let conda_packages = repo_data
+                    .conda_packages
+                    .iter()
+                    .map(|(name, _)| name.package);
+                tar_baz2_packages.merge(conda_packages).dedup().count()
+            }
+            PackageFormatSelection::Both => {
+                self.inner.borrow_repo_data().packages.len()
+                    + self.inner.borrow_repo_data().conda_packages.len()
+            }
+            PackageFormatSelection::OnlyTarBz2 => self.inner.borrow_repo_data().packages.len(),
+            PackageFormatSelection::OnlyConda => self.inner.borrow_repo_data().conda_packages.len(),
+        }
+    }
+
+    /// Returns all the records that matches any of the specified match spec.
+    pub fn load_matching_records(
+        &self,
+        spec: impl IntoIterator<Item = impl Borrow<MatchSpec>>,
+        variant_consolidation: PackageFormatSelection,
+    ) -> io::Result<Vec<RepoDataRecord>> {
+        let mut result = Vec::new();
+        let repo_data = self.inner.borrow_repo_data();
+        let base_url = repo_data.info.as_ref().and_then(|i| i.base_url.as_deref());
+        for (package_name, specs) in &spec.into_iter().chunk_by(|spec| spec.borrow().name.clone()) {
+            let grouped_specs = specs.into_iter().collect::<Vec<_>>();
+            let mut parsed_records = parse_records(
+                package_name.as_ref(),
+                &repo_data.packages,
+                &repo_data.conda_packages,
+                variant_consolidation,
+                base_url,
+                &self.channel,
+                &self.subdir,
+                self.patch_record_fn,
+                |record| {
+                    grouped_specs
+                        .iter()
+                        .any(|spec| spec.borrow().matches(&record.package_record))
+                },
+            )?;
+            result.append(&mut parsed_records);
+        }
+
+        Ok(result)
     }
 
     /// Returns all the records for the specified package name.
@@ -216,7 +283,27 @@ impl SparseRepoData {
         let repo_data = self.inner.borrow_repo_data();
         let base_url = repo_data.info.as_ref().and_then(|i| i.base_url.as_deref());
         parse_records(
-            package_name,
+            Some(package_name),
+            &repo_data.packages,
+            &repo_data.conda_packages,
+            variant_consolidation,
+            base_url,
+            &self.channel,
+            &self.subdir,
+            self.patch_record_fn,
+            |_| true, // Dont filter anything out
+        )
+    }
+
+    /// Returns all the records for the specified package format(s).
+    pub fn load_all_records(
+        &self,
+        variant_consolidation: PackageFormatSelection,
+    ) -> io::Result<Vec<RepoDataRecord>> {
+        let repo_data = self.inner.borrow_repo_data();
+        let base_url = repo_data.info.as_ref().and_then(|i| i.base_url.as_deref());
+        parse_records(
+            None,
             &repo_data.packages,
             &repo_data.conda_packages,
             variant_consolidation,
@@ -224,6 +311,7 @@ impl SparseRepoData {
             &self.channel,
             &self.subdir,
             self.patch_record_fn,
+            |_| true,
         )
     }
 
@@ -262,7 +350,7 @@ impl SparseRepoData {
 
                 // Get all records from the repodata
                 let mut records = parse_records(
-                    &next_package,
+                    Some(&next_package),
                     &repo_data_packages.packages,
                     &repo_data_packages.conda_packages,
                     variant_consolidation,
@@ -270,6 +358,7 @@ impl SparseRepoData {
                     &repo_data.channel,
                     &repo_data.subdir,
                     patch_function,
+                    |_| true,
                 )?;
 
                 // Iterate over all packages to find recursive dependencies.
@@ -328,10 +417,15 @@ struct LazyRepoData<'i> {
 /// package name.
 fn find_package_in_slice<'a, 'i: 'a>(
     slice: &'a [(PackageFilename<'i>, &'i RawValue)],
-    package_name: &PackageName,
+    package_name: Option<&PackageName>,
 ) -> impl Iterator<Item = (PackageFilename<'i>, &'i RawValue)> + 'a {
-    let range =
-        slice.equal_range_by(|(package, _)| package.package.cmp(package_name.as_normalized()));
+    let range = match package_name {
+        None => 0..slice.len(),
+        Some(package_name) => {
+            slice.equal_range_by(|(package, _)| package.package.cmp(package_name.as_normalized()))
+        }
+    };
+
     slice[range]
         .iter()
         .map(|(filename, raw_json)| (*filename, *raw_json))
@@ -357,8 +451,8 @@ fn add_stripped_filename<'i>(
 
 /// Parse the records for the specified package from the raw index
 #[allow(clippy::too_many_arguments)]
-fn parse_records<'i>(
-    package_name: &PackageName,
+fn parse_records<'i, F: Fn(&RepoDataRecord) -> bool>(
+    package_name: Option<&PackageName>,
     tar_bz2_packages: &[(PackageFilename<'i>, &'i RawValue)],
     conda_packages: &[(PackageFilename<'i>, &'i RawValue)],
     variant_consolidation: PackageFormatSelection,
@@ -366,6 +460,7 @@ fn parse_records<'i>(
     channel: &Channel,
     subdir: &str,
     patch_function: Option<fn(&mut PackageRecord)>,
+    filter_function: F,
 ) -> io::Result<Vec<RepoDataRecord>> {
     match variant_consolidation {
         PackageFormatSelection::PreferConda => {
@@ -393,6 +488,7 @@ fn parse_records<'i>(
                 channel,
                 subdir,
                 patch_function,
+                filter_function,
             )
         }
         PackageFormatSelection::Both => {
@@ -404,59 +500,92 @@ fn parse_records<'i>(
                 channel,
                 subdir,
                 patch_function,
+                filter_function,
             )
         }
         PackageFormatSelection::OnlyTarBz2 => {
             let tar_bz2_packages = find_package_in_slice(tar_bz2_packages, package_name);
-            parse_records_raw(tar_bz2_packages, base_url, channel, subdir, patch_function)
+            parse_records_raw(
+                tar_bz2_packages,
+                base_url,
+                channel,
+                subdir,
+                patch_function,
+                filter_function,
+            )
         }
         PackageFormatSelection::OnlyConda => {
             let conda_packages = find_package_in_slice(conda_packages, package_name);
-            parse_records_raw(conda_packages, base_url, channel, subdir, patch_function)
+            parse_records_raw(
+                conda_packages,
+                base_url,
+                channel,
+                subdir,
+                patch_function,
+                filter_function,
+            )
         }
     }
 }
 
-fn parse_records_raw<'i>(
-    packages: impl Iterator<Item = (PackageFilename<'i>, &'i RawValue)>,
+fn parse_record_raw<'i>(
+    (filename, raw_json): (PackageFilename<'i>, &'i RawValue),
     base_url: Option<&str>,
     channel: &Channel,
+    channel_name: Option<String>,
     subdir: &str,
     patch_function: Option<fn(&mut PackageRecord)>,
-) -> io::Result<Vec<RepoDataRecord>> {
-    let channel_name = channel.base_url.clone();
-    let expected_capacity = packages.size_hint();
-    let mut result = Vec::with_capacity(expected_capacity.1.unwrap_or(expected_capacity.0));
-    for (filename, raw_json) in packages {
-        let mut package_record: PackageRecord = serde_json::from_str(raw_json.get())?;
-        // Overwrite subdir if its empty
-        if package_record.subdir.is_empty() {
-            package_record.subdir = subdir.to_owned();
-        }
-        result.push(RepoDataRecord {
-            url: compute_package_url(
-                &channel
-                    .base_url
-                    .url()
-                    .join(&format!("{}/", &package_record.subdir))
-                    .expect("failed determine repo_base_url"),
-                base_url,
-                filename.filename,
-            ),
-            channel: Some(channel_name.url().clone().redact().to_string()),
-            package_record,
-            file_name: filename.filename.to_owned(),
-        });
+) -> io::Result<RepoDataRecord> {
+    let mut package_record: PackageRecord = serde_json::from_str(raw_json.get())?;
+    // Overwrite subdir if its empty
+    if package_record.subdir.is_empty() {
+        package_record.subdir = subdir.to_owned();
     }
+    let mut record = RepoDataRecord {
+        url: compute_package_url(
+            &channel
+                .base_url
+                .url()
+                .join(&format!("{}/", &package_record.subdir))
+                .expect("failed determine repo_base_url"),
+            base_url,
+            filename.filename,
+        ),
+        channel: channel_name.clone(),
+        package_record,
+        file_name: filename.filename.to_owned(),
+    };
 
     // Apply the patch function if one was specified
     if let Some(patch_fn) = patch_function {
-        for record in &mut result {
-            patch_fn(&mut record.package_record);
-        }
+        patch_fn(&mut record.package_record);
     }
 
-    Ok(result)
+    Ok(record)
+}
+
+fn parse_records_raw<'i, F: Fn(&RepoDataRecord) -> bool>(
+    packages: impl Iterator<Item = (PackageFilename<'i>, &'i RawValue)>,
+    base_url: Option<&str>,
+    channel: &Channel,
+    subdir: &str,
+    patch_function: Option<fn(&mut PackageRecord)>,
+    filter_function: F,
+) -> io::Result<Vec<RepoDataRecord>> {
+    let channel_name = channel.base_url.url().clone().redact().to_string();
+    packages
+        .map(move |record| {
+            parse_record_raw(
+                record,
+                base_url,
+                channel,
+                Some(channel_name.clone()),
+                subdir,
+                patch_function,
+            )
+        })
+        .filter_ok(filter_function)
+        .collect()
 }
 
 /// A helper function that immediately loads the records for the given packages
@@ -621,7 +750,9 @@ mod test {
     use bytes::Bytes;
     use fs_err as fs;
     use itertools::Itertools;
-    use rattler_conda_types::{Channel, ChannelConfig, PackageName, RepoData, RepoDataRecord};
+    use rattler_conda_types::{
+        Channel, ChannelConfig, MatchSpec, PackageName, ParseStrictness, RepoData, RepoDataRecord,
+    };
     use rstest::rstest;
 
     use super::{
@@ -866,23 +997,33 @@ mod test {
         .unwrap();
 
         assert_eq!(repo_data.packages.len(), 0);
-        assert_eq!(sparse_repodata.package_names().try_len().unwrap(), 0);
+        assert_eq!(
+            sparse_repodata
+                .package_names(PackageFormatSelection::default())
+                .try_len()
+                .unwrap(),
+            0
+        );
     }
 
-    #[test]
-    fn dedup_packages() {
+    #[rstest]
+    #[case::both(PackageFormatSelection::Both)]
+    #[case::prefer_conda(PackageFormatSelection::PreferConda)]
+    #[case::only_tar_bz2(PackageFormatSelection::OnlyTarBz2)]
+    #[case::only_conda(PackageFormatSelection::OnlyConda)]
+    fn dedup_packages(#[case] variant: PackageFormatSelection) {
         let (channel, platform, path) = dummy_repo_data();
         let sparse = SparseRepoData::from_file(channel, platform, path, None).unwrap();
-        let names = sparse.package_names().collect_vec();
+        let names = sparse.package_names(variant).collect_vec();
         let deduped_names = names.iter().copied().collect::<HashSet<_>>();
         assert_eq!(names.len(), deduped_names.len());
     }
 
     #[rstest]
-    #[case(PackageFormatSelection::Both)]
-    #[case(PackageFormatSelection::PreferConda)]
-    #[case(PackageFormatSelection::OnlyTarBz2)]
-    #[case(PackageFormatSelection::OnlyConda)]
+    #[case::both(PackageFormatSelection::Both)]
+    #[case::prefer_conda(PackageFormatSelection::PreferConda)]
+    #[case::only_tar_bz2(PackageFormatSelection::OnlyTarBz2)]
+    #[case::only_conda(PackageFormatSelection::OnlyConda)]
     fn test_only_conda(#[case] variant: PackageFormatSelection) {
         let (channel, platform, path) = dummy_repo_data();
         let sparse = SparseRepoData::from_file(channel, platform, path, None).unwrap();
@@ -897,4 +1038,46 @@ mod test {
             insta::assert_snapshot!(records.join("\n"));
         });
     }
+
+    #[test]
+    fn test_query() {
+        let (channel, platform, path) = dummy_repo_data();
+        let sparse = SparseRepoData::from_file(channel, platform, path, None).unwrap();
+        let records = sparse
+            .load_matching_records(
+                vec![
+                    MatchSpec::from_str("bors 1.*", ParseStrictness::Lenient).unwrap(),
+                    MatchSpec::from_str("issue_717", ParseStrictness::Lenient).unwrap(),
+                ],
+                PackageFormatSelection::default(),
+            )
+            .unwrap()
+            .into_iter()
+            .map(|record| record.file_name)
+            .collect::<Vec<_>>();
+
+        insta::assert_snapshot!(records.join("\n"), @r###"
+        bors-1.0-bla_1.tar.bz2
+        bors-1.1-bla_1.conda
+        bors-1.2.1-bla_1.tar.bz2
+        issue_717-2.1-bla_1.conda
+        "###);
+    }
+
+    #[test]
+    fn test_nameless_query() {
+        let (channel, platform, path) = dummy_repo_data();
+        let sparse = SparseRepoData::from_file(channel, platform, path, None).unwrap();
+        let records = sparse
+            .load_matching_records(
+                vec![MatchSpec::from_str("* 12.5", ParseStrictness::Lenient).unwrap()],
+                PackageFormatSelection::default(),
+            )
+            .unwrap()
+            .into_iter()
+            .map(|record| record.file_name)
+            .collect::<Vec<_>>();
+
+        insta::assert_snapshot!(records.join("\n"), @"cuda-version-12.5-hd4f0392_3.conda");
+    }
 }
diff --git a/py-rattler/rattler/repo_data/sparse.py b/py-rattler/rattler/repo_data/sparse.py
index 83e1ea0b63..c1a30d96e0 100644
--- a/py-rattler/rattler/repo_data/sparse.py
+++ b/py-rattler/rattler/repo_data/sparse.py
@@ -1,8 +1,10 @@
 from __future__ import annotations
 import os
 from pathlib import Path
-from typing import List, Optional, Type, Literal
+from typing import List, Optional, Type, Literal, Iterable
 from types import TracebackType
+
+from rattler.match_spec.match_spec import MatchSpec
 from rattler.channel.channel import Channel
 from rattler.package.package_name import PackageName
 from enum import Enum
@@ -101,7 +103,9 @@ def close(self) -> None:
         """
         self._sparse.close()
 
-    def package_names(self) -> List[str]:
+    def package_names(
+        self, package_format_selection: PackageFormatSelection = PackageFormatSelection.PREFER_CONDA
+    ) -> List[str]:
         """
         Returns a list over all package names in this repodata file.
         This works by iterating over all elements in the `packages` and
@@ -123,10 +127,21 @@ def package_names(self) -> List[str]:
         >>>
         ```
         """
-        return self._sparse.package_names()
+        return self._sparse.package_names(package_format_selection.value)
+
+    def record_count(
+        self, package_format_selection: PackageFormatSelection = PackageFormatSelection.PREFER_CONDA
+    ) -> int:
+        """
+        Returns the total number of packages in this repodata file.
+        :return:
+        """
+        return self._sparse.record_count(package_format_selection.value)
 
     def load_records(
-        self, package_name: PackageName, variant_selection: PackageFormatSelection = PackageFormatSelection.PREFER_CONDA
+        self,
+        package_name: str | PackageName,
+        package_format_selection: PackageFormatSelection = PackageFormatSelection.PREFER_CONDA,
     ) -> List[RepoDataRecord]:
         """
         Returns all the records for the specified package name.
@@ -147,10 +162,65 @@ def load_records(
         >>>
         ```
         """
+        if not isinstance(package_name, PackageName):
+            package_name = PackageName(package_name)
+        return [
+            RepoDataRecord._from_py_record(record)
+            for record in self._sparse.load_records(package_name._name, package_format_selection.value)
+        ]
+
+    def load_all_records(
+        self, package_format_selection: PackageFormatSelection = PackageFormatSelection.PREFER_CONDA
+    ) -> List[RepoDataRecord]:
+        """
+        Returns all the records for the specified package name.
+
+        Examples
+        --------
+        ```python
+        >>> from rattler import Channel, ChannelConfig, RepoDataRecord, PackageName
+        >>> channel = Channel("dummy", ChannelConfig())
+        >>> path = "../test-data/channels/dummy/linux-64/repodata.json"
+        >>> sparse_data = SparseRepoData(channel, "linux-64", path)
+        >>> records = sparse_data.load_all_records()
+        >>> records
+        [...]
+        >>> isinstance(records[0], RepoDataRecord)
+        True
+        >>>
+        ```
+        """
         # maybe change package_name to Union[str, PackageName]
         return [
             RepoDataRecord._from_py_record(record)
-            for record in self._sparse.load_records(package_name._name, variant_selection.value)
+            for record in self._sparse.load_all_records(package_format_selection.value)
+        ]
+
+    def load_matching_records(
+        self,
+        specs: Iterable[MatchSpec],
+        package_format_selection: PackageFormatSelection = PackageFormatSelection.PREFER_CONDA,
+    ) -> List[RepoDataRecord]:
+        """
+        Returns all the records that match any of the specified MatchSpecs.
+
+        Examples
+        --------
+        ```python
+        >>> from rattler import Channel, ChannelConfig, RepoDataRecord, PackageName
+        >>> channel = Channel("dummy", ChannelConfig())
+        >>> path = "../test-data/channels/dummy/linux-64/repodata.json"
+        >>> sparse_data = SparseRepoData(channel, "linux-64", path)
+        >>> [record.file_name for record in sparse_data.load_matching_records([MatchSpec("* 12.5")])]
+        ['cuda-version-12.5-hd4f0392_3.conda']
+        >>>
+        ```
+        """
+        return [
+            RepoDataRecord._from_py_record(record)
+            for record in self._sparse.load_matching_records(
+                [spec._match_spec for spec in specs], package_format_selection.value
+            )
         ]
 
     @property
@@ -176,7 +246,7 @@ def subdir(self) -> str:
     def load_records_recursive(
         repo_data: List[SparseRepoData],
         package_names: List[PackageName],
-        variant_selection: PackageFormatSelection = PackageFormatSelection.PREFER_CONDA,
+        package_format_selection: PackageFormatSelection = PackageFormatSelection.PREFER_CONDA,
     ) -> List[List[RepoDataRecord]]:
         """
         Given a set of [`SparseRepoData`]s load all the records
@@ -202,7 +272,7 @@ def load_records_recursive(
         return [
             [RepoDataRecord._from_py_record(record) for record in list_of_records]
             for list_of_records in PySparseRepoData.load_records_recursive(
-                [r._sparse for r in repo_data], [p._name for p in package_names], variant_selection.value
+                [r._sparse for r in repo_data], [p._name for p in package_names], package_format_selection.value
             )
         ]
 
diff --git a/py-rattler/src/match_spec.rs b/py-rattler/src/match_spec.rs
index d4d3680488..8e669c7d57 100644
--- a/py-rattler/src/match_spec.rs
+++ b/py-rattler/src/match_spec.rs
@@ -1,3 +1,4 @@
+use std::borrow::Borrow;
 use std::sync::Arc;
 
 use pyo3::{pyclass, pymethods, types::PyBytes, Bound, PyResult, Python};
@@ -27,6 +28,12 @@ impl From<PyMatchSpec> for MatchSpec {
     }
 }
 
+impl Borrow<MatchSpec> for PyMatchSpec {
+    fn borrow(&self) -> &MatchSpec {
+        &self.inner
+    }
+}
+
 #[pymethods]
 impl PyMatchSpec {
     #[new]
diff --git a/py-rattler/src/repo_data/sparse.rs b/py-rattler/src/repo_data/sparse.rs
index 8b0dccee02..a3858fc635 100644
--- a/py-rattler/src/repo_data/sparse.rs
+++ b/py-rattler/src/repo_data/sparse.rs
@@ -1,10 +1,11 @@
 use std::{path::PathBuf, sync::Arc};
 
-use pyo3::{pyclass, pymethods, Bound, PyResult, Python};
+use pyo3::{pyclass, pymethods, Bound, PyRef, PyResult, Python};
 
 use rattler_repodata_gateway::sparse::{PackageFormatSelection, SparseRepoData};
 
 use crate::channel::PyChannel;
+use crate::match_spec::PyMatchSpec;
 use crate::package_name::PyPackageName;
 use crate::record::PyRecord;
 use parking_lot::RwLock;
@@ -87,12 +88,29 @@ impl PySparseRepoData {
         Ok(SparseRepoData::from_file(channel.into(), subdir, path, None)?.into())
     }
 
-    pub fn package_names(&self) -> PyResult<Vec<String>> {
+    pub fn package_names(
+        &self,
+        package_format_selection: PyPackageFormatSelection,
+    ) -> PyResult<Vec<String>> {
         let lock = self.inner.read();
         let Some(sparse) = lock.as_ref() else {
             return Err(PyValueError::new_err("I/O operation on closed file."));
         };
-        Ok(sparse.package_names().map(Into::into).collect::<Vec<_>>())
+        Ok(sparse
+            .package_names(package_format_selection.into())
+            .map(Into::into)
+            .collect::<Vec<_>>())
+    }
+
+    pub fn record_count(
+        &self,
+        package_format_selection: PyPackageFormatSelection,
+    ) -> PyResult<usize> {
+        let lock = self.inner.read();
+        let Some(sparse) = lock.as_ref() else {
+            return Err(PyValueError::new_err("I/O operation on closed file."));
+        };
+        Ok(sparse.record_count(package_format_selection.into()))
     }
 
     pub fn load_records(
@@ -111,6 +129,40 @@ impl PySparseRepoData {
             .collect::<Vec<_>>())
     }
 
+    pub fn load_all_records(
+        &self,
+        package_format_selection: PyPackageFormatSelection,
+    ) -> PyResult<Vec<PyRecord>> {
+        let lock = self.inner.read();
+        let Some(sparse) = lock.as_ref() else {
+            return Err(PyValueError::new_err("I/O operation on closed file."));
+        };
+        Ok(sparse
+            .load_all_records(package_format_selection.into())?
+            .into_iter()
+            .map(Into::into)
+            .collect::<Vec<_>>())
+    }
+
+    pub fn load_matching_records(
+        &self,
+        specs: Vec<PyRef<'_, PyMatchSpec>>,
+        package_format_selection: PyPackageFormatSelection,
+    ) -> PyResult<Vec<PyRecord>> {
+        let lock = self.inner.read();
+        let Some(sparse) = lock.as_ref() else {
+            return Err(PyValueError::new_err("I/O operation on closed file."));
+        };
+        Ok(sparse
+            .load_matching_records(
+                specs.iter().map(|spec| &spec.inner),
+                package_format_selection.into(),
+            )?
+            .into_iter()
+            .map(Into::into)
+            .collect::<Vec<_>>())
+    }
+
     #[getter]
     pub fn subdir(&self) -> String {
         self.subdir.clone()


Reviews:
user: login: baszalmstra
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-22T14:45:11Z
user: login: jaimergp
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-22T14:53:44Z
user: login: baszalmstra
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-22T14:57:15Z
user: login: wolfv
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T15:14:04Z
user: login: baszalmstra
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-22T15:15:59Z
user: login: wolfv
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T15:17:18Z

Review comments:
diff_hunk: @@ -197,14 +197,81 @@ impl SparseRepoData {
     /// This works by iterating over all elements in the `packages` and
     /// `conda_packages` fields of the repodata and returning the unique
     /// package names.
-    pub fn package_names(&self) -> impl Iterator<Item = &'_ str> {
+    pub fn package_names(
+        &self,
+        package_format_selection: PackageFormatSelection,
path: crates/rattler_repodata_gateway/src/sparse/mod.rs
user: login: baszalmstra
body: This is new! This functionality was missing in previous implementations.
created_at: 2025-05-22T14:45:11Z
author_association: COLLABORATOR
diff_hunk: @@ -357,15 +451,16 @@ fn add_stripped_filename<'i>(
 
 /// Parse the records for the specified package from the raw index
 #[allow(clippy::too_many_arguments)]
-fn parse_records<'i>(
-    package_name: &PackageName,
+fn parse_records<'i, F: Fn(&RepoDataRecord) -> bool>(
+    package_name: Option<&PackageName>,
     tar_bz2_packages: &[(PackageFilename<'i>, &'i RawValue)],
     conda_packages: &[(PackageFilename<'i>, &'i RawValue)],
     variant_consolidation: PackageFormatSelection,
     base_url: Option<&str>,
     channel: &Channel,
     subdir: &str,
     patch_function: Option<fn(&mut PackageRecord)>,
+    filter_function: F,
path: crates/rattler_repodata_gateway/src/sparse/mod.rs
user: login: baszalmstra
body: I choose to implement this as a generic to generate two code paths for performance. If no filtering is required this will likely be optimized out.
created_at: 2025-05-22T14:46:00Z
author_association: COLLABORATOR
diff_hunk: @@ -216,14 +283,35 @@ impl SparseRepoData {
         let repo_data = self.inner.borrow_repo_data();
         let base_url = repo_data.info.as_ref().and_then(|i| i.base_url.as_deref());
         parse_records(
-            package_name,
+            Some(package_name),
+            &repo_data.packages,
+            &repo_data.conda_packages,
+            variant_consolidation,
+            base_url,
+            &self.channel,
+            &self.subdir,
+            self.patch_record_fn,
+            |_| true, // Dont filter anything out
+        )
+    }
+
+    /// Returns all the records for the specified package name.
path: crates/rattler_repodata_gateway/src/sparse/mod.rs
user: login: jaimergp
body: ```suggestion
    /// Returns all the records for the specified package format(s).
```

I think?
created_at: 2025-05-22T14:53:44Z
author_association: MEMBER
diff_hunk: @@ -197,14 +197,81 @@ impl SparseRepoData {
     /// This works by iterating over all elements in the `packages` and
     /// `conda_packages` fields of the repodata and returning the unique
     /// package names.
-    pub fn package_names(&self) -> impl Iterator<Item = &'_ str> {
+    pub fn package_names(
+        &self,
+        package_format_selection: PackageFormatSelection,
+    ) -> impl Iterator<Item = &'_ str> {
         let repo_data = self.inner.borrow_repo_data();
-        let tar_baz2_packages = repo_data.packages.iter().map(|(name, _)| name.package);
-        let conda_packages = repo_data
-            .conda_packages
-            .iter()
-            .map(|(name, _)| name.package);
-        tar_baz2_packages.merge(conda_packages).dedup()
+
+        fn select_package_name<'i>((filename, _): &(PackageFilename<'i>, &'i RawValue)) -> &'i str {
+            filename.package
+        }
+
+        let tar_baz2_packages = repo_data.packages.iter().map(select_package_name);
+        let conda_packages = repo_data.conda_packages.iter().map(select_package_name);
+
+        match package_format_selection {
+            PackageFormatSelection::Both | PackageFormatSelection::PreferConda => {
+                itertools::Either::Left(tar_baz2_packages.merge(conda_packages).dedup())
+            }
+            PackageFormatSelection::OnlyTarBz2 => {
+                itertools::Either::Right(tar_baz2_packages.dedup())
+            }
+            PackageFormatSelection::OnlyConda => itertools::Either::Right(conda_packages.dedup()),
+        }
+    }
+
+    /// Returns the number of packages in this instance.
+    pub fn package_count(&self, package_format_selection: PackageFormatSelection) -> usize {
path: crates/rattler_repodata_gateway/src/sparse/mod.rs
user: login: wolfv
body: I wonder if this is confusing. I think I initially thought it would be a count of package names.

Maybe package_record_count would be more explicit?
created_at: 2025-05-22T15:14:04Z
author_association: CONTRIBUTOR
diff_hunk: @@ -197,14 +197,81 @@ impl SparseRepoData {
     /// This works by iterating over all elements in the `packages` and
     /// `conda_packages` fields of the repodata and returning the unique
     /// package names.
-    pub fn package_names(&self) -> impl Iterator<Item = &'_ str> {
+    pub fn package_names(
+        &self,
+        package_format_selection: PackageFormatSelection,
+    ) -> impl Iterator<Item = &'_ str> {
         let repo_data = self.inner.borrow_repo_data();
-        let tar_baz2_packages = repo_data.packages.iter().map(|(name, _)| name.package);
-        let conda_packages = repo_data
-            .conda_packages
-            .iter()
-            .map(|(name, _)| name.package);
-        tar_baz2_packages.merge(conda_packages).dedup()
+
+        fn select_package_name<'i>((filename, _): &(PackageFilename<'i>, &'i RawValue)) -> &'i str {
+            filename.package
+        }
+
+        let tar_baz2_packages = repo_data.packages.iter().map(select_package_name);
+        let conda_packages = repo_data.conda_packages.iter().map(select_package_name);
+
+        match package_format_selection {
+            PackageFormatSelection::Both | PackageFormatSelection::PreferConda => {
+                itertools::Either::Left(tar_baz2_packages.merge(conda_packages).dedup())
+            }
+            PackageFormatSelection::OnlyTarBz2 => {
+                itertools::Either::Right(tar_baz2_packages.dedup())
+            }
+            PackageFormatSelection::OnlyConda => itertools::Either::Right(conda_packages.dedup()),
+        }
+    }
+
+    /// Returns the number of packages in this instance.
+    pub fn package_count(&self, package_format_selection: PackageFormatSelection) -> usize {
path: crates/rattler_repodata_gateway/src/sparse/mod.rs
user: login: baszalmstra
body: Or maybe just `record_count`?
created_at: 2025-05-22T15:15:59Z
author_association: COLLABORATOR

PR comments:
user: login: jaimergp
created_at: 2025-05-23T10:03:15Z
body: Looks like `.record_count()` is buggy in some cases. See this reproducer:

```python
from tempfile import NamedTemporaryFile
from urllib.request import urlretrieve

from rattler import Channel, SparseRepoData, PackageFormatSelection

channel_url = "https://conda.anaconda.org/conda-test"
subdir = "noarch"
repodata_url = "https://conda.anaconda.org/conda-test/noarch/repodata.json"
# This channel only contains .tar.bz2 files, a total of 12

with NamedTemporaryFile(suffix=".json", delete=False) as tmp:
    urlretrieve(repodata_url, tmp.name)

repo = SparseRepoData(Channel(channel_url), subdir, tmp.name)
for pkg_format in PackageFormatSelection:
    print(pkg_format)
    print(repo.record_count(pkg_format))
    print(len(repo.load_all_records(pkg_format)))
    print("---")
```

It should always report 12 (this channel only lists .tar.bz2 artifacts), but for some reason `.record_count()` reports 10 with PackageFormatSelection.PREFER_CONDA. `.load_all_records()` does report the correct thing all the time.

cc @baszalmstra 
author_association: MEMBER
user: login: baszalmstra
created_at: 2025-05-23T11:42:00Z
body: Let me fix that!
author_association: COLLABORATOR
user: login: baszalmstra
created_at: 2025-05-23T12:29:40Z
body: Should be fixed in: https://github.com/conda/rattler/pull/1350

Ill do a patch release when its merged!
author_association: COLLABORATOR