use rete_core::ingest::RawQuad;
use rete_core::Rete;
use std::path::Path;
use crate::commands::buildinfo;
use crate::commands::card::{self, CardArgs};
use crate::commands::progress;
use rete_core::build_progress::{self as live, Stat, Ticker, Unit};
const BATCH: usize = 10_000;
fn reader_share(budget: u64) -> u64 {
(budget / 4).clamp(64 << 20, 1 << 30)
}
fn open_input(path: &str, share: u64) -> anyhow::Result<Rete> {
crate::commands::range_source::open_local_ranged(path, Some(share))
}
fn stream_file<E>(
rete: &mut Rete,
label: &str,
share: u64,
visit: &mut dyn FnMut(RawQuad) -> Result<(), E>,
) -> Result<u64, E> {
let mut slots: Vec<Option<String>> = vec![None];
slots.extend(rete.graph_names().iter().map(|g| Some((*g).to_string())));
rete.set_memory_budget(Some(share));
let mut total: u64 = 0;
for slot in slots {
let mut cursor = 0u32;
loop {
let (triples, next, done) = rete.dump_batch(slot.as_deref(), cursor, BATCH);
for (s, p, o) in triples {
visit((s, p, o, slot.clone()))?;
total += 1;
}
cursor = next;
if done {
break;
}
}
if let Some(g) = &slot {
rete.release_named_graph(g);
}
}
eprintln!("merge: {label} contributed {total} quad(s)");
Ok(total)
}
fn hash_input(path: &str, tap: &buildinfo::InputTap) -> anyhow::Result<()> {
let started = std::time::Instant::now();
let file = std::fs::File::open(path).map_err(|e| anyhow::anyhow!("{path}: {e}"))?;
let mut reader = std::io::BufReader::with_capacity(1 << 20, progress::CountingRead(file));
let mut hasher = tap.hasher.borrow_mut();
let n = std::io::copy(&mut reader, &mut *hasher).map_err(|e| anyhow::anyhow!("{path}: {e}"))?;
tap.bytes.set(n);
eprintln!(
"merge: hashed {path} ({:.1} MB) in {:.1} s",
n as f64 / 1e6,
started.elapsed().as_secs_f64()
);
Ok(())
}
fn inputs_with_curated_cards(inputs: &[String]) -> Vec<String> {
inputs
.iter()
.filter(|path| {
crate::commands::range_source::LocalRangeReader::open(path)
.ok()
.and_then(|r| card::load_card_and_build_ranged(&r).ok())
.and_then(|read| read.card)
.is_some_and(|c| !card::CardInput::from_card(&c).is_empty())
})
.map(|p| buildinfo::file_name(p))
.collect()
}
pub(crate) fn merge_cmd(
inputs: &[String],
output: &str,
memory_budget_mb: u64,
tmp_dir: Option<&str>,
card_args: CardArgs,
measure_costs: bool,
) -> anyhow::Result<()> {
if inputs.is_empty() {
anyhow::bail!("merge needs at least one input .rete file");
}
for i in inputs {
if i == "-" {
anyhow::bail!(
"merge reads .rete files, not stdin — a .rete is read by range, not streamed"
);
}
if !Path::new(i).exists() {
anyhow::bail!("no such file: {i}");
}
}
if inputs.iter().any(|i| i == output) {
anyhow::bail!("refusing to overwrite an input ({output} is also an input)");
}
let tapped: Vec<(&str, &'static str)> = inputs.iter().map(|i| (i.as_str(), "rete")).collect();
let rec = std::rc::Rc::new(buildinfo::Recorder::start(&tapped[..]));
let mut perms_bits = rete_core::PermSet::CORE.bits();
let mut label_index = false;
let graphs = rete_core::GraphLayout::Quad;
let budget = memory_budget_mb.saturating_mul(1 << 20);
let share = reader_share(budget);
let mut statements = 0u64;
let mut named = false;
for path in inputs {
let rete = open_input(path, share)?;
statements += rete.header().quad_count;
named |= rete.header().has_quads();
perms_bits |= rete.header().perms.bits();
label_index |= rete.has_label_index();
}
let curated_inputs = inputs_with_curated_cards(inputs);
if !curated_inputs.is_empty() {
let what = if card_args.requested() {
"the merged card's curated fields come from --card-file and the curated flags only"
} else {
"the merged file has no card unless --card-file (or --card) gives it one"
};
eprintln!(
"merge: the cards of {} carry curated fields; a merge does not carry an input's \
card: {what}",
curated_inputs.join(", ")
);
}
let perms = rete_core::PermSet::from_bits(perms_bits).map_err(|e| anyhow::anyhow!("{e}"))?;
let curated = if card_args.requested() {
Some(card::load_curated(&card_args)?)
} else {
None
};
let flags = crate::commands::build::RecordFlags {
command: "merge",
format: None,
no_pyramid: true,
pyramid_algo: rete_core::PyramidAlgo::Louvain,
text_index: None,
materialize: false,
reason: false,
sections: crate::commands::build::BuildSections {
label_index,
schema: true,
},
perms,
graphs,
collapse_graphs: false,
strict: false,
surface: Default::default(),
type_predicate: None,
measure_costs,
memory_budget_mb: Some(memory_budget_mb),
card_args: &card_args,
};
let wants_card = curated.is_some();
let info = {
let mut info = buildinfo::new_build_info(flags.params(None));
rec.begin(&mut info);
std::rc::Rc::new(std::cell::RefCell::new(info))
};
let outcome = std::rc::Rc::new(crate::commands::build::CardOutcome::default());
let card_from_image = Some(crate::commands::build::full_card_hook(
info.clone(),
curated,
measure_costs,
memory_budget_mb,
rec.clone(),
rete_core::IriReport::default,
outcome.clone(),
"merge",
));
progress::plan(crate::commands::build::external_plan(
progress::Kind::Merge,
!rec.taps.is_empty(),
named,
false,
flags.sections,
wants_card,
measure_costs,
perms,
));
if !rec.taps.is_empty() {
let bytes = inputs
.iter()
.take(buildinfo::MAX_INPUTS_LISTED)
.map(|p| std::fs::metadata(p).map(|m| m.len()).unwrap_or(0))
.sum();
live::begin(
rete_core::build_trace::Phase::Read,
Some(bytes),
Unit::Bytes,
);
}
for (path, tap) in inputs
.iter()
.zip(&rec.taps)
.take(buildinfo::MAX_INPUTS_LISTED)
{
hash_input(path, tap)?;
}
if !rec.taps.is_empty() {
rec.lap(rete_core::build_trace::Phase::Read);
}
eprintln!("merge: {} input file(s) -> {output}", inputs.len());
if perms != rete_core::PermSet::ALL {
eprintln!(
"merge: every input is {}-permutation; the merged file will be too",
perms.len()
);
}
let out_path = Path::new(output).to_path_buf();
let read = rec.statements_read.clone();
let stats = rete_core::extbuild::build_external(
|visit| {
live::set_total(Some(statements), Unit::Statements);
let mut ticker = Ticker::new();
let mut counted = |q: RawQuad| {
read.set(read.get() + 1);
ticker.tick(1);
if read.get() & 0xFFF == 0 {
live::set_stat(Stat::Statements, read.get());
}
visit(q)
};
for path in inputs {
let mut rete = open_input(path, share).map_err(|e| {
rete_core::extbuild::ExtBuildError::Ingest(rete_core::ingest::IngestError::Io(
format!("{path}: {e}"),
))
})?;
stream_file(&mut rete, path, share, &mut counted)?;
if rete.index_incomplete() {
return Err(rete_core::extbuild::ExtBuildError::Ingest(
rete_core::ingest::IngestError::Io(format!(
"{path}: reads failed during the walk — the merge would be short"
)),
));
}
}
Ok(())
},
&out_path,
rete_core::extbuild::ExternalBuildOptions {
memory_budget: budget,
tmp_dir: tmp_dir.map(std::path::PathBuf::from),
metadata: Box::new(|_| Vec::new()),
build_info: Vec::new(),
perms,
label_index,
schema: true,
type_override: None,
text_index: None,
graphs,
input_memory: share,
card_from_image,
},
)
.map_err(|e| anyhow::anyhow!("{e}"))?;
crate::commands::build::finish_full_card(output, &info, &rec, &outcome)?;
let size = std::fs::metadata(&out_path).map(|m| m.len()).unwrap_or(0);
progress::finish(stats.statements as u64, output);
eprintln!(
"merge: wrote {output}: {} statement(s), {} term(s), {:.2} MB",
stats.statements,
stats.terms,
size as f64 / 1e6
);
Ok(())
}