metrale_bench/benchmarks/bfcl/
dataset.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: Reading the materialised BFCL table and applying the draw to it.
4//! The draw takes the first `n` rows of each subset in file order, so the row
5//! order within a subset is the sample selection.
6//!
7//! Owner: bench, BFCL benchmark.
8//! Invariants: none beyond the types.
9
10use std::collections::BTreeMap;
11use std::path::Path;
12
13use anyhow::{Context, Result, bail};
14use serde::Deserialize;
15use serde_json::Value;
16
17use super::draw::{self, DrawSpec};
18
19/// 2026-09-26: One materialised BFCL sample.
20#[derive(Clone, Debug, Deserialize)]
21pub struct Sample {
22    pub subset: String,
23    pub sample_id: String,
24    pub messages: Vec<Value>,
25    pub tools: Vec<Value>,
26    pub tool_choice: Value,
27}
28
29/// 2026-09-26: One shard of a draw: `index` (0-based) of `count`.
30#[derive(Clone, Copy, Debug, PartialEq, Eq)]
31pub struct Shard {
32    /// 2026-09-26: Which shard, in `0..count`.
33    pub index: usize,
34    /// 2026-09-26: How many shards the draw is split into.
35    pub count: usize,
36}
37
38impl Shard {
39    /// 2026-09-26: Parses `"i/n"`, the `--param shard=2/7` surface. `i` is
40    /// 0-based, as in `draw::shard_owns`, so the whole draw is `0/1` and `1/1`
41    /// is refused. The error is the message the operator sees, naming the
42    /// value and what is wrong with it.
43    pub fn parse(s: &str) -> std::result::Result<Self, String> {
44        let (i, n) = s
45            .split_once('/')
46            .ok_or_else(|| format!("shard {s:?} is not `index/count`, e.g. `2/4`"))?;
47        let index: usize = i
48            .trim()
49            .parse()
50            .map_err(|_| format!("shard index {i:?} is not a number"))?;
51        let count: usize = n
52            .trim()
53            .parse()
54            .map_err(|_| format!("shard count {n:?} is not a number"))?;
55        if count == 0 {
56            return Err("shard count must be at least 1".to_string());
57        }
58        if index >= count {
59            return Err(format!(
60                "shard index {index} is out of range for {count} shards — \
61                 indices are 0-based, so the last is {}",
62                count - 1
63            ));
64        }
65        Ok(Self { index, count })
66    }
67}
68
69pub fn load(path: &Path, spec: &DrawSpec) -> Result<Vec<Sample>> {
70    load_shard(path, spec, None)
71}
72
73/// 2026-09-26: Loads one shard of the draw, or all of it when `shard` is
74/// `None`. The plan decides how many rows of each subset the draw takes, and
75/// the shard then keeps every `count`-th of those; filtering before the plan
76/// would change which rows the draw contains. An empty selection is an error.
77pub fn load_shard(path: &Path, spec: &DrawSpec, shard: Option<Shard>) -> Result<Vec<Sample>> {
78    let text = std::fs::read_to_string(path).with_context(|| {
79        format!(
80            "reading {} — delete ~/.metrale/artifacts/bfcl to re-provision",
81            path.display()
82        )
83    })?;
84    let totals = totals_of(&text)?;
85    let plan: BTreeMap<String, usize> = draw::plan(spec, &totals).into_iter().collect();
86
87    let mut taken: BTreeMap<String, usize> = BTreeMap::new();
88    let mut out = Vec::new();
89    for (i, line) in text.lines().enumerate() {
90        if line.trim().is_empty() {
91            continue;
92        }
93        let sample: Sample = serde_json::from_str(line)
94            .with_context(|| format!("{}:{}: malformed sample", path.display(), i + 1))?;
95        let Some(limit) = plan.get(&sample.subset) else {
96            continue;
97        };
98        let count = taken.entry(sample.subset.clone()).or_insert(0);
99        if *count >= *limit {
100            continue;
101        }
102        // 2026-09-26: `*count` is this row's 0-based position within the drawn
103        // rows of its subset, which `shard_owns` strides over. It advances
104        // whether or not this shard keeps the row, so every shard agrees on
105        // each row's position.
106        let position = *count;
107        *count += 1;
108        if let Some(sh) = shard
109            && !draw::shard_owns(position, sh.index, sh.count)
110        {
111            continue;
112        }
113        out.push(sample);
114    }
115    if out.is_empty() {
116        match shard {
117            None => bail!("the draw selected no samples — check the categories and percentages"),
118            Some(sh) => bail!(
119                "shard {} of {} selected no samples — the draw is smaller than the \
120                 shard count, or the plan is empty",
121                sh.index,
122                sh.count
123            ),
124        }
125    }
126    // 2026-09-26: A stable sort: by subset name, then file order.
127    out.sort_by(|a, b| a.subset.cmp(&b.subset));
128    Ok(out)
129}
130
131/// 2026-09-26: Per-subset row counts, without parsing whole samples.
132pub fn totals(path: &Path) -> Result<BTreeMap<String, usize>> {
133    totals_of(&std::fs::read_to_string(path)?)
134}
135
136fn totals_of(text: &str) -> Result<BTreeMap<String, usize>> {
137    #[derive(Deserialize)]
138    struct SubsetOnly {
139        subset: String,
140    }
141    let mut totals = BTreeMap::new();
142    for (i, line) in text.lines().enumerate() {
143        if line.trim().is_empty() {
144            continue;
145        }
146        // 2026-09-26: `load_shard` counts before it parses, so this pass meets a
147        // corrupt line first and must name the line.
148        let row: SubsetOnly = serde_json::from_str(line)
149            .with_context(|| format!("line {}: malformed sample", i + 1))?;
150        *totals.entry(row.subset).or_insert(0) += 1;
151    }
152    Ok(totals)
153}
154
155#[cfg(test)]
156mod tests {
157    use super::*;
158
159    fn write(name: &str, rows: &[(&str, usize)]) -> std::path::PathBuf {
160        let p = std::env::temp_dir().join(format!(
161            "metrale-bfcl-ds-{name}-{}.jsonl",
162            std::process::id()
163        ));
164        let mut text = String::new();
165        for (subset, n) in rows {
166            for i in 0..*n {
167                text.push_str(&format!(
168                    r#"{{"subset":"{subset}","sample_id":"{subset}_{i}","messages":[],"tools":[],"tool_choice":"auto"}}"#
169                ));
170                text.push('\n');
171            }
172        }
173        std::fs::write(&p, text).unwrap();
174        p
175    }
176
177    #[test]
178    fn the_draw_takes_the_first_n_of_each_subset() {
179        let p = write(
180            "head",
181            &[
182                ("simple_python", 10),
183                ("simple_java", 10),
184                ("irrelevance", 10),
185            ],
186        );
187        let spec = DrawSpec {
188            categories: vec!["non_live".into()],
189            category_pct: [("non_live".to_string(), 50.0)].into_iter().collect(),
190            subset_floor: None,
191        };
192        let s = load(&p, &spec).unwrap();
193        assert_eq!(
194            s.len(),
195            10,
196            "50% of both selected subsets, and irrelevance is excluded"
197        );
198        let ids: Vec<&str> = s.iter().map(|x| x.sample_id.as_str()).collect();
199        assert_eq!(
200            ids,
201            vec![
202                "simple_java_0",
203                "simple_java_1",
204                "simple_java_2",
205                "simple_java_3",
206                "simple_java_4",
207                "simple_python_0",
208                "simple_python_1",
209                "simple_python_2",
210                "simple_python_3",
211                "simple_python_4"
212            ],
213            "head(n), not a random sample"
214        );
215    }
216
217    #[test]
218    fn totals_are_counted_without_parsing_whole_samples() {
219        let p = write("totals", &[]);
220        std::fs::write(
221            &p,
222            concat!(
223                "{\"subset\":\"multiple\"}\n",
224                "{\"subset\":\"multiple\"}\n",
225                "{\"subset\":\"multiple\"}\n",
226                "{\"subset\":\"live_simple\"}\n",
227                "{\"subset\":\"live_simple\"}\n",
228            ),
229        )
230        .unwrap();
231        let t = totals(&p).unwrap();
232        assert_eq!(t.len(), 2);
233        assert_eq!(t["multiple"], 3);
234        assert_eq!(t["live_simple"], 2);
235    }
236
237    #[test]
238    fn an_empty_draw_is_an_error_rather_than_a_zero_sample_run() {
239        let p = write("empty", &[("live_relevance", 4)]);
240        let err = load(&p, &DrawSpec::golden()).unwrap_err().to_string();
241        assert!(err.contains("selected no samples"), "{err}");
242    }
243
244    #[test]
245    fn a_malformed_line_names_its_line_number() {
246        let p = write("bad", &[("multiple", 1)]);
247        let mut text = std::fs::read_to_string(&p).unwrap();
248        text.push_str("{not json}\n");
249        std::fs::write(&p, text).unwrap();
250        let err = load(&p, &DrawSpec::full()).unwrap_err().to_string();
251        assert!(err.contains("line 2: malformed sample"), "{err}");
252    }
253}
metrale_bench/benchmarks/bfcl/dataset.rs

metrale_bench/benchmarks/bfcl/
dataset.rs
