Skip to content

Commit d20a48a

Browse files
Use memchr for wheel parsing (#10620)
## Summary Before: ``` wheelname_parsing/numpy-compatible time: [106.90 ns 107.86 ns 108.86 ns] thrpt: [402.97 MiB/s 406.71 MiB/s 410.39 MiB/s] change: time: [-4.5360% -3.7694% -2.9179%] (p = 0.00 < 0.05) thrpt: [+3.0056% +3.9170% +4.7515%] Performance has improved. wheelname_parsing/flyte-short-incompatible time: [81.439 ns 82.209 ns 83.015 ns] thrpt: [390.59 MiB/s 394.42 MiB/s 398.15 MiB/s] change: time: [+6.2344% +7.5385% +8.8928%] (p = 0.00 < 0.05) thrpt: [-8.1666% -7.0101% -5.8685%] Performance has regressed. wheelname_parsing/flyte-short-compatible time: [78.909 ns 79.456 ns 80.031 ns] thrpt: [357.49 MiB/s 360.08 MiB/s 362.57 MiB/s] change: time: [+3.2653% +4.1733% +5.1062%] (p = 0.00 < 0.05) thrpt: [-4.8582% -4.0061% -3.1620%] Performance has regressed. Found 1 outliers among 100 measurements (1.00%) 1 (1.00%) high mild wheelname_parsing/flyte-long-incompatible time: [353.35 ns 357.49 ns 361.77 ns] thrpt: [340.06 MiB/s 344.13 MiB/s 348.17 MiB/s] change: time: [+1.4846% +2.0228% +2.6504%] (p = 0.00 < 0.05) thrpt: [-2.5820% -1.9827% -1.4629%] Performance has regressed. Found 10 outliers among 100 measurements (10.00%) 5 (5.00%) high mild 5 (5.00%) high severe wheelname_parsing/flyte-long-compatible time: [256.47 ns 258.12 ns 260.17 ns] thrpt: [417.88 MiB/s 421.20 MiB/s 423.90 MiB/s] change: time: [+0.4079% +1.8252% +3.6270%] (p = 0.02 < 0.05) thrpt: [-3.5001% -1.7925% -0.4063%] Change within noise threshold. Found 1 outliers among 100 measurements (1.00%) 1 (1.00%) high severe ``` After: ``` wheelname_parsing_fastest/numpy-compatible time: [61.500 ns 61.904 ns 62.350 ns] thrpt: [703.60 MiB/s 708.66 MiB/s 713.32 MiB/s] change: time: [+0.9879% +1.4542% +1.9311%] (p = 0.00 < 0.05) thrpt: [-1.8945% -1.4334% -0.9782%] Change within noise threshold. Found 8 outliers among 100 measurements (8.00%) 6 (6.00%) high mild 2 (2.00%) high severe wheelname_parsing_fastest/flyte-short-incompatible time: [49.341 ns 49.538 ns 49.769 ns] thrpt: [651.50 MiB/s 654.54 MiB/s 657.16 MiB/s] change: time: [+5.8750% +6.3656% +6.8338%] (p = 0.00 < 0.05) thrpt: [-6.3967% -5.9847% -5.5490%] Performance has regressed. Found 17 outliers among 100 measurements (17.00%) 5 (5.00%) low severe 1 (1.00%) low mild 6 (6.00%) high mild 5 (5.00%) high severe wheelname_parsing_fastest/flyte-short-compatible time: [49.425 ns 49.789 ns 50.193 ns] thrpt: [570.01 MiB/s 574.63 MiB/s 578.86 MiB/s] change: time: [+5.1267% +5.7418% +6.3476%] (p = 0.00 < 0.05) thrpt: [-5.9687% -5.4300% -4.8767%] Performance has regressed. Found 7 outliers among 100 measurements (7.00%) 7 (7.00%) high mild wheelname_parsing_fastest/flyte-long-incompatible time: [295.81 ns 298.01 ns 301.04 ns] thrpt: [408.66 MiB/s 412.82 MiB/s 415.89 MiB/s] change: time: [+0.5553% +1.0842% +1.7059%] (p = 0.00 < 0.05) thrpt: [-1.6772% -1.0726% -0.5523%] Change within noise threshold. Found 4 outliers among 100 measurements (4.00%) 3 (3.00%) high mild 1 (1.00%) high severe wheelname_parsing_fastest/flyte-long-compatible time: [214.80 ns 216.10 ns 217.50 ns] thrpt: [499.87 MiB/s 503.10 MiB/s 506.15 MiB/s] change: time: [+0.9003% +1.3207% +1.8268%] (p = 0.00 < 0.05) thrpt: [-1.7940% -1.3035% -0.8923%] Change within noise threshold. Found 8 outliers among 100 measurements (8.00%) 7 (7.00%) high mild 1 (1.00%) high severe ``` (Ignore the percent changes; they're not relative to one another.) So it's like >40% faster for the small case and >15% faster for the large case.
1 parent 82c08b8 commit d20a48a

File tree

5 files changed

+84
-34
lines changed

5 files changed

+84
-34
lines changed

Cargo.lock

+1
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

crates/uv-distribution-filename/Cargo.toml

+1
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ uv-pep440 = { workspace = true }
2121
uv-platform-tags = { workspace = true }
2222
uv-small-str = { workspace = true }
2323

24+
memchr = { workspace = true }
2425
rkyv = { workspace = true, features = ["smallvec-1"] }
2526
serde = { workspace = true }
2627
smallvec = { workspace = true }

crates/uv-distribution-filename/src/lib.rs

+1
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ mod build_tag;
1313
mod egg;
1414
mod extension;
1515
mod source_dist;
16+
mod splitter;
1617
mod wheel;
1718

1819
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
/// A simple splitter that uses `memchr` to find the next delimiter.
2+
pub(crate) struct MemchrSplitter<'a> {
3+
memchr: memchr::Memchr<'a>,
4+
haystack: &'a str,
5+
offset: usize,
6+
}
7+
8+
impl<'a> MemchrSplitter<'a> {
9+
#[inline]
10+
pub(crate) fn split(haystack: &'a str, delimiter: u8) -> Self {
11+
Self {
12+
memchr: memchr::Memchr::new(delimiter, haystack.as_bytes()),
13+
haystack,
14+
offset: 0,
15+
}
16+
}
17+
}
18+
19+
impl<'a> Iterator for MemchrSplitter<'a> {
20+
type Item = &'a str;
21+
22+
#[inline(always)]
23+
#[allow(clippy::inline_always)]
24+
fn next(&mut self) -> Option<Self::Item> {
25+
match self.memchr.next() {
26+
Some(index) => {
27+
let start = self.offset;
28+
self.offset = index + 1;
29+
Some(&self.haystack[start..index])
30+
}
31+
None if self.offset < self.haystack.len() => {
32+
let start = self.offset;
33+
self.offset = self.haystack.len();
34+
Some(&self.haystack[start..])
35+
}
36+
None => None,
37+
}
38+
}
39+
40+
#[inline]
41+
fn size_hint(&self) -> (usize, Option<usize>) {
42+
// We know we'll return at least one item if there's remaining text.
43+
let min = usize::from(self.offset < self.haystack.len());
44+
45+
// Maximum possible splits is remaining length divided by 2 (minimum one char between delimiters).
46+
let max = (self.haystack.len() - self.offset + 1) / 2 + min;
47+
48+
(min, Some(max))
49+
}
50+
}

