use quick_xml::events::Event;
use quick_xml::reader::Reader;
use std::fs::File;
use std::io::{BufReader, Read};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::thread;
pub const ROW_CAP: usize = 2_000_000;
#[derive(Default)]
pub struct SheetData {
pub rows: Vec<Vec<String>>,
pub ncols: usize,
pub done: bool,
pub capped: bool,
}
pub struct StreamBook {
file: String,
sheets: Vec<(String, String)>, sst: Arc<Vec<String>>,
cur: usize,
data: Arc<Mutex<SheetData>>,
stop: Arc<AtomicBool>,
}
impl StreamBook {
pub fn open(file: &str) -> Result<Self, String> {
let f = File::open(file).map_err(|e| e.to_string())?;
let mut zip = zip::ZipArchive::new(f).map_err(|e| e.to_string())?;
let sheets = read_workbook_sheets(&mut zip)?;
let sst = Arc::new(read_shared_strings(&mut zip));
let mut book = StreamBook {
file: file.to_string(),
sheets,
sst,
cur: 0,
data: Arc::new(Mutex::new(SheetData::default())),
stop: Arc::new(AtomicBool::new(false)),
};
book.spawn();
Ok(book)
}
fn spawn(&mut self) {
self.stop.store(true, Ordering::Relaxed); self.stop = Arc::new(AtomicBool::new(false));
self.data = Arc::new(Mutex::new(SheetData::default()));
let file = self.file.clone();
let path = self.sheets[self.cur].1.clone();
let sst = self.sst.clone();
let data = self.data.clone();
let stop = self.stop.clone();
thread::spawn(move || stream_sheet(&file, &path, &sst, &data, &stop));
}
pub fn names(&self) -> Vec<String> {
self.sheets.iter().map(|(n, _)| n.clone()).collect()
}
pub fn selected(&self) -> usize {
self.cur
}
pub fn select(&mut self, idx: usize) {
if idx < self.sheets.len() && idx != self.cur {
self.cur = idx;
self.spawn();
}
}
pub fn dims(&self) -> (usize, usize, bool, bool) {
let d = self.data.lock().unwrap();
(d.rows.len(), d.ncols, d.done, d.capped)
}
pub fn find(&self, query: &str) -> Vec<(usize, usize)> {
let needle = query.to_ascii_lowercase();
let d = self.data.lock().unwrap();
let mut hits = Vec::new();
for (r, row) in d.rows.iter().enumerate() {
for (c, cell) in row.iter().enumerate() {
if contains_ci(cell, &needle) {
hits.push((r, c));
}
}
}
hits
}
pub fn window(&self, r0: usize, r1: usize, c0: usize, c1: usize) -> Vec<Vec<String>> {
let d = self.data.lock().unwrap();
let mut out = Vec::new();
for r in r0..r1.min(d.rows.len()) {
let row = &d.rows[r];
let mut line = Vec::with_capacity(c1.saturating_sub(c0));
for c in c0..c1 {
line.push(row.get(c).cloned().unwrap_or_default());
}
out.push(line);
}
out
}
}
impl Drop for StreamBook {
fn drop(&mut self) {
self.stop.store(true, Ordering::Relaxed);
}
}
pub fn contains_ci(hay: &str, needle: &str) -> bool {
let (h, n) = (hay.as_bytes(), needle.as_bytes());
if n.is_empty() {
return true;
}
if h.len() < n.len() {
return false;
}
for i in 0..=h.len() - n.len() {
if h[i..i + n.len()]
.iter()
.zip(n)
.all(|(a, b)| a.to_ascii_lowercase() == *b)
{
return true;
}
}
false
}
fn col_index(r: &[u8]) -> usize {
let mut n = 0usize;
for &b in r {
if b.is_ascii_alphabetic() {
n = n * 26 + (b.to_ascii_uppercase() - b'A' + 1) as usize;
} else {
break;
}
}
n.saturating_sub(1)
}
fn read_workbook_sheets(zip: &mut zip::ZipArchive<File>) -> Result<Vec<(String, String)>, String> {
let wb = read_entry(zip, "xl/workbook.xml").ok_or("missing xl/workbook.xml (not an xlsx?)")?;
let mut order: Vec<(String, String)> = Vec::new(); let mut rd = Reader::from_str(&wb);
let mut buf = Vec::new();
loop {
match rd.read_event_into(&mut buf) {
Ok(Event::Empty(e)) | Ok(Event::Start(e)) if e.name().as_ref() == b"sheet" => {
let mut name = String::new();
let mut rid = String::new();
for a in e.attributes().flatten() {
match a.key.as_ref() {
b"name" => name = String::from_utf8_lossy(&a.value).into_owned(),
b"r:id" => rid = String::from_utf8_lossy(&a.value).into_owned(),
_ => {}
}
}
order.push((name, rid));
}
Ok(Event::Eof) | Err(_) => break,
_ => {}
}
buf.clear();
}
let rels = read_entry(zip, "xl/_rels/workbook.xml.rels").unwrap_or_default();
let mut rid_to_path = std::collections::HashMap::new();
let mut rd = Reader::from_str(&rels);
let mut buf = Vec::new();
loop {
match rd.read_event_into(&mut buf) {
Ok(Event::Empty(e)) | Ok(Event::Start(e)) if e.name().as_ref() == b"Relationship" => {
let mut id = String::new();
let mut target = String::new();
for a in e.attributes().flatten() {
match a.key.as_ref() {
b"Id" => id = String::from_utf8_lossy(&a.value).into_owned(),
b"Target" => target = String::from_utf8_lossy(&a.value).into_owned(),
_ => {}
}
}
if !id.is_empty() {
rid_to_path.insert(id, target);
}
}
Ok(Event::Eof) | Err(_) => break,
_ => {}
}
buf.clear();
}
let sheets: Vec<(String, String)> = order
.into_iter()
.filter_map(|(name, rid)| {
rid_to_path.get(&rid).map(|t| {
let t = t.trim_start_matches('/');
let path = if t.starts_with("xl/") {
t.to_string()
} else {
format!("xl/{t}")
};
(name, path)
})
})
.collect();
if sheets.is_empty() {
return Err("no worksheets found".into());
}
Ok(sheets)
}
fn read_shared_strings(zip: &mut zip::ZipArchive<File>) -> Vec<String> {
let Some(xml) = read_entry(zip, "xl/sharedStrings.xml") else {
return Vec::new();
};
let mut out = Vec::new();
let mut rd = Reader::from_str(&xml);
let mut buf = Vec::new();
let mut cur = String::new();
let mut in_si = false;
let mut in_t = false;
loop {
match rd.read_event_into(&mut buf) {
Ok(Event::Start(e)) => match e.name().as_ref() {
b"si" => {
cur.clear();
in_si = true;
}
b"t" => in_t = true,
_ => {}
},
Ok(Event::Text(t)) if in_t && in_si => {
cur.push_str(&crate::util::xml_text(&t));
}
Ok(Event::GeneralRef(r)) if in_t && in_si => {
cur.push_str(&crate::util::xml_ref(&r));
}
Ok(Event::End(e)) => match e.name().as_ref() {
b"t" => in_t = false,
b"si" => {
out.push(std::mem::take(&mut cur));
in_si = false;
}
_ => {}
},
Ok(Event::Eof) | Err(_) => break,
_ => {}
}
buf.clear();
}
out
}
fn read_entry(zip: &mut zip::ZipArchive<File>, name: &str) -> Option<String> {
let f = zip.by_name(name).ok()?;
crate::util::read_to_string_capped(f, crate::util::MAX_DECODE_BYTES).ok()
}
pub fn preview_rows(
file: &str,
max_rows: usize,
max_cols: usize,
) -> Result<Vec<Vec<String>>, String> {
let f = File::open(file).map_err(|e| e.to_string())?;
let mut zip = zip::ZipArchive::new(f).map_err(|e| e.to_string())?;
let sheets = read_workbook_sheets(&mut zip)?;
let sst = read_shared_strings(&mut zip);
let sheet_path = sheets.first().ok_or("workbook has no sheets")?.1.clone();
let entry = zip.by_name(&sheet_path).map_err(|e| e.to_string())?;
let bounded = entry.take(crate::util::MAX_DECODE_BYTES as u64);
let mut rows: Vec<Vec<String>> = Vec::new();
parse_sheet_xml(
bounded,
&sst,
|| false,
|mut row| {
row.truncate(max_cols);
rows.push(row);
rows.len() >= max_rows
},
);
Ok(rows)
}
fn stream_sheet(
file: &str,
sheet_path: &str,
sst: &[String],
data: &Arc<Mutex<SheetData>>,
stop: &Arc<AtomicBool>,
) {
let Ok(f) = File::open(file) else { return };
let Ok(mut zip) = zip::ZipArchive::new(f) else {
return;
};
let Ok(entry) = zip.by_name(sheet_path) else {
return;
};
parse_sheet_xml(
entry,
sst,
|| stop.load(Ordering::Relaxed),
|row| {
let mut d = data.lock().unwrap();
d.ncols = d.ncols.max(row.len());
d.rows.push(row);
if d.rows.len() >= ROW_CAP {
d.capped = true;
d.done = true;
true
} else {
false
}
},
);
if !stop.load(Ordering::Relaxed) {
if let Ok(mut d) = data.lock() {
d.done = true;
}
}
}
fn parse_sheet_xml<R, A, F>(reader: R, sst: &[String], should_abort: A, mut on_row: F)
where
R: Read,
A: Fn() -> bool,
F: FnMut(Vec<String>) -> bool,
{
let mut rd = Reader::from_reader(BufReader::with_capacity(1 << 20, reader));
let mut buf = Vec::new();
let mut row: Vec<String> = Vec::new();
let mut col = 0usize;
let mut ctype: u8 = b'n'; let mut val = String::new();
let mut in_v = false;
let mut in_t = false;
loop {
if should_abort() {
return;
}
match rd.read_event_into(&mut buf) {
Ok(Event::Start(e)) | Ok(Event::Empty(e)) => match e.name().as_ref() {
b"row" => {
row = Vec::new();
col = 0;
}
b"c" => {
ctype = b'n';
let mut cref: Option<Vec<u8>> = None;
for a in e.attributes().flatten() {
match a.key.as_ref() {
b"r" => cref = Some(a.value.into_owned()),
b"t" => {
ctype = match a.value.as_ref() {
b"s" => b's',
b"inlineStr" => b'i',
b"str" => b'l',
b"b" => b'b',
_ => b'n',
}
}
_ => {}
}
}
if let Some(r) = cref {
col = col_index(&r);
}
while row.len() < col {
row.push(String::new());
}
val.clear();
}
b"v" => {
in_v = true;
val.clear();
}
b"t" => in_t = true,
_ => {}
},
Ok(Event::Text(t)) => {
if in_v || in_t {
val.push_str(&crate::util::xml_text(&t));
}
}
Ok(Event::GeneralRef(r)) => {
if in_v || in_t {
val.push_str(&crate::util::xml_ref(&r));
}
}
Ok(Event::End(e)) => match e.name().as_ref() {
b"v" => in_v = false,
b"t" => in_t = false,
b"c" => {
let resolved = match ctype {
b's' => val
.trim()
.parse::<usize>()
.ok()
.and_then(|i| sst.get(i).cloned())
.unwrap_or_default(),
b'b' => {
if val.trim() == "1" {
"TRUE".into()
} else {
"FALSE".into()
}
}
_ => val.clone(),
};
if row.len() <= col {
row.resize(col + 1, String::new());
}
row[col] = resolved;
col += 1;
}
b"row" if on_row(std::mem::take(&mut row)) => {
return;
}
_ => {}
},
Ok(Event::Eof) | Err(_) => break,
_ => {}
}
buf.clear();
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::time::Instant;
#[test]
#[ignore]
fn big_xlsx() {
let Ok(path) = std::env::var("SUCHER_BIG") else {
eprintln!("set SUCHER_BIG=/path/to/file.xlsx to run this benchmark");
return;
};
if !std::path::Path::new(&path).exists() {
eprintln!("SUCHER_BIG file not found: {path}");
return;
}
let t = Instant::now();
let book = StreamBook::open(&path).expect("open");
println!("open() (sst+workbook parse): {:?}", t.elapsed());
println!("sheets: {:?}", book.names());
let t = Instant::now();
loop {
let (rows, _, done, _) = book.dims();
if rows >= 1000 || done {
println!("time to first {rows} rows: {:?}", t.elapsed());
break;
}
}
let first = book.window(0, 3, 0, 5);
println!("first rows sample: {first:?}");
let t = Instant::now();
loop {
let (rows, ncols, done, capped) = book.dims();
if done {
println!(
"done: {rows} rows x {ncols} cols, capped={capped}, in {:?}",
t.elapsed()
);
break;
}
}
let t = Instant::now();
let hits = book.find("REACH");
println!("find(\"REACH\"): {} hits in {:?}", hits.len(), t.elapsed());
let t = Instant::now();
let hits = book.find("zzznotfoundzzz");
println!("find(miss): {} hits in {:?}", hits.len(), t.elapsed());
}
#[test]
fn ci() {
assert!(contains_ci("Hello World", "world"));
assert!(contains_ci("ABC", "abc"));
assert!(!contains_ci("abc", "xyz"));
assert!(contains_ci("anything", ""));
}
fn parse_all(xml: &str, sst: &[String]) -> Vec<Vec<String>> {
let mut rows = Vec::new();
parse_sheet_xml(
xml.as_bytes(),
sst,
|| false,
|row| {
rows.push(row);
false
},
);
rows
}
#[test]
fn parses_shared_inline_and_typed_cells() {
let sst = vec!["Alpha".to_string(), "Beta".to_string()];
let xml = "<worksheet><sheetData>\
<row r=\"1\"><c r=\"A1\" t=\"s\"><v>1</v></c>\
<c r=\"B1\" t=\"inlineStr\"><is><t>Inline</t></is></c></row>\
<row r=\"2\"><c r=\"A2\" t=\"n\"><v>42</v></c>\
<c r=\"B2\" t=\"b\"><v>1</v></c></row>\
</sheetData></worksheet>";
let rows = parse_all(xml, &sst);
assert_eq!(
rows,
vec![
vec!["Beta".to_string(), "Inline".to_string()],
vec!["42".to_string(), "TRUE".to_string()],
]
);
}
#[test]
fn parser_stops_when_on_row_signals_full() {
let xml = "<worksheet><sheetData>\
<row r=\"1\"><c r=\"A1\" t=\"n\"><v>1</v></c></row>\
<row r=\"2\"><c r=\"A2\" t=\"n\"><v>2</v></c></row>\
<row r=\"3\"><c r=\"A3\" t=\"n\"><v>3</v></c></row>\
</sheetData></worksheet>";
let mut rows = Vec::new();
parse_sheet_xml(
xml.as_bytes(),
&[],
|| false,
|row| {
rows.push(row);
rows.len() >= 2 },
);
assert_eq!(rows.len(), 2);
assert_eq!(rows[1], vec!["2".to_string()]);
}
#[test]
fn preview_rows_reads_and_caps() {
let g = preview_rows("samples/sample.xlsx", 2, 3).expect("first worksheet rows");
assert_eq!(g.len(), 2, "row cap");
assert!(g.iter().all(|r| r.len() <= 3), "col cap");
assert_eq!(g[0], vec!["Item", "Qty", "Unit"]);
assert_eq!(g[1], vec!["Coffee", "3", "4.5"]);
}
}