crates/uv-distribution-filename/src/wheel.rs

+31-34
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
use std::fmt::{Display, Formatter};
22

3+
use memchr::memchr;
34
use serde::{de, Deserialize, Deserializer, Serialize, Serializer};
45
use std::str::FromStr;
56
use thiserror::Error;
@@ -12,6 +13,7 @@ use uv_platform_tags::{
1213
PlatformTag, TagCompatibility, Tags,
1314
};
1415

16+
use crate::splitter::MemchrSplitter;
1517
use crate::{BuildTag, BuildTagError};
1618

1719
#[derive(
@@ -145,64 +147,66 @@ impl WheelFilename {
145147
// The wheel filename should contain either five or six entries. If six, then the third
146148
// entry is the build tag. If five, then the third entry is the Python tag.
147149
// https://www.python.org/dev/peps/pep-0427/#file-name-convention
148-
let mut parts = stem.split('-');
150+
let mut splitter = memchr::Memchr::new(b'-', stem.as_bytes());
149151

150-
let name = parts
151-
.next()
152-
.expect("split always yields 1 or more elements");
153-
154-
let Some(version) = parts.next() else {
152+
let Some(version) = splitter.next() else {
155153
return Err(WheelFilenameError::InvalidWheelFileName(
156154
filename.to_string(),
157155
"Must have a version".to_string(),
158156
));
159157
};
160158

161-
let Some(build_tag_or_python_tag) = parts.next() else {
159+
let Some(build_tag_or_python_tag) = splitter.next() else {
162160
return Err(WheelFilenameError::InvalidWheelFileName(
163161
filename.to_string(),
164162
"Must have a Python tag".to_string(),
165163
));
166164
};
167165

168-
let Some(python_tag_or_abi_tag) = parts.next() else {
166+
let Some(python_tag_or_abi_tag) = splitter.next() else {
169167
return Err(WheelFilenameError::InvalidWheelFileName(
170168
filename.to_string(),
171169
"Must have an ABI tag".to_string(),
172170
));
173171
};
174172

175-
let Some(abi_tag_or_platform_tag) = parts.next() else {
173+
let Some(abi_tag_or_platform_tag) = splitter.next() else {
176174
return Err(WheelFilenameError::InvalidWheelFileName(
177175
filename.to_string(),
178176
"Must have a platform tag".to_string(),
179177
));
180178
};
181179

182-
let (name, version, build_tag, python_tag, abi_tag, platform_tag) =
183-
if let Some(platform_tag) = parts.next() {
184-
if parts.next().is_some() {
180+
let (name, version, build_tag, python_tag, abi_tag, platform_tag, is_large) =
181+
if let Some(platform_tag) = splitter.next() {
182+
if splitter.next().is_some() {
185183
return Err(WheelFilenameError::InvalidWheelFileName(
186184
filename.to_string(),
187185
"Must have 5 or 6 components, but has more".to_string(),
188186
));
189187
}
190188
(
191-
name,
192-
version,
193-
Some(build_tag_or_python_tag),
194-
python_tag_or_abi_tag,
195-
abi_tag_or_platform_tag,
196-
platform_tag,
189+
&stem[..version],
190+
&stem[version + 1..build_tag_or_python_tag],
191+
Some(&stem[build_tag_or_python_tag + 1..python_tag_or_abi_tag]),
192+
&stem[python_tag_or_abi_tag + 1..abi_tag_or_platform_tag],
193+
&stem[abi_tag_or_platform_tag + 1..platform_tag],
194+
&stem[platform_tag + 1..],
195+
// Always take the slow path if a build tag is present.
196+
true,
197197
)
198198
} else {
199199
(
200-
name,
201-
version,
200+
&stem[..version],
201+
&stem[version + 1..build_tag_or_python_tag],
202202
None,
203-
build_tag_or_python_tag,
204-
python_tag_or_abi_tag,
205-
abi_tag_or_platform_tag,
203+
&stem[build_tag_or_python_tag + 1..python_tag_or_abi_tag],
204+
&stem[python_tag_or_abi_tag + 1..abi_tag_or_platform_tag],
205+
&stem[abi_tag_or_platform_tag + 1..],
206+
// Determine whether any of the tag types contain a period, which would indicate
207+
// that at least one of the tag types includes multiple tags (which in turn
208+
// necessitates taking the slow path).
209+
memchr(b'.', stem[build_tag_or_python_tag..].as_bytes()).is_some(),
206210
)
207211
};
208212

@@ -217,30 +221,23 @@ impl WheelFilename {
217221
})
218222
.transpose()?;
219223

220-
let tags = if build_tag.is_some()
221-
|| python_tag.contains('.')
222-
|| abi_tag.contains('.')
223-
|| platform_tag.contains('.')
224-
{
224+
let tags = if is_large {
225225
WheelTag::Large {
226226
large: Box::new(WheelTagLarge {
227227
build_tag,
228-
python_tag: python_tag
229-
.split('.')
228+
python_tag: MemchrSplitter::split(python_tag, b'.')
230229
.map(LanguageTag::from_str)
231230
.collect::<Result<_, _>>()
232231
.map_err(|err| {
233232
WheelFilenameError::InvalidLanguageTag(filename.to_string(), err)
234233
})?,
235-
abi_tag: abi_tag
236-
.split('.')
234+
abi_tag: MemchrSplitter::split(abi_tag, b'.')
237235
.map(AbiTag::from_str)
238236
.collect::<Result<_, _>>()
239237
.map_err(|err| {
240238
WheelFilenameError::InvalidAbiTag(filename.to_string(), err)
241239
})?,
242-
platform_tag: platform_tag
243-
.split('.')
240+
platform_tag: MemchrSplitter::split(platform_tag, b'.')
244241
.map(PlatformTag::from_str)
245242
.collect::<Result<_, _>>()
246243
.map_err(|err| {

0 commit comments

Comments
 (0)