nordisk/znippy-zoomies
osm-katana/src/xml_reader.rs
//! Two-pass XML reader via VTD index (xml_vtd).
//!
//! Pass 1 — parallel mmap scan, streamed to disk:
//! `build_elem_index_to_mmap` writes ElemIndex entries directly to a
//! mmap'd `.elem.idx` file — no Vec accumulation in RAM.
//! Sweden: 320 MB index stays warm in page cache.
//! Planet: 64 GB index streams through without touching heap.
//! Node records are then extracted in a second sequential pass over the index.
//!
//! Pass 2 — VTD-guided parallel parse:
//! ElemIndex filtered to ways+relations. `plan_chunks` groups entries into
//! ~4 MB batches with exact byte boundaries. Rayon distributes chunks across
//! all cores. Each worker parses its batch using the pre-known byte offsets —
//! no boundary scanning, no partial elements.
#![allow(
clippy::float_arithmetic,
clippy::arithmetic_side_effects,
reason = "e7 ↔ degree conversions and index arithmetic are intentional"
)]
#![allow(
unsafe_code,
reason = "find_ways_rels_offset mmaps the input read-only — safety documented at the call site"
)]
use std::sync::atomic::{AtomicU64, Ordering};
use std::{path::Path, sync::Arc};
use anyhow::Result;
use memchr::memchr;
use znippy_zoomies::gatling;
use crate::cancel::Checkpointed;
use crate::shared::gpu_backend::GpuBackend;
use crate::{
geometry,
node_store::{CoordFileWriter, NodeStore},
reader::{CollectedWay, MemberEntry, NodeRecord, RelationRecord, WayRecord},
tags,
writer::{self, NodeWriter, RelWriter, WayWriter},
xml_vtd::{self, ElemIndex, ElemKind},
};
// ~4 MB of element bytes per Rayon chunk (ways avg ~1 KB → ~4 000 ways/chunk)
const MAX_CHUNK_BYTES: usize = 4 * 1024 * 1024;
// ── Gatling typed pipeline for raw .osm → Parquet ────────────────────────────
/// Chunk/carry sizes for raw XML streaming through Gatling.
const RAW_XML_CHUNK_SIZE: usize = 200 * 1024 * 1024; // 200 MB per read
const RAW_XML_CARRY_HEADROOM: usize = 32 * 1024 * 1024; // 32 MB (largest element)
const RAW_XML_RING_SLOTS: usize = 6;
/// Like `find_top_level_start` but also matches `<changeset` and `<bound`.
/// Needed because the planet file starts with ~72 GB of changesets before any nodes.
fn find_any_top_level_start(bytes: &[u8], from: usize) -> usize {
let mut pos = from;
while pos < bytes.len() {
let Some(rel) = memchr::memchr(b'<', &bytes[pos..]) else {
break;
};
let lt = pos + rel;
let rest = bytes.get(lt + 1..).unwrap_or_default();
if rest.starts_with(b"node ")
|| rest.starts_with(b"node\t")
|| rest.starts_with(b"way ")
|| rest.starts_with(b"way\t")
|| rest.starts_with(b"relation ")
|| rest.starts_with(b"relation\t")
|| rest.starts_with(b"changeset ")
|| rest.starts_with(b"changeset\t")
|| rest.starts_with(b"bound ")
|| rest.starts_with(b"bound\t")
{
return lt;
}
pos = lt + 1;
}
bytes.len()
}
/// Byte offset of the first top-level `<way>`/`<relation>` element in an
/// uncompressed `.osm` file. OSM XML is ordered `bounds < nodes < ways <
/// relations`, so resolved pass 2 can start the reader here and never re-parse
/// the (node-heavy) node section. On a node-heavy planet the first `<way>` sits
/// ~80% into the file, so a single-threaded `memmem` scan reads ~80 GB on one
/// core — a serial trough that dragged the whole convert's core average down.
/// This scans the mmap in parallel (one gatling unit per file slice, global
/// minimum of the per-slice hits), all cores busy. Returns `0` when no
/// way/relation exists (pass 2 is then trivially empty but still correct) or on
/// any mmap error (safe fallback: parse all).
fn find_ways_rels_offset(path: &str) -> u64 {
find_first_offset(path, &[b"<way", b"<relation"])
}
/// Byte offset of the first top-level `<relation>` in an uncompressed `.osm` file —
/// where the relation clip prefilter's relations-only sub-pass starts, so it re-reads
/// none of the way section. `0` when there is none (or on any mmap error).
fn find_relations_offset(path: &str) -> u64 {
find_first_offset(path, &[b"<relation"])
}
/// The parallel `memmem` scan behind both offsets: the global minimum over every
/// slice of the first hit of any needle. `0` when nothing matches.
fn find_first_offset(path: &str, needles: &[&[u8]]) -> u64 {
use memchr::memmem;
let Ok(file) = std::fs::File::open(path) else {
return 0;
};
let Ok(mmap) = (unsafe { memmap2::Mmap::map(&file) }) else {
return 0;
};
let bytes: &[u8] = &mmap;
let n = bytes.len();
if n == 0 {
return 0;
}
// `<way`/`<relation` prefixes cover every delimiter form (`<way `, `<way>`,
// `<way\n`, …) since OSM has no other tag starting with these — one scan each
// instead of five per tag. `max_needle` bytes of inter-slice overlap so a
// needle straddling a boundary is still found by the earlier slice.
let max_needle = needles.iter().map(|n| n.len()).max().unwrap_or(1);
let n_workers = std::thread::available_parallelism().map_or(4, |w| w.get());
let units = (n_workers * 8).clamp(1, n);
let chunk = n.div_ceil(units);
let hits = crate::par::par_map(units, |i| {
let start = i * chunk;
if start >= n {
return usize::MAX;
}
let end = ((i + 1) * chunk + max_needle - 1).min(n);
let win = &bytes[start..end];
needles
.iter()
.filter_map(|nd| memmem::find(win, nd))
.map(|off| start + off)
.min()
.unwrap_or(usize::MAX)
});
hits.into_iter()
.min()
.filter(|&m| m != usize::MAX)
.map_or(0, |m| m as u64)
}
/// Split `data` into N worker segments at XML element boundaries.
/// Shared by all typed codecs in this module (and the typed PBF path).
pub(crate) fn xml_split_typed(
data: &[u8],
n_workers: usize,
is_last: bool,
) -> Option<gatling::Split<(usize, usize)>> {
let consumed = if is_last {
data.len()
} else {
xml_vtd::find_safe_slot_end(data)
};
if consumed == 0 {
return None;
}
// Oversubscribe: cut each slot into ~SPLIT× more segments than workers so the
// engine self-dispatches (a worker that drains its unit steals the next) —
// uneven node density can't strand a core, and finer units keep the reader→
// worker→collector pipeline full. Row-group size is unaffected: the per-worker
// NODE_ACC thread-local coalesces across a worker's segments into full groups.
// (Same trick the VTD pipelined reader uses; see vtd.rs SPLIT.)
const SPLIT: usize = 8;
let n = n_workers.max(1) * SPLIT;
let chunk_size = (consumed / n).max(1);
let mut borders: Vec<usize> = Vec::with_capacity(n + 1);
for i in 0..n {
borders.push(find_any_top_level_start(data, i * chunk_size));
}
borders.push(consumed);
let segments: Vec<(usize, usize)> = (0..n)
.filter_map(|i| {
let s = borders[i];
let e = borders[i + 1];
if s < e { Some((s, e)) } else { None }
})
.collect();
if segments.is_empty() {
return None;
}
Some(gatling::Split { segments, consumed })
}
/// TypedCodec for raw XML: split at element boundaries, workers VTD-parse → records.
///
/// PRIOR SERIAL FALLBACK (superseded 2026-07-14 by [`RawXmlParallelCodec`] +
/// [`RawParallelParquetSink`], which encode row groups on the workers). Retained
/// per the additive/no-regression law as the reference serial encoder; not on the
/// live `read_raw` path.
#[allow(dead_code)]
struct RawXmlTypedCodec {
tables: crate::Tables,
}
#[allow(dead_code)] // consumed only by the retained serial `RawXmlTypedCodec` fallback
type RawRecords = (Vec<NodeRecord>, Vec<WayRecord>, Vec<RelationRecord>);
impl gatling::TypedCodec for RawXmlTypedCodec {
type Seg = (usize, usize);
type Output = RawRecords;
fn split(
&self,
data: &[u8],
n_workers: usize,
is_last: bool,
) -> Option<gatling::Split<Self::Seg>> {
xml_split_typed(data, n_workers, is_last)
}
fn transform(&self, data: &[u8], seg: &Self::Seg) -> Self::Output {
let sub = &data[seg.0..seg.1];
let base = seg.0;
let mut nodes: Vec<NodeRecord> = Vec::new();
let mut ways: Vec<WayRecord> = Vec::new();
let mut rels: Vec<RelationRecord> = Vec::new();
xml_vtd::build_elem_index_slice(sub, base, &mut |e| {
let lo = (e.file_offset as usize).saturating_sub(base);
let hi = lo.saturating_add(e.file_length as usize).min(sub.len());
let b = sub.get(lo..hi).unwrap_or_default();
match e.kind {
ElemKind::Node if self.tables.nodes => {
let lat = e.lat_e7 as f32 / 1e7_f32;
let lon = e.lon_e7 as f32 / 1e7_f32;
nodes.push(NodeRecord {
id: e.id,
lon_lat: (f64::from(lon), f64::from(lat)),
tags_json: parse_tags_json(b),
version: parse_attr_i32(opening_tag(b), b"version"),
});
}
ElemKind::Way if self.tables.ways => ways.push(parse_way_raw(&e, b)),
ElemKind::Relation if self.tables.rels => rels.push(parse_relation(&e, b)),
_ => {}
}
});
(nodes, ways, rels)
}
}
/// TypedSink: pushes records to three Parquet writers on the collector thread —
/// Arrow-build + zstd compression are ALL serial here. This is the former
/// `read_raw` sink (the ~1.8-core all-core gate). Superseded 2026-07-14 by the
/// worker-encode [`RawParallelParquetSink`]; kept as the reference serial
/// fallback per the additive/no-regression law.
#[allow(dead_code)]
struct RawParquetTypedSink {
node_writer: Option<NodeWriter>,
way_writer: Option<WayWriter>,
rel_writer: Option<RelWriter>,
node_count: usize,
way_count: usize,
rel_count: usize,
pb: indicatif::ProgressBar,
}
impl gatling::TypedSink<RawRecords> for RawParquetTypedSink {
fn process(&mut self, output: RawRecords, _is_last: bool) -> Result<()> {
let (nodes, ways, rels) = output;
self.node_count += nodes.len();
self.way_count += ways.len();
self.rel_count += rels.len();
let progress = nodes.len() + ways.len() + rels.len();
if let Some(w) = &mut self.node_writer {
for r in &nodes {
w.push(r)?;
}
}
if let Some(w) = &mut self.way_writer {
for r in &ways {
w.push(r)?;
}
}
if let Some(w) = &mut self.rel_writer {
for r in &rels {
w.push(r)?;
}
}
self.pb.inc(progress as u64);
Ok(())
}
fn finish(&mut self) -> Result<()> {
if let Some(w) = self.node_writer.take() {
w.finish()?;
}
if let Some(w) = self.way_writer.take() {
w.finish()?;
}
if let Some(w) = self.rel_writer.take() {
w.finish()?;
}
Ok(())
}
}
/// Raw single-pass `.osm` → GeoParquet on the Gatling typed engine.
///
/// Streams the file through Gatling's ring buffer. N workers VTD-parse their
/// element-aligned segments AND Arrow-build + zstd-encode their own node/way/rel
/// Parquet row groups (parallel), exactly like the bz2/gz raw paths; the
/// collector ([`RawParallelParquetSink`]) only stitches the pre-compressed groups
/// in order. No mmap, no pre-built index — single sequential pass, all cores hot.
///
/// Before 2026-07-14 the collector compressed every row group on a single thread
/// (`RawParquetTypedSink`), pegging the all-core convert at ~1.8 cores despite the
/// parallel parse; moving the encode onto the workers lifts core-busy toward the
/// core-saturation-law target. Row VALUES and counts are unchanged (row ORDER
/// across segments coalesces per-worker, exactly as the shipped bz2/gz raw paths
/// already do — Parquet row order is not semantically meaningful).
///
/// Returns `(node_count, way_count, rel_count)`.
#[allow(clippy::too_many_arguments)]
pub fn read_raw(
path: &str,
_idx_path: &Path,
n_workers: usize,
pb: indicatif::ProgressBar,
out_dir: &Path,
compression: &str,
tables: crate::Tables,
cancel: &crate::CancelToken,
) -> Result<(usize, usize, usize)> {
let file = std::fs::File::open(path)?;
let reader = std::io::BufReader::with_capacity(8 * 1024 * 1024, file);
let (mut sink, encoders) = raw_typed_sink(out_dir, compression, tables, pb)?;
let codec = RawXmlParallelCodec {
tables,
node_encoder: encoders.node,
way_encoder: encoders.way,
rel_encoder: encoders.rel,
};
let cfg = gatling::Config {
chunk_size: RAW_XML_CHUNK_SIZE,
carry_headroom: RAW_XML_CARRY_HEADROOM,
ring_slots: RAW_XML_RING_SLOTS,
initial_carry: Vec::new(),
// Big: long sequential XML read, slots reused full → no resize churn.
slot_fill: gatling::SlotFill::Big,
};
gatling::run_typed(
reader,
codec,
&mut Checkpointed::new(&mut sink, cancel, "pass1"),
n_workers,
cfg,
)?;
sink.finish()
}
// ── Resolved .osm → GeoParquet on the Gatling typed engine (2-pass) ──────────
// Pass-1 uses `RawXmlPass1Codec` + `Pass1TypedSink` (workers Arrow-build + zstd
// their own node row groups; the collector only stitches + write_raws coords) —
// the same parallel machinery as the bz2/gz paths. The former serial
// `Pass1ResolvedCodec`/`Pass1ResolvedSink` (per-record `NodeAccumulator` on the
// single collector) was removed after it was found to peg raw `.osm` convert at
// ~3 cores.
/// Resolved 2-pass `.osm` → GeoParquet on the Gatling typed engine.
///
/// Pass 1: workers VTD-parse nodes → nodes.parquet + node_coords.bin (sorted).
/// Pass 2: workers VTD-parse ways/relations, resolve coords against the store
/// → ways/relations.parquet. Two file reads; no elem.idx, no in-RAM accumulation.
/// Returns `(node_count, way_count, rel_count)`.
#[allow(clippy::too_many_arguments)]
pub fn read_resolved_osm(
path: &str,
_idx_path: &Path,
n_workers: usize,
pb: indicatif::ProgressBar,
out_dir: &Path,
compression: &str,
tables: crate::Tables,
backend: &dyn GpuBackend,
sink: Option<&crate::phase_log::PhaseSink>,
clip: Option<Arc<crate::Clip>>,
cancel: &crate::CancelToken,
) -> Result<(usize, usize, usize)> {
let mk_cfg = || gatling::Config {
chunk_size: RAW_XML_CHUNK_SIZE,
carry_headroom: RAW_XML_CARRY_HEADROOM,
ring_slots: RAW_XML_RING_SLOTS,
initial_carry: Vec::new(),
// Big: long sequential XML read (resolved 2-pass), slots reused full.
slot_fill: gatling::SlotFill::Big,
};
// ── Pass 1: nodes ─────────────────────────────────────────────────────────
// Parallel node encode: workers Arrow-build + zstd their own row groups
// (RawXmlPass1Codec → pass1_parse_decoded, the same template the bz2/gz paths
// use); the collector only stitches pre-compressed groups + write_raws coords.
pb.set_message("pass 1/2 (nodes)");
let (node_encoder, node_writer) = if tables.nodes {
let (enc, w) = writer::parallel_node_writer(&out_dir.join("nodes.parquet"), compression)?;
(Some(enc), Some(w))
} else {
(None, None)
};
let node_acc = if tables.nodes {
Some(writer::NodeAccumulator::new())
} else {
None
};
let mut p1 = Pass1TypedSink {
coord_writer: CoordFileWriter::new(&out_dir.join("node_coords.bin"))?,
node_acc,
node_encoder: node_encoder.clone(),
node_writer,
node_count: 0,
pb: pb.clone(),
prev_stub: Vec::new(),
tables,
clip: clip.clone(),
busy_ns: None,
};
{
let file = std::fs::File::open(path)?;
let reader = std::io::BufReader::with_capacity(8 * 1024 * 1024, file);
let codec = RawXmlPass1Codec {
tables,
node_encoder,
clip: clip.clone(),
};
gatling::run_typed(
reader,
codec,
&mut Checkpointed::new(&mut p1, cancel, "pass1"),
n_workers,
mk_cfg(),
)?;
}
let (coord_writer, node_count) = p1.finish()?;
pb.set_message("sorting node coords …");
cancel.check("sort_and_stree", 0)?;
let store = Arc::new(coord_writer.finalize()?);
cancel.check("pass2", 0)?;
// ── Pass 2: ways + relations ──────────────────────────────────────────────
pb.set_message(format!("pass 2/2 ({} nodes)", store.len()));
// This reader historically logged no phases at all. The resolve-loss audit
// (finding (a)) needs a pass-2 `done` line to carry the loss keys on EVERY
// resolved convert, so pass 2 now logs one; pass 1 stays as it was.
let p2_phase = sink.map(|s| {
s.phase(
"pass2",
&[
("workers", &n_workers.to_string()),
("store_nodes", &store.len().to_string()),
],
)
});
let unresolved_refs = Arc::new(AtomicU64::new(0));
let shortened_ways = Arc::new(AtomicU64::new(0));
let report = run_pass2(
out_dir,
compression,
&store,
backend,
tables,
&pb,
p2_phase.as_ref().map(|p| Arc::clone(&p.busy_ns)),
clip.is_some(),
cancel,
&unresolved_refs,
&shortened_ways,
|sub, p2, c| {
let mut file = std::fs::File::open(path)?;
// Skip the node section — resolved pass 2 only needs ways + relations,
// and the codec VTD-parses every element in its slice — and, for the
// relation sub-pass, the way section too: OSM XML is ordered
// `nodes < ways < relations`, so the reader starts at the first element
// it will write.
let start = match sub {
SubPass::Relations => find_relations_offset(path),
SubPass::Whole | SubPass::Ways => find_ways_rels_offset(path),
};
if start > 0 {
use std::io::{Seek, SeekFrom};
file.seek(SeekFrom::Start(start))?;
}
let reader = std::io::BufReader::with_capacity(8 * 1024 * 1024, file);
let codec = RawXmlPass2Codec {
tables: c.tables,
store: Arc::clone(&store),
backend,
way_enc: c.way_enc,
rel_enc: c.rel_enc,
clip_active: c.clip_active,
unresolved_refs: Arc::clone(&unresolved_refs),
shortened_ways: Arc::clone(&shortened_ways),
sieve: c.sieve,
};
gatling::run_typed(
reader,
codec,
&mut Checkpointed::new(p2, cancel, "pass2"),
n_workers,
mk_cfg(),
)?;
Ok(match sub {
SubPass::Relations => vec![("rel_pass_start_offset", start.to_string())],
SubPass::Whole | SubPass::Ways => Vec::new(),
})
},
)?;
if let Some(p) = p2_phase {
p.done(&pass2_done_fields(
out_dir,
&report,
&unresolved_refs,
&shortened_ways,
Vec::new(),
));
}
Ok((node_count, report.ways, report.rels))
}
/// Raw way parser — collects node-ref IDs and tags but does NOT resolve coords.
/// `geometry` is always `None`; the node-ID list is preserved in `node_refs`.
fn parse_way_raw(e: &ElemIndex, elem: &[u8]) -> WayRecord {
let version = parse_attr_i32(opening_tag(elem), b"version");
let mut refs: Vec<i64> = Vec::new();
let mut tag_pairs: Vec<(String, String)> = Vec::new();
scan_children(elem, |name, child| {
if name == b"nd" {
if let Some(r) = xml_vtd::find_attr(child, b"ref") {
refs.push(xml_vtd::parse_i64(r));
}
} else if name == b"tag" {
if let (Some(k), Some(v)) = (
xml_vtd::find_attr(child, b"k"),
xml_vtd::find_attr(child, b"v"),
) {
tag_pairs.push((unescape(k), unescape(v)));
}
}
});
let tags_json = tags::serialize_tags(tag_pairs.iter().map(|(k, v)| (k.as_str(), v.as_str())));
WayRecord {
id: e.id,
geometry: None,
tags_json,
node_refs: refs,
version,
}
}
// ── VTD-guided chunk planner ──────────────────────────────────────────────────
/// Group ElemIndex entries into batches of at most `max_bytes` total element
/// bytes. Returns `(start, end)` index pairs (exclusive end) into `index`.
/// Every element fits entirely within its batch — no spanning, no partial parses.
fn plan_chunks(index: &[&ElemIndex], max_bytes: usize) -> Vec<(usize, usize)> {
let mut chunks = Vec::new();
let mut start = 0usize;
let mut accum = 0usize;
for (i, e) in index.iter().enumerate() {
accum = accum.saturating_add(e.file_length as usize);
if accum >= max_bytes {
chunks.push((start, i.saturating_add(1)));
start = i.saturating_add(1);
accum = 0;
}
}
if start < index.len() {
chunks.push((start, index.len()));
}
chunks
}
// ── Per-chunk parallel parser ─────────────────────────────────────────────────
/// Way metadata before coord resolution — refs, tags, version, is_area.
struct WayMeta {
id: i64,
refs: Vec<i64>,
tags_json: String,
is_area: bool,
version: Option<i32>,
}
fn parse_chunk(
bytes: &[u8],
chunk: &[&ElemIndex],
store: &Arc<NodeStore>,
base_offset: usize,
clip_active: bool,
oracle: Option<&crate::reader::RelationClipOracle>,
) -> (
Vec<CollectedWay>,
Vec<RelationRecord>,
crate::reader::ResolveLoss,
crate::reader::RelSieve,
) {
// Pass A: parse every way's metadata (refs + tags), collect all refs in
// one flat Vec for a single batch lookup. Relations parsed inline (no
// coord resolution needed) and, under a clip, judged by the oracle as they
// are parsed — the same `RelationClipOracle::admit` the PBF decoder calls.
let mut way_metas: Vec<WayMeta> = Vec::new();
let mut rels: Vec<RelationRecord> = Vec::new();
let mut sieve = crate::reader::RelSieve::default();
let mut all_refs: Vec<i64> = Vec::new();
let mut ref_offsets: Vec<usize> = Vec::with_capacity(chunk.len() + 1);
ref_offsets.push(0);
for e in chunk {
let start = (e.file_offset as usize).saturating_sub(base_offset);
let end = start.saturating_add(e.file_length as usize);
let elem = bytes.get(start..end).unwrap_or_default();
match e.kind {
ElemKind::Way => {
let meta = parse_way_meta(e, elem);
all_refs.extend_from_slice(&meta.refs);
ref_offsets.push(all_refs.len());
way_metas.push(meta);
}
ElemKind::Relation => {
let (members, tags_json, version) = parse_relation_parts(elem);
let kept = crate::reader::RelationClipOracle::admit(
oracle,
store,
&members,
|| RelationRecord {
id: e.id,
tags_json,
members_json: MemberEntry::json(&members),
version,
},
&mut sieve,
);
if let Some(r) = kept {
rels.push(r);
}
}
ElemKind::Node => {}
}
}
// Pass B: one batch lookup for every ref in the chunk. STree64Mmap sorts
// by leaf-block index and madvise()s the touched range — random page
// faults across the 144 GB coord mmap become a near-sequential walk
// with parallel kernel readahead.
let resolved = store.lookup_batch(&all_refs);
// Pass C: assemble CollectedWay using sliced coord results per way.
let mut ways: Vec<CollectedWay> = Vec::with_capacity(way_metas.len());
let mut loss = crate::reader::ResolveLoss::default();
for (w_idx, meta) in way_metas.into_iter().enumerate() {
let lo = ref_offsets[w_idx];
let hi = ref_offsets[w_idx + 1];
let coords: Vec<(f64, f64)> = resolved[lo..hi]
.iter()
.filter_map(|opt| opt.map(|(lat, lon)| (f64::from(lon), f64::from(lat))))
.collect();
// The XML twin of the PBF resolve seam (audit finding (a)): a ref the
// store cannot resolve WITHOUT a clip is a silently missing vertex —
// count it through the one shared rule so this path's convert summary
// can say so too (it was the disclosed uncounted gap of the first cut).
loss.note_way(hi - lo, coords.len(), clip_active);
// Region-clip way semantics: store holds only in-region nodes, so a way
// with zero resolvable coords is wholly outside the region → drop it.
if clip_active && coords.is_empty() {
continue;
}
ways.push(CollectedWay {
id: meta.id,
coords,
tags_json: meta.tags_json,
node_refs: meta.refs,
version: meta.version,
is_area: meta.is_area,
});
}
(ways, rels, loss, sieve)
}
// ── Element parsers ───────────────────────────────────────────────────────────
/// Parse a way's metadata without resolving node coords. Caller batches refs
/// across all ways in a slot/chunk and does one `store.lookup_batch` for the
/// whole batch — turns N serial mmap page faults into a sorted, prefetched
/// sequential walk.
fn parse_way_meta(e: &ElemIndex, elem: &[u8]) -> WayMeta {
let version = parse_attr_i32(opening_tag(elem), b"version");
let mut refs: Vec<i64> = Vec::new();
let mut tag_pairs: Vec<(String, String)> = Vec::new();
scan_children(elem, |name, child| {
if name == b"nd" {
if let Some(r) = xml_vtd::find_attr(child, b"ref") {
refs.push(xml_vtd::parse_i64(r));
}
} else if name == b"tag" {
if let (Some(k), Some(v)) = (
xml_vtd::find_attr(child, b"k"),
xml_vtd::find_attr(child, b"v"),
) {
tag_pairs.push((unescape(k), unescape(v)));
}
}
});
let tags_json = tags::serialize_tags(tag_pairs.iter().map(|(k, v)| (k.as_str(), v.as_str())));
// `is_closed_area` reads the pairs it is given; it used to take a
// `HashMap<String, String>`, so this line was a full clone of every tag
// String on every way (stinky `katana-5`; guarded by
// `tag_map_alloc_law::way_meta_does_not_clone_the_tag_list`).
let is_area = geometry::is_closed_area(&refs, &tag_pairs);
WayMeta {
id: e.id,
refs,
tags_json,
is_area,
version,
}
}
fn parse_relation(e: &ElemIndex, elem: &[u8]) -> RelationRecord {
let (members, tags_json, version) = parse_relation_parts(elem);
RelationRecord {
id: e.id,
tags_json,
members_json: MemberEntry::json(&members),
version,
}
}
/// A `<relation>`'s members, tags (already serialised) and version — the parts a
/// verdict needs BEFORE the members are serialised, so a dropped relation costs no
/// JSON (`parse_chunk`); [`parse_relation`] is the same parse with the record built.
fn parse_relation_parts(elem: &[u8]) -> (Vec<MemberEntry>, String, Option<i32>) {
let version = parse_attr_i32(opening_tag(elem), b"version");
let mut tag_pairs: Vec<(String, String)> = Vec::new();
let mut members: Vec<MemberEntry> = Vec::new();
scan_children(elem, |name, child| {
if name == b"tag" {
if let (Some(k), Some(v)) = (
xml_vtd::find_attr(child, b"k"),
xml_vtd::find_attr(child, b"v"),
) {
tag_pairs.push((unescape(k), unescape(v)));
}
} else if name == b"member" {
let member_type = xml_vtd::find_attr(child, b"type")
.map(|t| match t {
b"node" => "node",
b"way" => "way",
b"relation" => "relation",
_ => "node",
})
.unwrap_or("node");
let member_ref = xml_vtd::find_attr(child, b"ref")
.map(xml_vtd::parse_i64)
.unwrap_or(0);
let role = xml_vtd::find_attr(child, b"role")
.map(unescape)
.unwrap_or_default();
members.push(MemberEntry {
member_type,
member_ref,
role,
});
}
});
let tags_json = tags::serialize_tags(tag_pairs.iter().map(|(k, v)| (k.as_str(), v.as_str())));
(members, tags_json, version)
}
// ── XML element helpers ───────────────────────────────────────────────────────
/// Return the bytes between `<` and the first `>` in `elem` (the opening tag
/// attributes, without the angle brackets). Empty slice on malformed input.
fn opening_tag(elem: &[u8]) -> &[u8] {
let start = if elem.first() == Some(&b'<') { 1 } else { 0 };
match memchr(b'>', elem) {
Some(end) => elem.get(start..end).unwrap_or_default(),
None => elem.get(start..).unwrap_or_default(),
}
}
/// Call `f(name, child_tag)` for every direct child tag inside `elem`.
/// `name` — bytes of the tag name (e.g. `b"nd"`, `b"tag"`, `b"member"`)
/// `child_tag` — all bytes between `<` and `>` (exclusive), i.e. attributes
///
/// Stops at the first closing tag (`</`).
fn scan_children(elem: &[u8], mut f: impl FnMut(&[u8], &[u8])) {
// skip past the opening tag's `>`
let Some(first_close) = memchr(b'>', elem) else {
return;
};
let mut pos = first_close.saturating_add(1);
while let Some(rel) = memchr(b'<', elem.get(pos..).unwrap_or_default()) {
let child_open = pos.saturating_add(rel);
let tag_start = child_open.saturating_add(1);
let first = elem.get(tag_start).copied();
if first == Some(b'/') {
break; // closing tag — done
}
if first == Some(b'?') || first == Some(b'!') {
pos = tag_start; // skip PI / comment
continue;
}
let Some(rel2) = memchr(b'>', elem.get(tag_start..).unwrap_or_default()) else {
break;
};
let child_close = tag_start.saturating_add(rel2);
let child_tag = elem.get(tag_start..child_close).unwrap_or_default();
let name_end = memchr(b' ', child_tag).unwrap_or(child_tag.len());
let name = child_tag.get(..name_end).unwrap_or_default();
// strip self-closing `/` from end of child_tag before passing to f
let attr_end = if child_tag.last() == Some(&b'/') {
child_tag.len().saturating_sub(1)
} else {
child_tag.len()
};
let child_attrs = child_tag.get(..attr_end).unwrap_or(child_tag);
f(name, child_attrs);
pos = child_close.saturating_add(1);
}
}
/// Return the `name="..."` attribute value as an i32, or `None`.
fn parse_attr_i32(tag: &[u8], name: &[u8]) -> Option<i32> {
let raw = xml_vtd::find_attr(tag, name)?;
let s = std::str::from_utf8(raw).ok()?;
s.parse::<i32>().ok()
}
/// Scan child `<tag k="..." v="..."/>` elements and return a JSON object string.
fn parse_tags_json(elem: &[u8]) -> String {
let mut pairs: Vec<(String, String)> = Vec::new();
scan_children(elem, |name, child| {
if name == b"tag" {
if let (Some(k), Some(v)) = (
xml_vtd::find_attr(child, b"k"),
xml_vtd::find_attr(child, b"v"),
) {
pairs.push((unescape(k), unescape(v)));
}
}
});
tags::serialize_tags(pairs.iter().map(|(k, v)| (k.as_str(), v.as_str())))
}
// ── Streaming passes — for compressed inputs (.osm.bz2) ──────────────────────
//
// Both passes use the same pipelined reader loop (same as build_elem_index_pipelined)
// but with a generic Read source instead of a File.
//
// Key pattern: output variables are declared BEFORE thread::scope so the scope
// closure (which runs as the processor on the calling thread) can mutate them.
// Only the I/O reader is spawned as a thread; the processor stays on the caller.
// ── Raw bz2 → Parquet on the Gatling engine (fast path) ──────────────────────
// ── Raw single-pass → Parquet on the Gatling TYPED engine ────────────────────
//
// Each worker decodes its segment, VTD-parses the middle, and:
// • nodes → encodes Parquet row groups in-worker (NodeColumnEncoder, like
// resolved pass-1) so zstd runs in parallel across workers;
// • ways → collected as raw (unresolved) WayRecords;
// • rels → collected as RelationRecords.
// The collector ([`RawParallelParquetSink`]) only stitches the pre-encoded node row
// groups (ParallelNodeWriter) and ships way/rel batches to a single ParquetSink
// writer thread — so the parse+encode is no longer serial on the collector.
//
// Boundary stubs (`prev_stub + left_stub`) are parsed on the collector; volume
// is a handful of elements per segment boundary, negligible vs the middle.
/// One worker's output for one raw segment: pre-encoded node/way/relation row
/// groups + the partial edge bytes that straddle the boundary.
///
/// Ways and relations are now Arrow-built + zstd-encoded **on the worker** (via
/// the shared [`WAY_ACC`]/[`REL_ACC`] accumulators and `pass2_accumulate`),
/// exactly like nodes and exactly like the resolved 2-pass path — so way/rel
/// compression runs across all Gatling workers instead of funnelling through one
/// `ParquetSink` writer thread. The collector only stitches the pre-encoded row
/// groups in order (plus the tiny boundary-stub groups it encodes itself).
#[derive(Default)]
struct RawSegment {
node_row_groups: Vec<writer::EncodedRowGroup>,
node_count: usize,
way_row_groups: Vec<writer::EncodedRowGroup>,
rel_row_groups: Vec<writer::EncodedRowGroup>,
way_rows: usize,
rel_rows: usize,
left_stub: Vec<u8>,
right_stub: Vec<u8>,
}
/// Parse a decoded XML run into a [`RawSegment`]. Nodes, ways and relations are
/// all Arrow-built + zstd-encoded in the worker: nodes via the per-worker
/// [`NODE_ACC`] (`node_encoder`), ways/rels via the shared [`WAY_ACC`]/[`REL_ACC`]
/// accumulators + `pass2_accumulate` (`way_encoder`/`rel_encoder`) — so the
/// collector never compresses. Shared by the bz2 and gz raw codecs.
fn raw_parse_decoded(
decoded: &[u8],
tables: crate::Tables,
node_encoder: Option<&writer::NodeColumnEncoder>,
way_encoder: Option<&writer::ColumnEncoder>,
rel_encoder: Option<&writer::ColumnEncoder>,
) -> RawSegment {
if decoded.is_empty() {
return RawSegment::default();
}
let first = xml_vtd::find_top_level_start(decoded, 0);
let last = xml_vtd::find_safe_slot_end(decoded);
let last = last.max(first);
let left_stub = decoded[..first].to_vec();
let right_stub = decoded[last..].to_vec();
let middle = &decoded[first..last];
let (node_row_groups, node_count, p2) =
raw_encode_middle(middle, tables, node_encoder, way_encoder, rel_encoder);
RawSegment {
node_row_groups,
node_count,
way_row_groups: p2.way_row_groups,
rel_row_groups: p2.rel_row_groups,
way_rows: p2.way_rows,
rel_rows: p2.rel_rows,
left_stub,
right_stub,
}
}
/// Parse an ALREADY element-aligned XML run (no partial elements at either end,
/// so both boundary stubs are empty) into a [`RawSegment`], sharing the exact
/// worker-side encode body with [`raw_parse_decoded`] via [`raw_encode_middle`].
///
/// Used by the uncompressed `.osm` single-pass codec ([`RawXmlParallelCodec`]):
/// its gatling split (`xml_split_typed`) already cuts on top-level element
/// boundaries — so, unlike the bz2/gz decoded buffers, there are never partial
/// boundary elements to carry as stubs. This is what lets the uncompressed path
/// encode node/way/rel Parquet row groups ON THE WORKERS (parallel zstd) instead
/// of on a single collector thread (the former ~1.8-core serial gate).
fn raw_parse_aligned(
sub: &[u8],
tables: crate::Tables,
node_encoder: Option<&writer::NodeColumnEncoder>,
way_encoder: Option<&writer::ColumnEncoder>,
rel_encoder: Option<&writer::ColumnEncoder>,
) -> RawSegment {
if sub.is_empty() {
return RawSegment::default();
}
let (node_row_groups, node_count, p2) =
raw_encode_middle(sub, tables, node_encoder, way_encoder, rel_encoder);
RawSegment {
node_row_groups,
node_count,
way_row_groups: p2.way_row_groups,
rel_row_groups: p2.rel_row_groups,
way_rows: p2.way_rows,
rel_rows: p2.rel_rows,
left_stub: Vec::new(),
right_stub: Vec::new(),
}
}
/// Shared worker-side encode of a run of complete top-level elements: nodes via
/// the per-worker [`NODE_ACC`] (`node_encoder`), ways/rels via the shared
/// [`WAY_ACC`]/[`REL_ACC`] accumulators + `pass2_accumulate`. Returns the sealed
/// full node row groups, the node count (counted at parse time), and the pass-2
/// segment (full way/rel groups + their row counts). Compresses NOTHING on the
/// caller's behalf beyond what these per-worker encoders do — the collector only
/// stitches. Shared verbatim by [`raw_parse_decoded`] (stub path, bz2/gz) and
/// [`raw_parse_aligned`] (no-stub path, uncompressed `.osm`).
fn raw_encode_middle(
middle: &[u8],
tables: crate::Tables,
node_encoder: Option<&writer::NodeColumnEncoder>,
way_encoder: Option<&writer::ColumnEncoder>,
rel_encoder: Option<&writer::ColumnEncoder>,
) -> (Vec<writer::EncodedRowGroup>, usize, P2Segment) {
let mut node_count = 0usize;
let mut ways: Vec<WayRecord> = Vec::new();
let mut rels: Vec<RelationRecord> = Vec::new();
let mut node_row_groups: Vec<writer::EncodedRowGroup> = Vec::new();
NODE_ACC.with(|cell| {
let mut slot = cell.borrow_mut();
let mut acc = if tables.nodes && node_encoder.is_some() {
Some(slot.get_or_insert_with(writer::NodeAccumulator::new))
} else {
None
};
let enc = node_encoder;
xml_vtd::build_elem_index_slice(middle, 0, &mut |e| {
let lo = e.file_offset as usize;
let hi = lo.saturating_add(e.file_length as usize).min(middle.len());
let b = middle.get(lo..hi).unwrap_or_default();
match e.kind {
ElemKind::Node if tables.nodes => {
node_count += 1;
if let (Some(acc), Some(enc)) = (acc.as_deref_mut(), enc) {
let lat = e.lat_e7 as f32 / 1e7_f32;
let lon = e.lon_e7 as f32 / 1e7_f32;
acc.push(&NodeRecord {
id: e.id,
lon_lat: (f64::from(lon), f64::from(lat)),
tags_json: parse_tags_json(b),
version: parse_attr_i32(opening_tag(b), b"version"),
});
if let Ok(Some(batch)) = acc.take_if_full() {
node_row_groups
.push(enc.encode(&batch).expect("node row group encode"));
}
}
}
ElemKind::Way if tables.ways => ways.push(parse_way_raw(&e, b)),
ElemKind::Relation if tables.rels => rels.push(parse_relation(&e, b)),
_ => {}
}
});
});
// Ways/rels: seal full ROW_GROUP_SIZE groups on THIS worker (parallel zstd)
// via the shared per-worker accumulators — identical to the resolved path.
let mut p2 = P2Segment::default();
pass2_accumulate(&ways, &rels, way_encoder, rel_encoder, &mut p2, None);
(node_row_groups, node_count, p2)
}
/// Drain the per-worker [`NODE_ACC`]/[`WAY_ACC`]/[`REL_ACC`] remainders into one
/// final row group each. Mirrors [`pass1_finish_worker`] + [`pass2_finish_worker`]
/// but returns a single [`RawSegment`] carrying all three tails.
fn raw_finish_worker(
tables: crate::Tables,
node_encoder: Option<&writer::NodeColumnEncoder>,
way_encoder: Option<&writer::ColumnEncoder>,
rel_encoder: Option<&writer::ColumnEncoder>,
) -> Option<RawSegment> {
let mut seg = RawSegment::default();
let mut any = false;
if tables.nodes {
if let Some(enc) = node_encoder {
NODE_ACC.with(|cell| {
if let Some(mut acc) = cell.borrow_mut().take() {
if let Ok(Some(batch)) = acc.take_remaining() {
// Nodes are counted at PARSE time (node_count += 1 per node
// in raw_parse_decoded), so the tail row group must NOT add
// to node_count again — only emit the pre-encoded group.
seg.node_row_groups
.push(enc.encode(&batch).expect("node tail row group encode"));
any = true;
}
}
});
}
}
// Ways/rels: reuse the resolved path's worker-drain helper verbatim.
if let Some(p2) = pass2_finish_worker(way_encoder, rel_encoder) {
seg.way_row_groups = p2.way_row_groups;
seg.rel_row_groups = p2.rel_row_groups;
seg.way_rows = p2.way_rows;
seg.rel_rows = p2.rel_rows;
any = true;
}
if any { Some(seg) } else { None }
}
/// Raw bz2 → Parquet [`gatling::TypedCodec`]: decode a bz2 block range + parse +
/// (node) encode, all on the worker.
struct RawBz2ParquetTypedCodec {
max_blocksize: u32,
tables: crate::Tables,
node_encoder: Option<writer::NodeColumnEncoder>,
way_encoder: Option<writer::ColumnEncoder>,
rel_encoder: Option<writer::ColumnEncoder>,
}
impl gatling::TypedCodec for RawBz2ParquetTypedCodec {
type Seg = (u64, u64);
type Output = RawSegment;
fn split(
&self,
data: &[u8],
n_workers: usize,
is_last: bool,
) -> Option<gatling::Split<(u64, u64)>> {
// ONE bz2 slot split for every bz2 codec in the crate (LAW 5).
crate::xml_to_pbf::bz2_split(data, n_workers, self.max_blocksize, is_last)
}
fn transform(&self, data: &[u8], &(start_bit, end_bit): &(u64, u64)) -> RawSegment {
let decoded = lbzip2::chunk::decode_segment(data, start_bit, end_bit, self.max_blocksize);
let nenc = if self.tables.nodes {
self.node_encoder.as_ref()
} else {
None
};
let wenc = if self.tables.ways {
self.way_encoder.as_ref()
} else {
None
};
let renc = if self.tables.rels {
self.rel_encoder.as_ref()
} else {
None
};
raw_parse_decoded(&decoded, self.tables, nenc, wenc, renc)
}
fn finish_worker(&self) -> Option<RawSegment> {
raw_finish_worker(
self.tables,
self.node_encoder.as_ref(),
self.way_encoder.as_ref(),
self.rel_encoder.as_ref(),
)
}
}
/// Raw gz → Parquet [`gatling::TypedCodec`]: DEFLATE decode + parse + (node/way/
/// rel) encode on the worker. Mirrors [`RawBz2ParquetTypedCodec`] with lgz splits.
struct RawGzParquetTypedCodec {
tables: crate::Tables,
node_encoder: Option<writer::NodeColumnEncoder>,
way_encoder: Option<writer::ColumnEncoder>,
rel_encoder: Option<writer::ColumnEncoder>,
}
impl gatling::TypedCodec for RawGzParquetTypedCodec {
type Seg = (usize, usize);
type Output = RawSegment;
fn split(
&self,
data: &[u8],
n_workers: usize,
is_last: bool,
) -> Option<gatling::Split<(usize, usize)>> {
crate::xml_to_pbf::gz_split(data, n_workers, is_last)
}
fn transform(&self, data: &[u8], seg: &(usize, usize)) -> RawSegment {
let decoded = crate::xml_to_pbf::gz_decode_segment(data, seg);
let nenc = if self.tables.nodes {
self.node_encoder.as_ref()
} else {
None
};
let wenc = if self.tables.ways {
self.way_encoder.as_ref()
} else {
None
};
let renc = if self.tables.rels {
self.rel_encoder.as_ref()
} else {
None
};
raw_parse_decoded(&decoded, self.tables, nenc, wenc, renc)
}
fn finish_worker(&self) -> Option<RawSegment> {
raw_finish_worker(
self.tables,
self.node_encoder.as_ref(),
self.way_encoder.as_ref(),
self.rel_encoder.as_ref(),
)
}
}
/// Raw single-pass `.osm` (UNCOMPRESSED XML) → Parquet [`gatling::TypedCodec`]:
/// each worker VTD-parses its element-aligned segment and Arrow-builds +
/// zstd-encodes its own node/way/rel row groups (parallel), exactly like the
/// bz2/gz raw codecs. Because `xml_split_typed` cuts on top-level element
/// boundaries there are no partial boundary elements — [`raw_parse_aligned`]
/// carries empty stubs — so the collector ([`RawParallelParquetSink`]) only
/// stitches pre-compressed groups. This replaced the former serial
/// `RawXmlTypedCodec` + `RawParquetTypedSink`, whose collector compressed every
/// row group on ONE thread and pegged the all-core raw convert at ~1.8 cores.
struct RawXmlParallelCodec {
tables: crate::Tables,
node_encoder: Option<writer::NodeColumnEncoder>,
way_encoder: Option<writer::ColumnEncoder>,
rel_encoder: Option<writer::ColumnEncoder>,
}
impl gatling::TypedCodec for RawXmlParallelCodec {
type Seg = (usize, usize);
type Output = RawSegment;
fn split(
&self,
data: &[u8],
n_workers: usize,
is_last: bool,
) -> Option<gatling::Split<Self::Seg>> {
xml_split_typed(data, n_workers, is_last)
}
fn transform(&self, data: &[u8], seg: &Self::Seg) -> RawSegment {
let nenc = if self.tables.nodes {
self.node_encoder.as_ref()
} else {
None
};
let wenc = if self.tables.ways {
self.way_encoder.as_ref()
} else {
None
};
let renc = if self.tables.rels {
self.rel_encoder.as_ref()
} else {
None
};
raw_parse_aligned(&data[seg.0..seg.1], self.tables, nenc, wenc, renc)
}
fn finish_worker(&self) -> Option<RawSegment> {
raw_finish_worker(
self.tables,
self.node_encoder.as_ref(),
self.way_encoder.as_ref(),
self.rel_encoder.as_ref(),
)
}
}
/// Raw single-pass [`gatling::TypedSink`]: stitches the workers' pre-encoded
/// node/way/relation row groups in order (NO compression on this thread) and
/// encodes only the tiny boundary-stub batches itself. Ways/rels are now encoded
/// on the workers (like nodes and like the resolved 2-pass path) instead of
/// funnelling through a single `ParquetSink` writer thread.
struct RawParallelParquetSink {
node_acc: Option<writer::NodeAccumulator>,
node_encoder: Option<writer::NodeColumnEncoder>,
node_writer: Option<writer::ParallelNodeWriter>,
node_count: usize,
/// Tiny collector-side accumulators for boundary-stub ways/rels parsed HERE
/// (a handful per segment boundary); full groups are encoded with
/// `way_encoder`/`rel_encoder` — the SAME encoders the workers hold.
way_acc: Option<writer::WayAccumulator>,
rel_acc: Option<writer::RelAccumulator>,
way_encoder: Option<writer::ColumnEncoder>,
rel_encoder: Option<writer::ColumnEncoder>,
/// In-order row-group stitchers (no compression): append both the workers'
/// bulk row groups and the collector's boundary-stub groups.
way_writer: Option<writer::ParallelRowGroupWriter>,
rel_writer: Option<writer::ParallelRowGroupWriter>,
way_count: usize,
rel_count: usize,
tables: crate::Tables,
prev_stub: Vec<u8>,
pb: indicatif::ProgressBar,
}
impl RawParallelParquetSink {
fn flush_stub(&mut self, stub: &[u8]) -> Result<()> {
if stub.is_empty() {
return Ok(());
}
let (inc_n, inc_w, inc_r) = (self.tables.nodes, self.tables.ways, self.tables.rels);
let mut nodes: Vec<NodeRecord> = Vec::new();
let mut ways: Vec<WayRecord> = Vec::new();
let mut rels: Vec<RelationRecord> = Vec::new();
xml_vtd::build_elem_index_slice(stub, 0, &mut |e| {
let lo = e.file_offset as usize;
let hi = lo.saturating_add(e.file_length as usize).min(stub.len());
let b = stub.get(lo..hi).unwrap_or_default();
match e.kind {
ElemKind::Node if inc_n => {
let lat = e.lat_e7 as f32 / 1e7_f32;
let lon = e.lon_e7 as f32 / 1e7_f32;
nodes.push(NodeRecord {
id: e.id,
lon_lat: (f64::from(lon), f64::from(lat)),
tags_json: parse_tags_json(b),
version: parse_attr_i32(opening_tag(b), b"version"),
});
}
ElemKind::Way if inc_w => ways.push(parse_way_raw(&e, b)),
ElemKind::Relation if inc_r => rels.push(parse_relation(&e, b)),
_ => {}
}
});
self.node_count += nodes.len();
if let (Some(acc), Some(enc), Some(w)) = (
self.node_acc.as_mut(),
self.node_encoder.as_ref(),
self.node_writer.as_mut(),
) {
for r in &nodes {
acc.push(r);
}
while let Some(batch) = acc.take_if_full()? {
w.append(enc.encode(&batch)?)?;
}
}
self.way_count += ways.len();
if let (Some(acc), Some(enc), Some(w)) = (
self.way_acc.as_mut(),
self.way_encoder.as_ref(),
self.way_writer.as_mut(),
) {
for r in &ways {
acc.push(r);
}
while let Some(batch) = acc.take_if_full()? {
w.append(enc.encode(&batch)?)?;
}
}
self.rel_count += rels.len();
if let (Some(acc), Some(enc), Some(w)) = (
self.rel_acc.as_mut(),
self.rel_encoder.as_ref(),
self.rel_writer.as_mut(),
) {
for r in &rels {
acc.push(r);
}
while let Some(batch) = acc.take_if_full()? {
w.append(enc.encode(&batch)?)?;
}
}
Ok(())
}
fn finish(mut self) -> Result<(usize, usize, usize)> {
let stub = std::mem::take(&mut self.prev_stub);
self.flush_stub(&stub)?;
if let (Some(acc), Some(enc), Some(w)) = (
self.node_acc.as_mut(),
self.node_encoder.as_ref(),
self.node_writer.as_mut(),
) {
if let Some(batch) = acc.take_remaining()? {
w.append(enc.encode(&batch)?)?;
}
}
if let Some(w) = self.node_writer.take() {
w.finish()?;
}
if let (Some(acc), Some(enc), Some(w)) = (
self.way_acc.as_mut(),
self.way_encoder.as_ref(),
self.way_writer.as_mut(),
) {
if let Some(batch) = acc.take_remaining()? {
w.append(enc.encode(&batch)?)?;
}
}
if let Some(w) = self.way_writer.take() {
w.finish()?;
}
if let (Some(acc), Some(enc), Some(w)) = (
self.rel_acc.as_mut(),
self.rel_encoder.as_ref(),
self.rel_writer.as_mut(),
) {
if let Some(batch) = acc.take_remaining()? {
w.append(enc.encode(&batch)?)?;
}
}
if let Some(w) = self.rel_writer.take() {
w.finish()?;
}
Ok((self.node_count, self.way_count, self.rel_count))
}
}
impl gatling::TypedSink<RawSegment> for RawParallelParquetSink {
fn process(&mut self, seg: RawSegment, _is_last: bool) -> Result<()> {
let boundary = [self.prev_stub.as_slice(), seg.left_stub.as_slice()].concat();
self.flush_stub(&boundary)?;
// Nodes / ways / rels: stitch the pre-encoded row groups in order (no
// zstd on this thread — the workers already compressed them).
self.node_count += seg.node_count;
if let Some(w) = self.node_writer.as_mut() {
for rg in seg.node_row_groups {
w.append(rg)?;
}
}
self.way_count += seg.way_rows;
if let Some(w) = self.way_writer.as_mut() {
for rg in seg.way_row_groups {
w.append(rg)?;
}
}
self.rel_count += seg.rel_rows;
if let Some(w) = self.rel_writer.as_mut() {
for rg in seg.rel_row_groups {
w.append(rg)?;
}
}
self.pb.inc(
(seg.left_stub.len() + seg.node_count * 64 + seg.way_rows * 512 + seg.right_stub.len())
as u64,
);
self.prev_stub = seg.right_stub;
Ok(())
}
}
/// The three column encoders the raw codec hands to the Gatling workers so each
/// worker Arrow-builds + zstd-encodes its own node/way/rel row groups in parallel.
#[derive(Default, Clone)]
struct RawEncoders {
node: Option<writer::NodeColumnEncoder>,
way: Option<writer::ColumnEncoder>,
rel: Option<writer::ColumnEncoder>,
}
/// Build the raw single-pass typed sink + its writers, shared by the bz2 and gz
/// raw entries. Nodes, ways AND relations use the parallel (worker-encode) writer
/// pattern (one row-group stitcher per table on the collector, all zstd on the
/// workers) — no single funnel writer thread.
fn raw_typed_sink(
out_dir: &Path,
compression: &str,
tables: crate::Tables,
pb: indicatif::ProgressBar,
) -> Result<(RawParallelParquetSink, RawEncoders)> {
let (node_encoder, node_writer) = if tables.nodes {
let (enc, w) = writer::parallel_node_writer(&out_dir.join("nodes.parquet"), compression)?;
(Some(enc), Some(w))
} else {
(None, None)
};
let node_acc = if tables.nodes {
Some(writer::NodeAccumulator::new())
} else {
None
};
let (way_encoder, way_writer) = if tables.ways {
let (enc, w) = writer::parallel_way_writer_with(
&out_dir.join("ways.parquet"),
compression,
tables.node_refs,
)?;
(Some(enc), Some(w))
} else {
(None, None)
};
let (rel_encoder, rel_writer) = if tables.rels {
let (enc, w) =
writer::parallel_rel_writer(&out_dir.join("relations.parquet"), compression)?;
(Some(enc), Some(w))
} else {
(None, None)
};
let way_acc = if tables.ways {
Some(writer::WayAccumulator::with_node_refs(tables.node_refs))
} else {
None
};
let rel_acc = if tables.rels {
Some(writer::RelAccumulator::new())
} else {
None
};
let encoders = RawEncoders {
node: node_encoder.clone(),
way: way_encoder.clone(),
rel: rel_encoder.clone(),
};
let sink = RawParallelParquetSink {
node_acc,
node_encoder,
node_writer,
node_count: 0,
way_acc,
rel_acc,
way_encoder,
rel_encoder,
way_writer,
rel_writer,
way_count: 0,
rel_count: 0,
tables,
prev_stub: Vec::new(),
pb,
};
Ok((sink, encoders))
}
/// Raw single-pass `.osm.bz2` → GeoParquet on the Gatling engine.
///
/// The engine decodes bz2 in parallel (split + N decode workers, no barrier) and
/// hands contiguous XML to [`RawParquetSink`]. Single-pass, no node store.
/// Returns `(nodes, ways, rels)`.
#[allow(clippy::too_many_arguments)]
pub fn read_raw_bz2(
input: &Path,
out_dir: &Path,
compression: &str,
tables: crate::Tables,
n_workers: usize,
skip_to_nodes: bool,
pb: indicatif::ProgressBar,
cancel: &crate::CancelToken,
) -> Result<(usize, usize, usize)> {
let (mut sink, encoders) = raw_typed_sink(out_dir, compression, tables, pb)?;
let skip = if skip_to_nodes {
crate::xml_to_pbf::Bz2Skip::Changesets(crate::xml_to_pbf::DATA_NEEDLES)
} else {
crate::xml_to_pbf::Bz2Skip::None
};
let crate::xml_to_pbf::Bz2Prepared {
reader,
max_blocksize,
initial_carry,
..
} = crate::xml_to_pbf::bz2_prepare(input, skip)?;
let codec = RawBz2ParquetTypedCodec {
max_blocksize,
tables,
node_encoder: encoders.node,
way_encoder: encoders.way,
rel_encoder: encoders.rel,
};
gatling::run_typed(
reader,
codec,
&mut Checkpointed::new(&mut sink, cancel, "pass1"),
n_workers,
crate::xml_to_pbf::mk_bz2_cfg(initial_carry),
)?;
sink.finish()
}
// ── Resolved bz2 → Parquet on the Gatling engine (2-pass) ─────────────────────
// Live path: Bz2Pass1TypedCodec + Pass1TypedSink (and pass-2 equivalents) below.
// The byte-mode Pass1ParquetSink / Pass2ParquetSink + parse_nodes_slot +
// parse_ways_rels_slot were removed 2026-05-29 — unreachable after the
// TypedCodec migration; held ~200 lines of rayon par_iter in the dead path.
// ── TypedCodec path: bz2 decode + XML parse in one worker ────────────────────
//
// Each Gatling worker decodes its bz2 segment then immediately scans the decoded
// XML, producing typed records. No rayon pool inside the collector.
//
// Boundary handling: bz2 block boundaries do not align with XML element boundaries.
// Each decoded segment is trimmed with find_top_level_start / find_safe_slot_end to
// isolate complete elements ("middle"). The leading/trailing partial bytes
// (left_stub / right_stub) are carried in the output and reassembled by the sink
// between consecutive segments so no element is ever dropped.
/// Output of one bz2 segment's pass-1 parse (node coords + optional node row groups).
///
/// `node_row_groups` are encoded **in the worker**: each worker turns its
/// segment's nodes into one (or more) fully-compressed Parquet row groups, so
/// the zstd/page encoding runs in parallel across all Gatling workers. The
/// collector only stitches the pre-compressed chunks into the file — it never
/// compresses. Boundary records (in `left_stub` / `right_stub`) are still
/// handled by the collector via `flush_stub`, but their volume is tiny
/// (~1–10 records per segment vs ~50K in the middle).
struct P1Segment {
node_row_groups: Vec<writer::EncodedRowGroup>, // empty when tables.nodes=false
packed: Vec<u8>, // tightly packed (i64 id, f32 lat, f32 lon) per node
left_stub: Vec<u8>, // decoded bytes before first complete element
right_stub: Vec<u8>, // decoded bytes after last complete element
}
impl Default for P1Segment {
fn default() -> Self {
Self {
node_row_groups: Vec::new(),
packed: Vec::new(),
left_stub: Vec::new(),
right_stub: Vec::new(),
}
}
}
/// Output of one bz2 segment's pass-2 parse (resolved ways + relations).
struct P2Segment {
/// Full ROW_GROUP_SIZE way row groups, Arrow-built + zstd-encoded ON THE
/// WORKER (parallel). The collector only stitches these in order.
way_row_groups: Vec<writer::EncodedRowGroup>,
/// Full relation row groups, likewise pre-encoded on the worker.
rel_row_groups: Vec<writer::EncodedRowGroup>,
/// Count of way rows represented by `way_row_groups` (for progress/totals).
way_rows: usize,
/// Count of relation rows represented by `rel_row_groups`.
rel_rows: usize,
left_stub: Vec<u8>,
right_stub: Vec<u8>,
/// **The relation clip prefilter's carry** (all four empty unless a clip is
/// active and the PBF path split pass 2 in two — see `read_resolved_pbf`).
/// The way sub-pass reports the ids of the ways it KEPT, which become the
/// [`crate::reader::RelationClipOracle`]; the relation sub-pass reports the ids of
/// the relations it kept, the relations it could not settle on the worker, and
/// how many it dropped.
kept_way_ids: Vec<i64>,
kept_rel_ids: Vec<i64>,
pending_rels: Vec<crate::reader::PendingRelation>,
rels_dropped: u64,
}
impl Default for P2Segment {
fn default() -> Self {
Self {
way_row_groups: Vec::new(),
rel_row_groups: Vec::new(),
way_rows: 0,
rel_rows: 0,
left_stub: Vec::new(),
right_stub: Vec::new(),
kept_way_ids: Vec::new(),
kept_rel_ids: Vec::new(),
pending_rels: Vec::new(),
rels_dropped: 0,
}
}
}
thread_local! {
/// Per-worker pass-2 way accumulator that persists across the segments a
/// single worker thread processes — the exact mirror of pass-1's `NODE_ACC`.
/// Full ROW_GROUP_SIZE batches are Arrow-built + zstd-encoded inside the
/// worker (parallel); the sub-full remainder carries into this worker's next
/// segment so we don't emit a tiny group per segment. `pass2_finish_worker`
/// drains the final remainder once before the thread exits.
static WAY_ACC: std::cell::RefCell<Option<writer::WayAccumulator>> =
const { std::cell::RefCell::new(None) };
/// Per-worker pass-2 relation accumulator (mirror of `WAY_ACC`).
static REL_ACC: std::cell::RefCell<Option<writer::RelAccumulator>> =
const { std::cell::RefCell::new(None) };
}
/// Push finished way/rel records into the per-worker accumulators, sealing +
/// encoding any full ROW_GROUP_SIZE row groups on this worker thread. Returns
/// the encoded full groups (the sub-full remainder stays in the thread-locals).
/// `way_enc`/`rel_enc` are the shared (cloneable, Sync) column encoders.
fn pass2_accumulate(
ways: &[WayRecord],
rels: &[RelationRecord],
way_enc: Option<&writer::ColumnEncoder>,
rel_enc: Option<&writer::ColumnEncoder>,
out: &mut P2Segment,
sieve: Option<&mut P2Sieve>,
) {
// The prefilter's bookkeeping, only when the caller asked for it: the ids of
// what this segment KEPT, so the way sub-pass can hand the relation sub-pass its
// oracle and the relation sub-pass can hand the collector its fixpoint input.
if let Some(sv) = sieve {
if sv.collect_way_ids && way_enc.is_some() {
out.kept_way_ids.extend(ways.iter().map(|w| w.id));
}
if sv.oracle_active && rel_enc.is_some() {
out.kept_rel_ids.extend(rels.iter().map(|r| r.id));
}
}
// Count rows ONLY from the batches actually emitted (here as full groups,
// and in `pass2_finish_worker` for the tail) so every row is counted exactly
// once regardless of how it splits across segments / row groups.
if let Some(enc) = way_enc {
WAY_ACC.with(|cell| {
let mut slot = cell.borrow_mut();
// The WRITER's schema decides whether this batch carries `node_refs`
// (`katana-6`) — never a second copy of the flag. `WayAccumulator::new()`
// here is what made the first cut of that fix explode on 32 threads.
let acc = slot.get_or_insert_with(|| {
writer::WayAccumulator::with_node_refs(enc.has_column("node_refs"))
});
for r in ways {
acc.push(r);
if let Ok(Some(batch)) = acc.take_if_full() {
out.way_rows += batch.num_rows();
out.way_row_groups
.push(enc.encode(&batch).expect("way row group encode"));
}
}
});
}
if let Some(enc) = rel_enc {
REL_ACC.with(|cell| {
let mut slot = cell.borrow_mut();
let acc = slot.get_or_insert_with(writer::RelAccumulator::new);
for r in rels {
acc.push(r);
if let Ok(Some(batch)) = acc.take_if_full() {
out.rel_rows += batch.num_rows();
out.rel_row_groups
.push(enc.encode(&batch).expect("rel row group encode"));
}
}
});
}
}
/// What the relation clip prefilter asks a pass-2 worker to carry, per call.
struct P2Sieve {
/// The way sub-pass: report the ids of the ways kept.
collect_way_ids: bool,
/// The relation sub-pass: report the ids of the relations kept.
oracle_active: bool,
}
/// Which pass-2 elements a reader indexes at all: only the kinds the tables it is
/// writing ask for. A ways-only sub-pass does not parse a relation it will throw
/// away, and a relations-only sub-pass opened on the last way block does not resolve
/// the ways it starts among.
fn pass2_wants(tables: crate::Tables, kind: &ElemKind) -> bool {
match kind {
ElemKind::Way => tables.ways,
ElemKind::Relation => tables.rels,
ElemKind::Node => false,
}
}
/// Drain the per-worker `WAY_ACC`/`REL_ACC` remainders into one final row group
/// each. Mirror of pass-1's `pass1_finish_worker`.
fn pass2_finish_worker(
way_enc: Option<&writer::ColumnEncoder>,
rel_enc: Option<&writer::ColumnEncoder>,
) -> Option<P2Segment> {
let mut seg = P2Segment::default();
let mut any = false;
if let Some(enc) = way_enc {
WAY_ACC.with(|cell| {
if let Some(mut acc) = cell.borrow_mut().take() {
if let Ok(Some(batch)) = acc.take_remaining() {
seg.way_rows += batch.num_rows();
seg.way_row_groups
.push(enc.encode(&batch).expect("way tail row group encode"));
any = true;
}
}
});
}
if let Some(enc) = rel_enc {
REL_ACC.with(|cell| {
if let Some(mut acc) = cell.borrow_mut().take() {
if let Ok(Some(batch)) = acc.take_remaining() {
seg.rel_rows += batch.num_rows();
seg.rel_row_groups
.push(enc.encode(&batch).expect("rel tail row group encode"));
any = true;
}
}
});
}
if any { Some(seg) } else { None }
}
// ── Pass-1 TypedCodec ─────────────────────────────────────────────────────────
struct Bz2Pass1TypedCodec {
max_blocksize: u32,
tables: crate::Tables,
/// Per-worker Parquet encoder: each worker compresses its own segment's
/// node row group in parallel (no rayon, no shared writer lock).
node_encoder: Option<writer::NodeColumnEncoder>,
/// Optional region clip applied per-node inside the worker (no extra pass).
clip: Option<Arc<crate::Clip>>,
}
thread_local! {
/// Per-worker node accumulator that persists across the segments a single
/// worker thread processes. Full ROW_GROUP_SIZE batches are sealed mid-stream
/// (parallel zstd, in the worker); the sub-full remainder carries into this
/// worker's next segment instead of being flushed as its own tiny row group.
/// `finish_worker` drains the final remainder once, before the thread exits.
static NODE_ACC: std::cell::RefCell<Option<writer::NodeAccumulator>> =
const { std::cell::RefCell::new(None) };
}
/// **The one pass-1 node emit.** Applies the region `clip`, then — only for a
/// node that survives it — packs the 16-byte coord record (`id`, `lat`, `lon`)
/// the node store is built from and pushes the `NodeRecord` for nodes.parquet
/// into `acc`. Returns `true` iff the node was kept (so the caller can count it
/// and drain a full row group).
///
/// BOTH pass-1 writers route through this: the WORKER's middle
/// ([`pass1_parse_decoded`]) and the COLLECTOR's boundary stub
/// ([`Pass1TypedSink::flush_stub`]). They were hand-copied twins, and the stub
/// copy had **no clip test at all** — so with a clip active every slot boundary
/// leaked whatever element straddled it, unclipped, into both the coord store and
/// nodes.parquet. MEASURED on korp's Sweden cut of `planet-260427.osm.bz2`
/// (`~/.korp/work/sweden`, 2026-08-02): **16 260 of the 16 583 nodes in the last
/// row group of nodes.parquet — 98 % — lie outside the sweden.poly bbox**, ids
/// spanning the whole planet (337 966 … 13 765 167 641), and the same 16 260 sit
/// in `node_coords.arrow`; pass 2 then resolved them into 1 167 wholly-foreign
/// ways (Wellington, Toronto, Vienna, Kyiv). One writer means the clip cannot be
/// applied on one path and skipped on the other.
#[inline]
fn pass1_emit_node(
buf: &[u8],
e: &ElemIndex,
clip: Option<&crate::Clip>,
packed: &mut Vec<u8>,
acc: Option<&mut writer::NodeAccumulator>,
) -> bool {
// Region clip FIRST: an out-of-region node reaches neither the coord store
// nor the parquet. Tested on the full-precision e7 ints so the predicate
// matches the PBF path exactly.
if !crate::clip_keeps(clip, e.lon_e7 as f64 / 1e7, e.lat_e7 as f64 / 1e7) {
return false;
}
let lat = e.lat_e7 as f32 / 1e7_f32;
let lon = e.lon_e7 as f32 / 1e7_f32;
packed.extend_from_slice(&e.id.to_le_bytes());
packed.extend_from_slice(&lat.to_le_bytes());
packed.extend_from_slice(&lon.to_le_bytes());
if let Some(acc) = acc {
let lo = e.file_offset as usize;
let hi = lo.saturating_add(e.file_length as usize).min(buf.len());
let b = buf.get(lo..hi).unwrap_or_default();
acc.push(&NodeRecord {
id: e.id,
lon_lat: (f64::from(lon), f64::from(lat)),
tags_json: parse_tags_json(b),
version: parse_attr_i32(opening_tag(b), b"version"),
});
}
true
}
/// Parse decoded XML bytes into a pass-1 segment (node coords + node row
/// groups). Shared by the bz2 and gz typed codecs — only the decode step
/// upstream differs. Uses the per-worker [`NODE_ACC`] accumulator so full row
/// groups are sealed (parallel zstd) inside the worker.
fn pass1_parse_decoded(
decoded: &[u8],
tables: crate::Tables,
node_encoder: Option<&writer::NodeColumnEncoder>,
clip: Option<&crate::Clip>,
) -> P1Segment {
if decoded.is_empty() {
return P1Segment::default();
}
let first = xml_vtd::find_top_level_start(decoded, 0);
let last = xml_vtd::find_safe_slot_end(decoded);
let last = last.max(first);
let left_stub = decoded[..first].to_vec();
let right_stub = decoded[last..].to_vec();
let middle = &decoded[first..last];
let mut packed = Vec::new();
let mut node_row_groups: Vec<writer::EncodedRowGroup> = Vec::new();
NODE_ACC.with(|cell| {
let mut slot = cell.borrow_mut();
let mut acc = if tables.nodes && node_encoder.is_some() {
Some(slot.get_or_insert_with(writer::NodeAccumulator::new))
} else {
None
};
let enc = node_encoder;
xml_vtd::build_elem_index_slice(middle, 0, &mut |e| {
if e.kind != ElemKind::Node {
return;
}
if !pass1_emit_node(middle, &e, clip, &mut packed, acc.as_deref_mut()) {
return;
}
if let (Some(acc), Some(enc)) = (acc.as_deref_mut(), enc) {
if let Ok(Some(batch)) = acc.take_if_full() {
node_row_groups.push(enc.encode(&batch).expect("node row group encode"));
}
}
});
});
P1Segment {
node_row_groups,
packed,
left_stub,
right_stub,
}
}
/// Drain the per-worker [`NODE_ACC`] remainder into one final row group. Shared
/// `finish_worker` body for the bz2 and gz pass-1 codecs.
fn pass1_finish_worker(
tables: crate::Tables,
node_encoder: Option<&writer::NodeColumnEncoder>,
) -> Option<P1Segment> {
if !tables.nodes {
return None;
}
let enc = node_encoder?;
NODE_ACC.with(|cell| {
let mut acc = cell.borrow_mut().take()?;
let batch = acc.take_remaining().ok().flatten()?;
let rg = enc.encode(&batch).expect("node tail row group encode");
Some(P1Segment {
node_row_groups: vec![rg],
..P1Segment::default()
})
})
}
/// Parse one PBF blob into a pass-1 segment (node coords + node row groups).
/// PBF analogue of [`pass1_parse_decoded`]: blobs are self-contained, so there
/// are no boundary stubs. `data` is the full slot slice; `seg` indexes the blob.
/// Parse one COARSE candidate range into a pass-1 segment. `seg` is
/// `(cand_start, cand_end)`; the worker first aligns to its own blob boundary via
/// `pbf_io::align_and_collect` (boundary-finding is DISTRIBUTED across all workers,
/// not serialized on main) and then decodes every OSMData blob it owns into ONE
/// segment. Blobs are self-contained so there are no boundary stubs. `data` is
/// the full slot slice (borrowed, zero-copy).
fn pass1_parse_blob(
data: &[u8],
seg: &crate::xml_to_pbf::PbfSeg,
tables: crate::Tables,
node_encoder: Option<&writer::NodeColumnEncoder>,
clip: Option<&crate::Clip>,
plan: &crate::pbf_io::BlobPlanBuilder,
) -> P1Segment {
let blobs = crate::pbf_io::align_and_collect(data, seg.0, seg.1);
if blobs.is_empty() {
return P1Segment::default();
}
// ABSOLUTE file offset of this slot's byte 0 — the pass-2 plan's key.
let base = seg.2;
// Reused across the blobs this worker owns — no per-blob Vec churn.
let mut packed: Vec<u8> = Vec::new();
let mut node_row_groups: Vec<writer::EncodedRowGroup> = Vec::new();
let mut seen: Vec<u64> = Vec::with_capacity(blobs.len());
let mut with_ways: Vec<u64> = Vec::new();
let mut with_rels: Vec<u64> = Vec::new();
let mut first_rel_frame: Option<u64> = None;
NODE_ACC.with(|cell| {
let mut slot = cell.borrow_mut();
let mut acc = if tables.nodes && node_encoder.is_some() {
Some(slot.get_or_insert_with(writer::NodeAccumulator::new))
} else {
None
};
for pos in &blobs {
let (records, coords, kinds) =
crate::reader::blob_nodes_and_coords_probed(data, pos, clip);
let abs = base.saturating_add(pos.offset as u64);
seen.push(abs);
if kinds.ways {
with_ways.push(abs);
}
if kinds.rels {
with_rels.push(abs);
let frame = base.saturating_add(pos.frame as u64);
first_rel_frame = Some(first_rel_frame.map_or(frame, |f| f.min(frame)));
}
if coords.is_empty() {
continue;
}
packed.reserve(coords.len() * 16);
for (id, lat, lon) in &coords {
packed.extend_from_slice(&id.to_le_bytes());
packed.extend_from_slice(&lat.to_le_bytes());
packed.extend_from_slice(&lon.to_le_bytes());
}
if let (Some(acc), Some(enc)) = (acc.as_deref_mut(), node_encoder) {
for r in &records {
acc.push(r);
if let Ok(Some(batch)) = acc.take_if_full() {
node_row_groups.push(enc.encode(&batch).expect("node row group encode"));
}
}
}
}
});
plan.note(&seen, &with_ways, &with_rels, first_rel_frame);
P1Segment {
node_row_groups,
packed,
left_stub: Vec::new(),
right_stub: Vec::new(),
}
}
/// Parse one COARSE candidate range into a pass-2 segment (resolved ways +
/// relations). `seg` is `(cand_start, cand_end)`; see [`pass1_parse_blob`] for the
/// alignment contract. Arrow-build + zstd-encode of full row groups happens HERE
/// on the worker (via the per-worker `WAY_ACC`/`REL_ACC` accumulators); only the
/// sub-full remainder carries into this worker's next segment. No boundary stubs.
#[allow(clippy::too_many_arguments)]
fn pass2_parse_blob(
data: &[u8],
seg: &crate::xml_to_pbf::PbfSeg,
store: &Arc<NodeStore>,
tables: crate::Tables,
way_enc: Option<&writer::ColumnEncoder>,
rel_enc: Option<&writer::ColumnEncoder>,
clip_active: bool,
plan: &crate::pbf_io::BlobPlan,
opened: &AtomicU64,
skipped: &AtomicU64,
unresolved_refs: &AtomicU64,
shortened_ways: &AtomicU64,
sieve: Option<&P2SieveSetup>,
) -> P2Segment {
let blobs = crate::pbf_io::align_and_collect(data, seg.0, seg.1);
let mut out = P2Segment::default();
if blobs.is_empty() {
return out;
}
let base = seg.2;
let (mut n_open, mut n_skip) = (0u64, 0u64);
let mut loss = crate::reader::ResolveLoss::default();
let mut flags = sieve.map(P2SieveSetup::flags);
let oracle = sieve.and_then(|sv| sv.oracle.as_deref());
for pos in &blobs {
// The plan test happens BEFORE the inflate — that is the entire saving.
// `blob_ways_rels` would decompress ~700 KB of node deltas and then find
// no way group; the mmap path has skipped that since the plan landed.
// Asked per KIND: a sub-pass that writes only ways leaves the relation
// blobs closed, and one that writes only relations leaves the way blobs closed.
if !plan.needs_kind(base.saturating_add(pos.offset as u64), tables.ways, tables.rels) {
n_skip += 1;
continue;
}
n_open += 1;
let (ways, rels, blob_loss, rel_sieve) =
crate::reader::blob_ways_rels(data, pos, store, clip_active, oracle);
loss.add(blob_loss);
let ways = if tables.ways { ways } else { Vec::new() };
let rels = if tables.rels { rels } else { Vec::new() };
pass2_accumulate(&ways, &rels, way_enc, rel_enc, &mut out, flags.as_mut());
if tables.rels && oracle.is_some() {
out.pending_rels.extend(rel_sieve.pending);
out.rels_dropped += rel_sieve.dropped;
}
}
opened.fetch_add(n_open, Ordering::Relaxed);
skipped.fetch_add(n_skip, Ordering::Relaxed);
// Way-resolution loss (audit finding (a)) — carried to the pass-2 done line
// of the convert summary, so a silently shortened way is silent no more.
unresolved_refs.fetch_add(loss.unresolved_refs, Ordering::Relaxed);
shortened_ways.fetch_add(loss.shortened_ways, Ordering::Relaxed);
out
}
/// Parse decoded XML bytes into a pass-2 segment (resolved ways + relations).
/// Shared by the osm, bz2 and gz typed codecs. The segment's complete (middle)
/// ways/rels are Arrow-built + zstd-encoded ON THE WORKER (via the per-worker
/// accumulators); only the boundary stubs travel raw to the collector, which
/// stitches the straddling element(s) into its own (tiny) stub accumulator.
///
/// `sieve` is the relation clip prefilter: with an oracle, every relation is judged
/// as it is parsed (`parse_chunk` → `RelationClipOracle::admit`) and the unsettled
/// ones ride `P2Segment::pending_rels` to the collector's fixpoint — exactly what
/// `pass2_parse_blob` does for a PBF blob.
#[allow(clippy::too_many_arguments)]
fn pass2_parse_decoded(
decoded: &[u8],
store: &Arc<NodeStore>,
backend: &dyn GpuBackend,
tables: crate::Tables,
way_enc: Option<&writer::ColumnEncoder>,
rel_enc: Option<&writer::ColumnEncoder>,
clip_active: bool,
unresolved_refs: &AtomicU64,
shortened_ways: &AtomicU64,
sieve: Option<&P2SieveSetup>,
) -> P2Segment {
if decoded.is_empty() {
return P2Segment::default();
}
let first = xml_vtd::find_top_level_start(decoded, 0);
let last = xml_vtd::find_safe_slot_end(decoded);
let last = last.max(first);
let left_stub = decoded[..first].to_vec();
let right_stub = decoded[last..].to_vec();
let middle = &decoded[first..last];
let mut elems: Vec<ElemIndex> = Vec::new();
xml_vtd::build_elem_index_slice(middle, 0, &mut |e| {
if pass2_wants(tables, &e.kind) {
elems.push(e);
}
});
if elems.is_empty() {
return P2Segment {
left_stub,
right_stub,
..P2Segment::default()
};
}
let refs: Vec<&ElemIndex> = elems.iter().collect();
let chunks = plan_chunks(&refs, MAX_CHUNK_BYTES);
let oracle = sieve.and_then(|sv| sv.oracle.as_deref());
let mut all_ways: Vec<CollectedWay> = Vec::new();
let mut all_rels: Vec<RelationRecord> = Vec::new();
let mut loss = crate::reader::ResolveLoss::default();
let mut rel_sieve = crate::reader::RelSieve::default();
for (lo, hi) in chunks {
let (ways, rels, chunk_loss, chunk_sieve) =
parse_chunk(middle, &refs[lo..hi], store, 0, clip_active, oracle);
all_ways.extend(ways);
all_rels.extend(rels);
loss.add(chunk_loss);
rel_sieve.absorb(chunk_sieve);
}
// Way-resolution loss (audit finding (a)) — carried to the pass-2 done line
// of the bz2/gz convert summary, exactly like the PBF path's counters. This
// was the disclosed uncounted twin: production's planet-latest.osm.bz2 cut
// resolved HERE while only the PBF seam counted.
unresolved_refs.fetch_add(loss.unresolved_refs, Ordering::Relaxed);
shortened_ways.fetch_add(loss.shortened_ways, Ordering::Relaxed);
let ways = if tables.ways {
collected_to_way_records(all_ways, backend)
} else {
Vec::new()
};
let rels = if tables.rels { all_rels } else { Vec::new() };
let mut out = P2Segment::default();
let mut flags = sieve.map(P2SieveSetup::flags);
pass2_accumulate(&ways, &rels, way_enc, rel_enc, &mut out, flags.as_mut());
if tables.rels && oracle.is_some() {
out.pending_rels.extend(rel_sieve.pending);
out.rels_dropped += rel_sieve.dropped;
}
out.left_stub = left_stub;
out.right_stub = right_stub;
out
}
impl gatling::TypedCodec for Bz2Pass1TypedCodec {
type Seg = (u64, u64);
type Output = P1Segment;
fn split(
&self,
data: &[u8],
n_workers: usize,
is_last: bool,
) -> Option<gatling::Split<(u64, u64)>> {
// ONE bz2 slot split for every bz2 codec in the crate (LAW 5).
crate::xml_to_pbf::bz2_split(data, n_workers, self.max_blocksize, is_last)
}
fn transform(&self, data: &[u8], &(start_bit, end_bit): &(u64, u64)) -> P1Segment {
let decoded = lbzip2::chunk::decode_segment(data, start_bit, end_bit, self.max_blocksize);
let enc = if self.tables.nodes {
self.node_encoder.as_ref()
} else {
None
};
pass1_parse_decoded(&decoded, self.tables, enc, self.clip.as_deref())
}
fn finish_worker(&self) -> Option<P1Segment> {
pass1_finish_worker(self.tables, self.node_encoder.as_ref())
}
}
/// gzip pass-1 [`gatling::TypedCodec`]: DEFLATE decode + node parse fused in the
/// worker. Mirrors [`Bz2Pass1TypedCodec`] but with lgz flush-boundary splits.
struct GzPass1TypedCodec {
tables: crate::Tables,
node_encoder: Option<writer::NodeColumnEncoder>,
clip: Option<Arc<crate::Clip>>,
}
impl gatling::TypedCodec for GzPass1TypedCodec {
type Seg = (usize, usize);
type Output = P1Segment;
fn split(
&self,
data: &[u8],
n_workers: usize,
is_last: bool,
) -> Option<gatling::Split<(usize, usize)>> {
crate::xml_to_pbf::gz_split(data, n_workers, is_last)
}
fn transform(&self, data: &[u8], seg: &(usize, usize)) -> P1Segment {
let decoded = crate::xml_to_pbf::gz_decode_segment(data, seg);
let enc = if self.tables.nodes {
self.node_encoder.as_ref()
} else {
None
};
pass1_parse_decoded(&decoded, self.tables, enc, self.clip.as_deref())
}
fn finish_worker(&self) -> Option<P1Segment> {
pass1_finish_worker(self.tables, self.node_encoder.as_ref())
}
}
/// Uncompressed `.osm` pass-1 (the `read_resolved_osm` entry). Identical to the
/// bz2/gz pass-1 codecs but with NO decode step — the slot bytes are already XML,
/// so `transform` slices its element-aligned segment and runs the shared
/// [`pass1_parse_decoded`]. This replaced the old serial `Pass1ResolvedCodec` +
/// per-record `NodeAccumulator` on the collector, which pegged raw `.osm` convert
/// at ~3 cores (workers starved while the single sink built + wrote every node
/// batch). Nodes are now Arrow-built + zstd-encoded on the workers; the sink only
/// stitches pre-compressed row groups + `write_raw`s the packed coords.
struct RawXmlPass1Codec {
tables: crate::Tables,
node_encoder: Option<writer::NodeColumnEncoder>,
clip: Option<Arc<crate::Clip>>,
}
impl gatling::TypedCodec for RawXmlPass1Codec {
type Seg = (usize, usize);
type Output = P1Segment;
fn split(
&self,
data: &[u8],
n_workers: usize,
is_last: bool,
) -> Option<gatling::Split<(usize, usize)>> {
xml_split_typed(data, n_workers, is_last)
}
fn transform(&self, data: &[u8], seg: &(usize, usize)) -> P1Segment {
let enc = if self.tables.nodes {
self.node_encoder.as_ref()
} else {
None
};
pass1_parse_decoded(&data[seg.0..seg.1], self.tables, enc, self.clip.as_deref())
}
fn finish_worker(&self) -> Option<P1Segment> {
pass1_finish_worker(self.tables, self.node_encoder.as_ref())
}
}
// ── Pass-2 TypedCodec ─────────────────────────────────────────────────────────
struct Bz2Pass2TypedCodec<'a> {
max_blocksize: u32,
tables: crate::Tables,
store: Arc<NodeStore>,
backend: &'a dyn GpuBackend,
way_enc: Option<writer::ColumnEncoder>,
rel_enc: Option<writer::ColumnEncoder>,
clip_active: bool,
/// Way-resolution loss (audit finding (a)) — the bz2 twin of
/// [`PbfPass2TypedCodec`]'s counters; lands on the pass-2 `done` line as
/// `unresolved_node_refs` / `shortened_ways`.
unresolved_refs: Arc<AtomicU64>,
shortened_ways: Arc<AtomicU64>,
/// The relation clip prefilter's setup, `None` when no clip is active.
sieve: Option<P2SieveSetup>,
}
impl gatling::TypedCodec for Bz2Pass2TypedCodec<'_> {
type Seg = (u64, u64);
type Output = P2Segment;
fn split(
&self,
data: &[u8],
n_workers: usize,
is_last: bool,
) -> Option<gatling::Split<(u64, u64)>> {
// ONE bz2 slot split for every bz2 codec in the crate (LAW 5).
crate::xml_to_pbf::bz2_split(data, n_workers, self.max_blocksize, is_last)
}
fn transform(&self, data: &[u8], &(start_bit, end_bit): &(u64, u64)) -> P2Segment {
let decoded = lbzip2::chunk::decode_segment(data, start_bit, end_bit, self.max_blocksize);
pass2_parse_decoded(
&decoded,
&self.store,
self.backend,
self.tables,
self.way_enc.as_ref(),
self.rel_enc.as_ref(),
self.clip_active,
&self.unresolved_refs,
&self.shortened_ways,
self.sieve.as_ref(),
)
}
fn finish_worker(&self) -> Option<P2Segment> {
pass2_finish_worker(self.way_enc.as_ref(), self.rel_enc.as_ref())
}
}
/// gzip pass-2 [`gatling::TypedCodec`]: DEFLATE decode + way/relation resolve
/// fused in the worker. Mirrors [`Bz2Pass2TypedCodec`] with lgz splits.
struct GzPass2TypedCodec<'a> {
tables: crate::Tables,
store: Arc<NodeStore>,
backend: &'a dyn GpuBackend,
way_enc: Option<writer::ColumnEncoder>,
rel_enc: Option<writer::ColumnEncoder>,
clip_active: bool,
/// Way-resolution loss (audit finding (a)) — the gz twin of
/// [`PbfPass2TypedCodec`]'s counters; lands on the pass-2 `done` line as
/// `unresolved_node_refs` / `shortened_ways`.
unresolved_refs: Arc<AtomicU64>,
shortened_ways: Arc<AtomicU64>,
/// The relation clip prefilter's setup, `None` when no clip is active.
sieve: Option<P2SieveSetup>,
}
impl gatling::TypedCodec for GzPass2TypedCodec<'_> {
type Seg = (usize, usize);
type Output = P2Segment;
fn split(
&self,
data: &[u8],
n_workers: usize,
is_last: bool,
) -> Option<gatling::Split<(usize, usize)>> {
crate::xml_to_pbf::gz_split(data, n_workers, is_last)
}
fn transform(&self, data: &[u8], seg: &(usize, usize)) -> P2Segment {
let decoded = crate::xml_to_pbf::gz_decode_segment(data, seg);
pass2_parse_decoded(
&decoded,
&self.store,
self.backend,
self.tables,
self.way_enc.as_ref(),
self.rel_enc.as_ref(),
self.clip_active,
&self.unresolved_refs,
&self.shortened_ways,
self.sieve.as_ref(),
)
}
fn finish_worker(&self) -> Option<P2Segment> {
pass2_finish_worker(self.way_enc.as_ref(), self.rel_enc.as_ref())
}
}
/// Plain-XML pass-2 [`gatling::TypedCodec`]: element-boundary split + way/relation
/// resolve on the worker. The `.osm` reader used to have its own serial sink
/// (`WayWriter::push` per row on the collector); it now runs the SAME
/// `pass2_parse_decoded` → `Pass2TypedSink` path as bz2/gz/pbf, which is what puts
/// the relation clip prefilter on it by construction rather than by a second copy.
struct RawXmlPass2Codec<'a> {
tables: crate::Tables,
store: Arc<NodeStore>,
backend: &'a dyn GpuBackend,
way_enc: Option<writer::ColumnEncoder>,
rel_enc: Option<writer::ColumnEncoder>,
clip_active: bool,
unresolved_refs: Arc<AtomicU64>,
shortened_ways: Arc<AtomicU64>,
sieve: Option<P2SieveSetup>,
}
impl gatling::TypedCodec for RawXmlPass2Codec<'_> {
type Seg = (usize, usize);
type Output = P2Segment;
fn split(
&self,
data: &[u8],
n_workers: usize,
is_last: bool,
) -> Option<gatling::Split<(usize, usize)>> {
xml_split_typed(data, n_workers, is_last)
}
fn transform(&self, data: &[u8], seg: &(usize, usize)) -> P2Segment {
pass2_parse_decoded(
&data[seg.0..seg.1],
&self.store,
self.backend,
self.tables,
self.way_enc.as_ref(),
self.rel_enc.as_ref(),
self.clip_active,
&self.unresolved_refs,
&self.shortened_ways,
self.sieve.as_ref(),
)
}
fn finish_worker(&self) -> Option<P2Segment> {
pass2_finish_worker(self.way_enc.as_ref(), self.rel_enc.as_ref())
}
}
// ── PBF TypedCodecs ───────────────────────────────────────────────────────────
//
// PBF blobs are self-contained: `split` cuts the slot at blob frame boundaries
// (pbf_io), `transform` zlib+protobuf-decodes one blob into a P1/P2 segment.
// No boundary stubs (blobs never straddle a logical element), so the shared
// Pass1TypedSink / Pass2TypedSink consume them with empty stubs — unifying PBF
// onto the same coord-sort / node-store path as bz2 and gz.
struct PbfPass1TypedCodec {
tables: crate::Tables,
node_encoder: Option<writer::NodeColumnEncoder>,
clip: Option<Arc<crate::Clip>>,
/// Absolute file offset of the slot being split — the key the pass-2 plan
/// is written in. See [`crate::xml_to_pbf::SlotCursor`].
cursor: crate::xml_to_pbf::SlotCursor,
/// Filled in as pass 1 decodes: which blobs carry ways or relations.
plan: Arc<crate::pbf_io::BlobPlanBuilder>,
}
impl gatling::TypedCodec for PbfPass1TypedCodec {
type Seg = crate::xml_to_pbf::PbfSeg;
type Output = P1Segment;
fn split(
&self,
data: &[u8],
n_workers: usize,
is_last: bool,
) -> Option<gatling::Split<Self::Seg>> {
self.cursor.split(data, n_workers, is_last)
}
fn transform(&self, data: &[u8], seg: &Self::Seg) -> P1Segment {
let enc = if self.tables.nodes {
self.node_encoder.as_ref()
} else {
None
};
pass1_parse_blob(
data,
seg,
self.tables,
enc,
self.clip.as_deref(),
&self.plan,
)
}
fn finish_worker(&self) -> Option<P1Segment> {
pass1_finish_worker(self.tables, self.node_encoder.as_ref())
}
}
struct PbfPass2TypedCodec {
tables: crate::Tables,
store: Arc<NodeStore>,
/// Shared column encoders handed to every worker so the Arrow build + zstd
/// encode of the way/rel row groups runs in PARALLEL on the workers, not on
/// the single collector (the pass-2 saturation fix). `None` when the
/// corresponding output is disabled.
way_enc: Option<writer::ColumnEncoder>,
rel_enc: Option<writer::ColumnEncoder>,
/// True when a region clip is active: pass-2 then drops ways with zero
/// resolvable (in-region) coords (node-membership clip semantics).
clip_active: bool,
/// Same cursor discipline as pass 1 — the engine re-reads the file, so the
/// base has to be re-derived, and blob offsets are file-intrinsic either way.
cursor: crate::xml_to_pbf::SlotCursor,
/// What pass 1 learned. `BlobPlan::open_everything()` reproduces the old
/// visit-every-blob behaviour exactly.
plan: Arc<crate::pbf_io::BlobPlan>,
/// Applied-output counters, so the saving is a measured number in the phase
/// log and not a claim: blobs opened vs blobs skipped by the plan.
opened: Arc<AtomicU64>,
skipped: Arc<AtomicU64>,
/// Way-resolution loss (audit finding (a)): refs the store could not
/// resolve, and the ways emitted shortened because of them. Reported on
/// the pass-2 `done` line as `unresolved_node_refs` / `shortened_ways`.
unresolved_refs: Arc<AtomicU64>,
shortened_ways: Arc<AtomicU64>,
/// The relation clip prefilter's setup, `None` when no clip is active (every
/// relation is kept, byte-identical to before the prefilter existed).
sieve: Option<P2SieveSetup>,
}
/// The relation clip prefilter as a pass-2 codec — and the collector's stub path —
/// sees it. `None` when no clip is active.
#[derive(Clone)]
struct P2SieveSetup {
/// The way sub-pass collects the ids it keeps.
collect_way_ids: bool,
/// The relation sub-pass judges every relation against this.
oracle: Option<Arc<crate::reader::RelationClipOracle>>,
}
impl P2SieveSetup {
/// The per-call flags a worker carries into `pass2_accumulate`.
fn flags(&self) -> P2Sieve {
P2Sieve {
collect_way_ids: self.collect_way_ids,
oracle_active: self.oracle.is_some(),
}
}
}
impl gatling::TypedCodec for PbfPass2TypedCodec {
type Seg = crate::xml_to_pbf::PbfSeg;
type Output = P2Segment;
fn split(
&self,
data: &[u8],
n_workers: usize,
is_last: bool,
) -> Option<gatling::Split<Self::Seg>> {
self.cursor.split(data, n_workers, is_last)
}
fn transform(&self, data: &[u8], seg: &Self::Seg) -> P2Segment {
pass2_parse_blob(
data,
seg,
&self.store,
self.tables,
self.way_enc.as_ref(),
self.rel_enc.as_ref(),
self.clip_active,
&self.plan,
&self.opened,
&self.skipped,
&self.unresolved_refs,
&self.shortened_ways,
self.sieve.as_ref(),
)
}
fn finish_worker(&self) -> Option<P2Segment> {
pass2_finish_worker(self.way_enc.as_ref(), self.rel_enc.as_ref())
}
}
struct Pass1TypedSink {
coord_writer: CoordFileWriter,
node_acc: Option<writer::NodeAccumulator>,
/// Encodes the (tiny) boundary-stub node batches on the collector thread.
node_encoder: Option<writer::NodeColumnEncoder>,
/// Stitches pre-compressed row groups (from workers + boundary stubs) into
/// the file. No compression on this thread.
node_writer: Option<writer::ParallelNodeWriter>,
node_count: usize,
pb: indicatif::ProgressBar,
prev_stub: Vec<u8>, // right_stub from last segment; combined with next left_stub
tables: crate::Tables,
/// The SAME region clip the workers hold. The boundary-stub nodes this
/// collector parses are production output — they go into the coord store and
/// nodes.parquet — so they must pass the identical predicate. See
/// [`pass1_emit_node`] for what happened when they did not.
clip: Option<Arc<crate::Clip>>,
/// Optional shared `busy_ns` counter for the surrounding `Phase`.
busy_ns: Option<Arc<std::sync::atomic::AtomicU64>>,
}
impl Pass1TypedSink {
fn flush_stub(&mut self, stub: &[u8]) -> Result<()> {
if stub.is_empty() {
return Ok(());
}
let mut packed = Vec::new();
let mut kept = 0usize;
{
// Disjoint field borrows: `clip` + `tables.nodes` read, `node_acc`
// written. `coord_writer` is touched only after this block.
let clip = self.clip.as_deref();
let mut acc = if self.tables.nodes {
self.node_acc.as_mut()
} else {
None
};
xml_vtd::build_elem_index_slice(stub, 0, &mut |e| {
if e.kind != ElemKind::Node {
return;
}
if pass1_emit_node(stub, &e, clip, &mut packed, acc.as_deref_mut()) {
kept += 1;
}
});
}
self.node_count += kept;
self.coord_writer.write_raw(&packed)?;
if let (Some(acc), Some(enc), Some(w)) = (
self.node_acc.as_mut(),
self.node_encoder.as_ref(),
self.node_writer.as_mut(),
) {
while let Some(batch) = acc.take_if_full()? {
w.append(enc.encode(&batch)?)?;
}
}
Ok(())
}
fn finish(mut self) -> Result<(CoordFileWriter, usize)> {
let stub = std::mem::take(&mut self.prev_stub);
self.flush_stub(&stub)?;
if let (Some(acc), Some(enc), Some(w)) = (
self.node_acc.as_mut(),
self.node_encoder.as_ref(),
self.node_writer.as_mut(),
) {
if let Some(batch) = acc.take_remaining()? {
w.append(enc.encode(&batch)?)?;
}
}
if let Some(w) = self.node_writer.take() {
w.finish()?;
}
Ok((self.coord_writer, self.node_count))
}
}
impl gatling::TypedSink<P1Segment> for Pass1TypedSink {
fn process(&mut self, seg: P1Segment, _is_last: bool) -> Result<()> {
let t0 = std::time::Instant::now();
// Combine right_stub from previous segment with this segment's left_stub,
// then parse any complete element that straddles the segment boundary.
let boundary = [self.prev_stub.as_slice(), seg.left_stub.as_slice()].concat();
self.flush_stub(&boundary)?;
// Main records: write the packed coords, then stitch the pre-compressed
// node row groups the worker already encoded. No Arrow append, no zstd
// on this thread.
self.coord_writer.write_raw(&seg.packed)?;
self.node_count += seg.packed.len() / 16;
if let Some(w) = self.node_writer.as_mut() {
for rg in seg.node_row_groups {
w.append(rg)?;
}
}
self.pb
.inc((seg.left_stub.len() + seg.packed.len() / 16 * 220 + seg.right_stub.len()) as u64);
self.prev_stub = seg.right_stub;
if let Some(b) = &self.busy_ns {
b.fetch_add(
t0.elapsed().as_nanos() as u64,
std::sync::atomic::Ordering::Relaxed,
);
}
Ok(())
}
}
// ── Pass-2 TypedSink ──────────────────────────────────────────────────────────
//
// Workers build WayRecord + RelationRecord per segment (in parallel, including
// WKB encode via the GpuBackend). The collector pushes them through an
// Arrow-builder accumulator and ships full row-group RecordBatches over a
// channel to a dedicated `ParquetSink` writer thread — the actual parquet
// encode + zstd + disk write happens off the hot path.
//
// Without this routing the collector calls `WayWriter::push` per row, which
// hits `writer.write(&batch)` synchronously every ROW_GROUP_SIZE rows and
// pegged pass-2 at ~5 cores on planet (workers blocked on collector_tx.send
// while collector was inside parquet flush). See backlog.md / commit history.
struct Pass2TypedSink<'a> {
store: Arc<NodeStore>,
backend: &'a dyn GpuBackend,
/// Tiny accumulator for boundary-stub ways parsed ON THE COLLECTOR (a few
/// rows per segment boundary); full groups are encoded with `way_encoder`
/// here, but this is negligible work next to the workers' bulk.
way_acc: Option<writer::WayAccumulator>,
rel_acc: Option<writer::RelAccumulator>,
/// Column encoders: encode the (rare) stub batches on the collector and are
/// the SAME encoders the workers hold for the bulk row groups.
way_encoder: Option<writer::ColumnEncoder>,
rel_encoder: Option<writer::ColumnEncoder>,
/// In-order row-group stitchers (no compression here). They append both the
/// workers' pre-encoded bulk row groups and the collector's stub groups.
way_writer: Option<writer::ParallelRowGroupWriter>,
rel_writer: Option<writer::ParallelRowGroupWriter>,
way_count: usize,
rel_count: usize,
pb: indicatif::ProgressBar,
prev_stub: Vec<u8>,
tables: crate::Tables,
busy_ns: Option<Arc<std::sync::atomic::AtomicU64>>,
/// Region clip active — boundary-stub ways with zero in-region coords are
/// dropped here too, matching the worker path.
clip_active: bool,
/// Way-resolution loss (audit finding (a)) for STUB ways — the elements
/// that straddle a decode-slot boundary resolve HERE on the collector, not
/// in a worker, so they need their own tap into the same two counters the
/// workers feed (the boundary-stub clip leak taught exactly this lesson).
unresolved_refs: Arc<AtomicU64>,
shortened_ways: Arc<AtomicU64>,
/// The relation clip prefilter's carry, gathered from every segment and settled
/// in [`Self::finish_sieved`]. Empty on every path but a clipped convert.
kept_way_ids: Vec<i64>,
kept_rel_ids: Vec<i64>,
pending_rels: Vec<crate::reader::PendingRelation>,
rels_dropped: u64,
/// The prefilter's setup for the elements the COLLECTOR parses — the boundary
/// stubs of the XML readers. PBF has no stubs; an XML relation that straddles a
/// decode-slot boundary resolves here and must meet the same oracle.
sieve: Option<P2SieveSetup>,
}
/// What pass 2 leaves behind besides its row counts — the relation clip prefilter's
/// oracle input and its ledger.
#[derive(Default)]
pub(crate) struct Pass2Done {
pub ways: usize,
pub rels: usize,
/// The ids of every way written (unsorted, as the workers reported them).
/// Only collected when the codec's sieve asked for it.
pub kept_way_ids: Vec<i64>,
/// Relations the oracle dropped on the workers — no member survived the clip.
pub rels_dropped: u64,
/// Relations the collector's fixpoint promoted — kept only via a relation member.
pub rels_via_relations: u64,
/// Relations the fixpoint could NOT promote and therefore dropped.
pub rels_pending_dropped: u64,
}
impl Pass2TypedSink<'_> {
/// Append the pre-encoded bulk row groups the worker produced for one
/// segment — pure stitch + I/O, no compression on the collector.
fn append_segment_groups(&mut self, seg: &mut P2Segment) -> Result<()> {
if let Some(w) = self.way_writer.as_mut() {
for rg in std::mem::take(&mut seg.way_row_groups) {
w.append(rg)?;
}
}
if let Some(w) = self.rel_writer.as_mut() {
for rg in std::mem::take(&mut seg.rel_row_groups) {
w.append(rg)?;
}
}
self.way_count += seg.way_rows;
self.rel_count += seg.rel_rows;
self.kept_way_ids.append(&mut seg.kept_way_ids);
self.kept_rel_ids.append(&mut seg.kept_rel_ids);
self.pending_rels.append(&mut seg.pending_rels);
self.rels_dropped += seg.rels_dropped;
Ok(())
}
/// [`gatling::TypedSink::finish`] plus the prefilter's carry. The pending
/// relations are settled HERE, after every worker has reported the relations it
/// kept, by the same fixpoint `relgeom` runs — then written through the
/// collector's own accumulator as the last row groups of the file.
fn finish_sieved(mut self) -> Result<Pass2Done> {
let stub = std::mem::take(&mut self.prev_stub);
self.flush_stub(&stub)?;
let mut done = Pass2Done::default();
if !self.pending_rels.is_empty() {
let mut kept: std::collections::HashSet<i64> =
self.kept_rel_ids.iter().copied().collect();
let pending = std::mem::take(&mut self.pending_rels);
let n_pending = pending.len() as u64;
let promoted = crate::reader::RelationClipOracle::settle(pending, &mut kept);
done.rels_via_relations = promoted.len() as u64;
done.rels_pending_dropped = n_pending - done.rels_via_relations;
if let (Some(acc), Some(enc), Some(w)) = (
self.rel_acc.as_mut(),
self.rel_encoder.as_ref(),
self.rel_writer.as_mut(),
) {
for r in &promoted {
acc.push(r);
self.rel_count += 1;
if let Some(batch) = acc.take_if_full()? {
w.append(enc.encode(&batch)?)?;
}
}
}
}
done.rels_dropped = self.rels_dropped + done.rels_pending_dropped;
done.kept_way_ids = std::mem::take(&mut self.kept_way_ids);
// Drain the collector's stub remainders.
if let (Some(acc), Some(enc), Some(w)) = (
self.way_acc.as_mut(),
self.way_encoder.as_ref(),
self.way_writer.as_mut(),
) {
if let Some(batch) = acc.take_remaining()? {
w.append(enc.encode(&batch)?)?;
}
}
if let (Some(acc), Some(enc), Some(w)) = (
self.rel_acc.as_mut(),
self.rel_encoder.as_ref(),
self.rel_writer.as_mut(),
) {
if let Some(batch) = acc.take_remaining()? {
w.append(enc.encode(&batch)?)?;
}
}
if let Some(w) = self.way_writer.take() {
w.finish()?;
}
if let Some(w) = self.rel_writer.take() {
w.finish()?;
}
done.ways = self.way_count;
done.rels = self.rel_count;
Ok(done)
}
fn flush_stub(&mut self, stub: &[u8]) -> Result<()> {
if stub.is_empty() {
return Ok(());
}
let mut elems: Vec<ElemIndex> = Vec::new();
xml_vtd::build_elem_index_slice(stub, 0, &mut |e| {
if pass2_wants(self.tables, &e.kind) {
elems.push(e);
}
});
if elems.is_empty() {
return Ok(());
}
let refs: Vec<&ElemIndex> = elems.iter().collect();
let oracle = self.sieve.as_ref().and_then(|sv| sv.oracle.as_deref());
let (ways, rels, loss, rel_sieve) =
parse_chunk(stub, &refs, &self.store, 0, self.clip_active, oracle);
// Stub ways resolve on the collector — count their loss through the
// SAME counters the workers feed, or a boundary-straddling way could
// lose a ref invisibly (the exact silence this audit fix ends).
self.unresolved_refs
.fetch_add(loss.unresolved_refs, Ordering::Relaxed);
self.shortened_ways
.fetch_add(loss.shortened_ways, Ordering::Relaxed);
let ways = if self.tables.ways {
collected_to_way_records(ways, self.backend)
} else {
Vec::new()
};
// The stub's share of the prefilter's carry, same books as a worker segment.
if let Some(sv) = &self.sieve {
if sv.collect_way_ids && self.tables.ways {
self.kept_way_ids.extend(ways.iter().map(|w| w.id));
}
if sv.oracle.is_some() && self.tables.rels {
self.kept_rel_ids.extend(rels.iter().map(|r| r.id));
self.pending_rels.extend(rel_sieve.pending);
self.rels_dropped += rel_sieve.dropped;
}
}
self.way_count += ways.len();
self.rel_count += rels.len();
if let (Some(acc), Some(enc), Some(w)) = (
self.way_acc.as_mut(),
self.way_encoder.as_ref(),
self.way_writer.as_mut(),
) {
for r in &ways {
acc.push(r);
}
while let Some(batch) = acc.take_if_full()? {
w.append(enc.encode(&batch)?)?;
}
}
if let (Some(acc), Some(enc), Some(w)) = (
self.rel_acc.as_mut(),
self.rel_encoder.as_ref(),
self.rel_writer.as_mut(),
) {
for r in &rels {
acc.push(r);
}
while let Some(batch) = acc.take_if_full()? {
w.append(enc.encode(&batch)?)?;
}
}
Ok(())
}
fn finish(self) -> Result<(usize, usize)> {
let done = self.finish_sieved()?;
Ok((done.ways, done.rels))
}
}
/// Build a parallel pass-2 sink (workers encode, collector stitches) together
/// with the shared way/rel column encoders to hand to the codec's workers.
/// Shared by all four resolved readers (pbf / bz2 / gz / osm), through [`run_pass2`].
#[allow(clippy::too_many_arguments)]
fn build_pass2_sink<'a>(
out_dir: &Path,
compression: &str,
store: &Arc<NodeStore>,
backend: &'a dyn GpuBackend,
tables: crate::Tables,
pb: indicatif::ProgressBar,
busy_ns: Option<Arc<std::sync::atomic::AtomicU64>>,
clip_active: bool,
unresolved_refs: Arc<AtomicU64>,
shortened_ways: Arc<AtomicU64>,
sieve: Option<P2SieveSetup>,
) -> Result<(
Pass2TypedSink<'a>,
Option<writer::ColumnEncoder>,
Option<writer::ColumnEncoder>,
)> {
let (way_encoder, way_writer) = if tables.ways {
let (e, w) = writer::parallel_way_writer_with(
&out_dir.join("ways.parquet"),
compression,
tables.node_refs,
)?;
(Some(e), Some(w))
} else {
(None, None)
};
let (rel_encoder, rel_writer) = if tables.rels {
let (e, w) = writer::parallel_rel_writer(&out_dir.join("relations.parquet"), compression)?;
(Some(e), Some(w))
} else {
(None, None)
};
let way_acc = if tables.ways {
Some(writer::WayAccumulator::with_node_refs(tables.node_refs))
} else {
None
};
let rel_acc = if tables.rels {
Some(writer::RelAccumulator::new())
} else {
None
};
let sink = Pass2TypedSink {
store: Arc::clone(store),
backend,
way_acc,
rel_acc,
way_encoder: way_encoder.clone(),
rel_encoder: rel_encoder.clone(),
way_writer,
rel_writer,
way_count: 0,
rel_count: 0,
pb,
prev_stub: Vec::new(),
tables,
busy_ns,
clip_active,
unresolved_refs,
shortened_ways,
kept_way_ids: Vec::new(),
kept_rel_ids: Vec::new(),
pending_rels: Vec::new(),
rels_dropped: 0,
sieve,
};
Ok((sink, way_encoder, rel_encoder))
}
// ── THE relation clip prefilter's pass 2, for every reader ───────────────────
//
// With a clip active, pass 2 runs as TWO sub-passes over the same input: ways first,
// then relations. The way sub-pass collects the ids it kept; those and the node store
// are the oracle the relation sub-pass judges every relation against — kept when ANY
// member survived the clip, the same admission `relgeom::resolve_relations` applies.
// Without a clip every relation is kept and there is one pass, byte-identical to
// before the prefilter existed.
//
// This is the ONE orchestration (LAW 5). What differs per format is only how the
// relation bytes are reached — a PBF opens at the first relation blob's frame, a
// `.osm` seeks to its first `<relation`, a `.osm.bz2` block-seeks to its last `<way `
// block, a `.osm.gz` cannot seek and re-decodes from byte 0 — and that is the
// closure's whole job. The sink, the encoders, the sieve wiring, the fixpoint and
// the ledger live here, once.
//
// MEASURED on the live Sweden cut 2026-08-31 (PBF, lane R): 14 632 472 relations
// written for a clip whose resolve stage kept 274 206 — 1.005 GB of `members` +
// `tags` per cut for rows that could never resolve. The XML readers — including the
// `.osm.bz2` planet cut oden runs — wrote the same rows until 2026-09-01; see
// `tests/relation_clip_prefilter.rs` and `.nornir/osm-katana-v2-design.md`.
/// Which sub-pass a codec is being built for.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum SubPass {
/// No clip: ways and relations in one pass, as always.
Whole,
/// Clip, first half: ways only, collecting the ids kept.
Ways,
/// Clip, second half: relations only, judged against the oracle. The reader may
/// open the input where the relations start.
Relations,
}
/// What [`run_pass2`] hands its closure to build a codec from: the tables THIS
/// sub-pass writes, the column encoders the sink shares with the workers, and the
/// prefilter's setup.
struct Pass2Codec {
tables: crate::Tables,
way_enc: Option<writer::ColumnEncoder>,
rel_enc: Option<writer::ColumnEncoder>,
clip_active: bool,
sieve: Option<P2SieveSetup>,
}
/// What pass 2 reports up: the counts, and the prefilter's ledger as phase-log fields.
struct Pass2Report {
ways: usize,
rels: usize,
fields: Vec<(&'static str, String)>,
}
/// Run pass 2 — one pass, or the prefilter's two sub-passes — with `run` building the
/// format's codec and driving the engine over `sink` for each. `run` returns the
/// format-specific fields for the phase log (where a relation sub-pass started, …).
#[allow(clippy::too_many_arguments)]
fn run_pass2<'a>(
out_dir: &Path,
compression: &str,
store: &Arc<NodeStore>,
backend: &'a dyn GpuBackend,
tables: crate::Tables,
pb: &indicatif::ProgressBar,
busy_ns: Option<Arc<AtomicU64>>,
clip_active: bool,
cancel: &crate::CancelToken,
unresolved_refs: &Arc<AtomicU64>,
shortened_ways: &Arc<AtomicU64>,
mut run: impl FnMut(SubPass, &mut Pass2TypedSink<'a>, Pass2Codec) -> Result<Vec<(&'static str, String)>>,
) -> Result<Pass2Report> {
let sieve = clip_active && tables.rels;
let tables_a = if sieve {
crate::Tables {
rels: false,
..tables
}
} else {
tables
};
let sieve_a = sieve.then_some(P2SieveSetup {
collect_way_ids: true,
oracle: None,
});
let (mut p2, way_enc, rel_enc) = build_pass2_sink(
out_dir,
compression,
store,
backend,
tables_a,
pb.clone(),
busy_ns.clone(),
clip_active,
Arc::clone(unresolved_refs),
Arc::clone(shortened_ways),
sieve_a.clone(),
)?;
let mut fields = run(
if sieve { SubPass::Ways } else { SubPass::Whole },
&mut p2,
Pass2Codec {
tables: tables_a,
way_enc,
rel_enc,
clip_active,
sieve: sieve_a,
},
)?;
let done_a = p2.finish_sieved()?;
let way_count = done_a.ways;
let mut rel_count = done_a.rels;
fields.insert(0, ("rel_prefilter", if sieve { "clip" } else { "none" }.to_string()));
if sieve {
cancel.check("pass2", 0)?;
let oracle = Arc::new(crate::reader::RelationClipOracle::new(done_a.kept_way_ids));
pb.set_message(format!(
"pass 2/2 (relations, {} ways kept)",
oracle.way_count()
));
let tables_b = crate::Tables {
nodes: false,
ways: false,
rels: true,
node_refs: tables.node_refs,
};
let sieve_b = Some(P2SieveSetup {
collect_way_ids: false,
oracle: Some(Arc::clone(&oracle)),
});
let (mut p2b, way_enc_b, rel_enc_b) = build_pass2_sink(
out_dir,
compression,
store,
backend,
tables_b,
pb.clone(),
busy_ns,
true,
Arc::clone(unresolved_refs),
Arc::clone(shortened_ways),
sieve_b.clone(),
)?;
debug_assert!(way_enc_b.is_none(), "the relation sub-pass must not reopen ways.parquet");
fields.push(("rel_oracle_ways", oracle.way_count().to_string()));
let more = run(
SubPass::Relations,
&mut p2b,
Pass2Codec {
tables: tables_b,
way_enc: way_enc_b,
rel_enc: rel_enc_b,
clip_active: true,
sieve: sieve_b,
},
)?;
let done_b = p2b.finish_sieved()?;
rel_count = done_b.rels;
fields.extend(more);
fields.push(("rels_dropped_by_clip", done_b.rels_dropped.to_string()));
fields.push(("rels_kept_via_relations", done_b.rels_via_relations.to_string()));
}
Ok(Pass2Report {
ways: way_count,
rels: rel_count,
fields,
})
}
/// The pass-2 `done` line, one spelling for every reader: the counts, the bytes on
/// disk, `extra` (a format's own counters), the resolve-loss counters — always
/// emitted, never gated on >0, so an absent key can never be mistaken for a zero —
/// and the prefilter's ledger.
fn pass2_done_fields(
out_dir: &Path,
report: &Pass2Report,
unresolved_refs: &AtomicU64,
shortened_ways: &AtomicU64,
extra: Vec<(&'static str, String)>,
) -> Vec<(&'static str, String)> {
let bytes = |name: &str| {
std::fs::metadata(out_dir.join(name))
.map(|m| m.len())
.unwrap_or(0)
};
let mut f = vec![
("ways", report.ways.to_string()),
("rels", report.rels.to_string()),
("way_bytes", bytes("ways.parquet").to_string()),
("rel_bytes", bytes("relations.parquet").to_string()),
];
f.extend(extra);
// Way-resolution loss (audit finding (a)): a way whose node ref did not
// resolve is emitted SHORTENED — `node_refs` keeps the full authored list
// beside the shorter geometry, which is the per-way signal the schema already
// supports; dropping the way instead would discard the vertices that DID
// resolve. These two counters are the aggregate that used to not exist: 0
// unless the input references nodes it does not contain.
f.push(("unresolved_node_refs", unresolved_refs.load(Ordering::Relaxed).to_string()));
f.push(("shortened_ways", shortened_ways.load(Ordering::Relaxed).to_string()));
f.extend(report.fields.iter().cloned());
f
}
impl gatling::TypedSink<P2Segment> for Pass2TypedSink<'_> {
fn process(&mut self, mut seg: P2Segment, _is_last: bool) -> Result<()> {
let t0 = std::time::Instant::now();
let boundary = [self.prev_stub.as_slice(), seg.left_stub.as_slice()].concat();
self.flush_stub(&boundary)?;
let prog = seg.left_stub.len() + seg.way_rows * 512 + seg.right_stub.len();
self.append_segment_groups(&mut seg)?;
self.pb.inc(prog as u64);
self.prev_stub = seg.right_stub;
if let Some(b) = &self.busy_ns {
b.fetch_add(
t0.elapsed().as_nanos() as u64,
std::sync::atomic::Ordering::Relaxed,
);
}
Ok(())
}
}
/// Resolved 2-pass `.osm.bz2` → GeoParquet on the Gatling engine.
///
/// **The ONE reader of an on-disk byte figure for a phase field** (LAW 5 + LAW 2).
///
/// # The defect this replaces — `"coord_bytes":0` over a 1.65 GB file
///
/// All three convert paths (bz2, gz, pbf) carried a byte-identical copy of
///
/// ```text
/// let coord_bytes = std::fs::metadata(out_dir.join("node_coords.bin"))
/// .map(|m| m.len())
/// .unwrap_or(0);
/// ```
///
/// and `node_coords.bin` is a name the coord store stopped writing when it became a
/// single-batch Arrow IPC file. [`crate::node_store::CoordFileWriter::new`] keeps only
/// the PARENT of the path it is handed — its own comment says so, "`path` is the legacy
/// `node_coords.bin` location" — and `finalize()` writes `node_coords.arrow`. So the
/// `metadata()` call was ENOENT on every real run, and `unwrap_or(0)` turned "I could
/// not look" into a plausible-looking zero. MEASURED on the 2026-08-27 Sweden cut: the
/// phase log wrote `"coord_bytes":0` while `node_coords.arrow` beside it was
/// **1 655 818 282 B**. A gate on `coord_bytes > 0` would have failed a pass 1 that
/// worked, and a reader took the zero for "pass 1 produced no coords".
///
/// Three things had to change, because any two of them still lie:
///
/// 1. **One place asks for the name**, so a rename cannot leave three stale stat calls
/// behind (it already had, three times over).
/// 2. **A missing file reports `absent`, never `0`.** Zero now means a zero-byte file
/// and nothing else. The identity value can no longer stand in for "I could not
/// look" — that substitution is the whole shape of the defect.
/// 3. **The figure is reported by the phase that produced the artefact.** Pass 1 cannot
/// report `node_coords.arrow`: at `p1.finish()` the file does not exist yet. Pass 1
/// reports the payload it accepted; `sort_and_stree`, which runs the merge that
/// writes the file, reports the file.
fn on_disk_bytes_field(dir: &Path, name: &str) -> String {
match std::fs::metadata(dir.join(name)) {
Ok(m) => m.len().to_string(),
// NOT `"0"`. See point 2 above.
Err(e) => {
eprintln!(
" warning: {} is not on disk where the phase log promises it ({e}) — \
reporting `absent` rather than 0, because a zero here reads as \
\"nothing was written\"",
dir.join(name).display()
);
"absent".to_string()
}
}
}
/// **Pass 1's `done` fields — one declaration, three call sites** (LAW 5).
///
/// `coord_payload_bytes` is the coord payload pass 1 accepted, exactly
/// `nodes * 16`; `node_coords.arrow` is NOT reported here because it does not exist
/// yet — see [`on_disk_bytes_field`] point 3. The pbf path appends its two blob-plan
/// fields to what this returns.
fn pass1_done_fields(
out_dir: &Path,
node_count: usize,
coord_writer: &crate::node_store::CoordFileWriter,
) -> Vec<(&'static str, String)> {
vec![
("nodes", node_count.to_string()),
(
"coord_payload_bytes",
coord_writer.coord_payload_bytes().to_string(),
),
(
"parquet_bytes",
on_disk_bytes_field(out_dir, "nodes.parquet"),
),
]
}
/// **`sort_and_stree`'s `done` fields — one declaration, three call sites** (LAW 5).
///
/// This phase runs `CoordFileWriter::finalize()`, which is what writes
/// `node_coords.arrow`, so this is the first and only moment `coord_bytes` can be a
/// true statement about a file. `coord_framing_bytes` is the file minus the payload
/// pass 1 reported — 954 B on the Sweden cut — so the two phases' numbers can be
/// checked against each other and against the artefact instead of merely being printed.
fn sort_and_stree_done_fields(
out_dir: &Path,
payload_bytes: u64,
) -> Vec<(&'static str, String)> {
let coord_bytes = on_disk_bytes_field(out_dir, "node_coords.arrow");
let framing = coord_bytes
.parse::<u64>()
.ok()
.map(|b| b.saturating_sub(payload_bytes).to_string())
.unwrap_or_else(|| "absent".to_string());
vec![
("note", "psort-samplesort".to_string()),
("coord_bytes", coord_bytes),
("coord_payload_bytes", payload_bytes.to_string()),
("coord_framing_bytes", framing),
]
}
/// Pass 1 (engine pass): stream node coords to `node_coords.bin` + nodes.parquet, then
/// sort into a `NodeStore`. Pass 2 (engine pass): resolve way/relation geometry against
/// the store → ways/relations.parquet. Each pass is one parallel bz2 decode through the
/// engine. Returns `(node_count, way_count, rel_count)`.
#[allow(clippy::too_many_arguments)]
pub fn read_resolved_bz2(
input: &Path,
out_dir: &Path,
compression: &str,
tables: crate::Tables,
n_workers: usize,
skip_to_nodes: bool,
pb: indicatif::ProgressBar,
backend: &dyn GpuBackend,
sink: Option<&crate::phase_log::PhaseSink>,
clip: Option<Arc<crate::Clip>>,
cancel: &crate::CancelToken,
reuse_coords: bool,
) -> Result<(usize, usize, usize)> {
let n_workers_str = n_workers.to_string();
// ── Pass 1, or the resume that skips it ──────────────────────────────────
//
// Pass 1 exists to produce `node_coords.arrow` (+ `nodes.parquet`). If a
// previous run already produced them from THIS input under THIS clip, the
// store can be opened straight off disk — `coords_ipc::open_mmap` +
// `NodeStore::from_arrow_mmap` is exactly what `finalize()` ends with. On
// Sweden that is 1 365 s not paid a second time, which is the difference
// between a resumable chain and one that only pretends to resume.
//
// `reusable_store` refuses on anything less than a complete, provably
// matching artefact set, and names the reason. A refusal is logged rather
// than swallowed: a silent fallback to a full pass 1 is indistinguishable
// from normal work, and "resume is quietly never taken" is the exact bug
// this would otherwise grow.
let resumed = if reuse_coords {
match crate::coords_resume::reusable_store(out_dir, input, clip.as_deref(), tables.nodes) {
Ok((store, n)) => {
pb.set_message(format!(
"pass 1 skipped — {n} nodes reused from node_coords.arrow"
));
if let Some(s) = sink {
let p = s.phase("pass1_reused", &[("nodes", &n.to_string())]);
p.done(&[
("source", "node_coords.arrow".into()),
("provenance", "matched".into()),
]);
}
Some((Arc::new(store), n))
}
Err(verdict) => {
eprintln!(" pass 1 resume refused: {verdict} — running pass 1");
if let Some(s) = sink {
let p = s.phase("pass1_resume_refused", &[("why", &verdict.to_string())]);
p.done(&[("action", "full pass 1".into())]);
}
None
}
}
} else {
None
};
let (store, node_count) = if let Some((store, n)) = resumed {
(store, n)
} else {
read_resolved_bz2_pass1(
input,
out_dir,
compression,
tables,
n_workers,
skip_to_nodes,
&pb,
sink,
clip.clone(),
cancel,
&n_workers_str,
)?
};
cancel.check("pass2", 0)?;
read_resolved_bz2_pass2(
input,
out_dir,
compression,
tables,
n_workers,
skip_to_nodes,
pb,
backend,
sink,
clip,
cancel,
store,
node_count,
&n_workers_str,
)
}
/// Pass 1 proper: stream node coords + node records, then sort into a store.
///
/// Split out of [`read_resolved_bz2`] so the resume path has something to skip.
/// Returns the store and the node count, and writes the provenance record that
/// authorises a later run to skip it.
#[allow(clippy::too_many_arguments)]
fn read_resolved_bz2_pass1(
input: &Path,
out_dir: &Path,
compression: &str,
tables: crate::Tables,
n_workers: usize,
skip_to_nodes: bool,
pb: &indicatif::ProgressBar,
sink: Option<&crate::phase_log::PhaseSink>,
clip: Option<Arc<crate::Clip>>,
cancel: &crate::CancelToken,
n_workers_str: &str,
) -> Result<(Arc<NodeStore>, usize)> {
use crate::xml_to_pbf::{Bz2Prepared, Bz2Skip, bz2_prepare, mk_bz2_cfg};
// ── Pass 1: node coords + node records (TypedCodec — decode+parse per worker) ─
pb.set_message("pass 1/2 (coords)");
let p1_phase = sink.map(|s| s.phase("pass1", &[("workers", &n_workers_str)]));
let (node_encoder, node_writer) = if tables.nodes {
let (enc, w) = writer::parallel_node_writer(&out_dir.join("nodes.parquet"), compression)?;
(Some(enc), Some(w))
} else {
(None, None)
};
let node_acc = if tables.nodes {
Some(writer::NodeAccumulator::new())
} else {
None
};
let mut p1 = Pass1TypedSink {
coord_writer: CoordFileWriter::new(&out_dir.join("node_coords.bin"))?,
node_acc,
node_encoder: node_encoder.clone(),
node_writer,
node_count: 0,
pb: pb.clone(),
prev_stub: Vec::new(),
tables,
clip: clip.clone(),
busy_ns: p1_phase.as_ref().map(|p| Arc::clone(&p.busy_ns)),
};
{
let skip = if skip_to_nodes {
Bz2Skip::Changesets(crate::xml_to_pbf::DATA_NEEDLES)
} else {
Bz2Skip::None
};
let Bz2Prepared {
reader,
max_blocksize,
initial_carry,
..
} = bz2_prepare(input, skip)?;
let codec = Bz2Pass1TypedCodec {
max_blocksize,
tables,
node_encoder,
clip: clip.clone(),
};
gatling::run_typed(
reader,
codec,
&mut Checkpointed::new(&mut p1, cancel, "pass1"),
n_workers,
mk_bz2_cfg(initial_carry),
)?;
}
let (coord_writer, node_count) = p1.finish()?;
let coord_payload_bytes = coord_writer.coord_payload_bytes();
if let Some(p) = p1_phase {
p.done(&pass1_done_fields(out_dir, node_count, &coord_writer));
}
pb.set_message("sorting node coords …");
let sort_phase = sink.map(|s| s.phase("sort_and_stree", &[("nodes", &node_count.to_string())]));
// The k-way merge inside `finalize` is the ONE step with no per-work-unit
// checkpoint, so it is bracketed instead: cancelled before, we never enter it;
// cancelled during, we stop the moment it returns and pass 2 never opens a
// writer. Measured 0.862 s of a 2 664.653 s Sweden cut — 0.03 % of the wall.
cancel.check("sort_and_stree", 0)?;
let store = Arc::new(coord_writer.finalize()?);
if let Some(p) = sort_phase {
p.done(&sort_and_stree_done_fields(out_dir, coord_payload_bytes));
}
// Written only now, after `finalize()` has closed a complete
// `node_coords.arrow`. A record that exists therefore implies the artefact
// it describes finished; crash in between and there is simply no record,
// which reads as "cannot resume" — the safe verdict, and the one a resume
// must default to.
if let Ok(prov) =
crate::coords_resume::CoordsProvenance::of(input, clip.as_deref(), tables.nodes, node_count)
{
if let Err(e) = prov.write(out_dir) {
eprintln!(" warning: could not write resume provenance: {e}");
}
}
Ok((store, node_count))
}
/// Pass 2 proper: resolve way/relation geometry against the node store.
///
/// Split out of [`read_resolved_bz2`] so it can run against a store that came
/// from disk instead of from a pass 1 in the same process.
#[allow(clippy::too_many_arguments)]
fn read_resolved_bz2_pass2(
input: &Path,
out_dir: &Path,
compression: &str,
tables: crate::Tables,
n_workers: usize,
skip_to_nodes: bool,
pb: indicatif::ProgressBar,
backend: &dyn GpuBackend,
sink: Option<&crate::phase_log::PhaseSink>,
clip: Option<Arc<crate::Clip>>,
cancel: &crate::CancelToken,
store: Arc<NodeStore>,
node_count: usize,
n_workers_str: &str,
) -> Result<(usize, usize, usize)> {
use crate::xml_to_pbf::{Bz2Prepared, Bz2Skip, bz2_prepare, mk_bz2_cfg};
// ── Pass 2: ways + relations resolved (TypedCodec — decode+parse per worker) ──
// Uses seek_past_nodes (finds last <node> block + 1) rather than seeking for <way>,
// which fails on large ways that span entire bz2 blocks (no <way> opening tag in block).
pb.set_message(format!("pass 2/2 ({} nodes)", store.len()));
let p2_phase = sink.map(|s| {
s.phase(
"pass2",
&[
("workers", &n_workers_str),
("store_nodes", &store.len().to_string()),
],
)
});
let unresolved_refs = Arc::new(AtomicU64::new(0));
let shortened_ways = Arc::new(AtomicU64::new(0));
let report = run_pass2(
out_dir,
compression,
&store,
backend,
tables,
&pb,
p2_phase.as_ref().map(|p| Arc::clone(&p.busy_ns)),
clip.is_some(),
cancel,
&unresolved_refs,
&shortened_ways,
|sub, p2, c| {
let skip = match sub {
// The relation sub-pass opens the stream at the last block that holds a
// `<way ` — every block before it is node or way bytes the way sub-pass
// already decoded. This is the path oden's planet cut runs, so this
// seek is what keeps the prefilter from costing a second full decode.
SubPass::Relations => Bz2Skip::Ways,
SubPass::Whole | SubPass::Ways if skip_to_nodes => Bz2Skip::Nodes,
SubPass::Whole | SubPass::Ways => Bz2Skip::None,
};
let Bz2Prepared {
reader,
max_blocksize,
initial_carry,
start,
} = bz2_prepare(input, skip)?;
let codec = Bz2Pass2TypedCodec {
max_blocksize,
tables: c.tables,
store: Arc::clone(&store),
backend,
way_enc: c.way_enc,
rel_enc: c.rel_enc,
clip_active: c.clip_active,
unresolved_refs: Arc::clone(&unresolved_refs),
shortened_ways: Arc::clone(&shortened_ways),
sieve: c.sieve,
};
gatling::run_typed(
reader,
codec,
&mut Checkpointed::new(p2, cancel, "pass2"),
n_workers,
mk_bz2_cfg(initial_carry),
)?;
Ok(match sub {
SubPass::Relations => vec![("rel_pass_start_offset", start.to_string())],
SubPass::Whole | SubPass::Ways => Vec::new(),
})
},
)?;
if let Some(p) = p2_phase {
// The bz2 twin of the PBF path's keys, i.e. the path oden's
// planet-latest.osm.bz2 cut runs.
p.done(&pass2_done_fields(
out_dir,
&report,
&unresolved_refs,
&shortened_ways,
Vec::new(),
));
}
Ok((node_count, report.ways, report.rels))
}
/// Raw single-pass `.osm.gz` → GeoParquet on the Gatling engine.
///
/// The engine decodes DEFLATE in parallel where the gzip stream has full-flush
/// boundaries (pigz/bgzf); standard single-stream gzip decodes as one segment.
/// Hands contiguous XML to [`RawParquetSink`]. Mirrors [`read_raw_bz2`].
#[allow(clippy::too_many_arguments)]
pub fn read_raw_gz(
input: &Path,
out_dir: &Path,
compression: &str,
tables: crate::Tables,
n_workers: usize,
pb: indicatif::ProgressBar,
cancel: &crate::CancelToken,
) -> Result<(usize, usize, usize)> {
let (mut sink, encoders) = raw_typed_sink(out_dir, compression, tables, pb)?;
let codec = RawGzParquetTypedCodec {
tables,
node_encoder: encoders.node,
way_encoder: encoders.way,
rel_encoder: encoders.rel,
};
crate::xml_to_pbf::gz_gatling_run_typed(
input,
codec,
&mut Checkpointed::new(&mut sink, cancel, "pass1"),
n_workers,
)?;
sink.finish()
}
/// Resolved 2-pass `.osm.gz` → GeoParquet on the Gatling engine.
///
/// Mirrors [`read_resolved_bz2`]: pass 1 streams node coords + nodes.parquet and
/// sorts into a `NodeStore`; pass 2 resolves way/relation geometry. Each pass is
/// one parallel DEFLATE decode through the engine via the gz typed codecs.
#[allow(clippy::too_many_arguments)]
pub fn read_resolved_gz(
input: &Path,
out_dir: &Path,
compression: &str,
tables: crate::Tables,
n_workers: usize,
pb: indicatif::ProgressBar,
backend: &dyn GpuBackend,
sink: Option<&crate::phase_log::PhaseSink>,
clip: Option<Arc<crate::Clip>>,
cancel: &crate::CancelToken,
) -> Result<(usize, usize, usize)> {
let n_workers_str = n_workers.to_string();
// ── Pass 1: node coords + node records ───────────────────────────────────
pb.set_message("pass 1/2 (coords)");
let p1_phase = sink.map(|s| s.phase("pass1", &[("workers", &n_workers_str)]));
let (node_encoder, node_writer) = if tables.nodes {
let (enc, w) = writer::parallel_node_writer(&out_dir.join("nodes.parquet"), compression)?;
(Some(enc), Some(w))
} else {
(None, None)
};
let node_acc = if tables.nodes {
Some(writer::NodeAccumulator::new())
} else {
None
};
let mut p1 = Pass1TypedSink {
coord_writer: CoordFileWriter::new(&out_dir.join("node_coords.bin"))?,
node_acc,
node_encoder: node_encoder.clone(),
node_writer,
node_count: 0,
pb: pb.clone(),
prev_stub: Vec::new(),
tables,
clip: clip.clone(),
busy_ns: p1_phase.as_ref().map(|p| Arc::clone(&p.busy_ns)),
};
{
let codec = GzPass1TypedCodec {
tables,
node_encoder,
clip: clip.clone(),
};
crate::xml_to_pbf::gz_gatling_run_typed(
input,
codec,
&mut Checkpointed::new(&mut p1, cancel, "pass1"),
n_workers,
)?;
}
let (coord_writer, node_count) = p1.finish()?;
let coord_payload_bytes = coord_writer.coord_payload_bytes();
if let Some(p) = p1_phase {
p.done(&pass1_done_fields(out_dir, node_count, &coord_writer));
}
pb.set_message("sorting node coords …");
let sort_phase = sink.map(|s| s.phase("sort_and_stree", &[("nodes", &node_count.to_string())]));
// The k-way merge inside `finalize` is the ONE step with no per-work-unit
// checkpoint, so it is bracketed instead: cancelled before, we never enter it;
// cancelled during, we stop the moment it returns and pass 2 never opens a
// writer. Measured 0.862 s of a 2 664.653 s Sweden cut — 0.03 % of the wall.
cancel.check("sort_and_stree", 0)?;
let store = Arc::new(coord_writer.finalize()?);
if let Some(p) = sort_phase {
p.done(&sort_and_stree_done_fields(out_dir, coord_payload_bytes));
}
cancel.check("pass2", 0)?;
// ── Pass 2: ways + relations resolved ────────────────────────────────────
pb.set_message(format!("pass 2/2 ({} nodes)", store.len()));
let p2_phase = sink.map(|s| {
s.phase(
"pass2",
&[
("workers", &n_workers_str),
("store_nodes", &store.len().to_string()),
],
)
});
let unresolved_refs = Arc::new(AtomicU64::new(0));
let shortened_ways = Arc::new(AtomicU64::new(0));
let report = run_pass2(
out_dir,
compression,
&store,
backend,
tables,
&pb,
p2_phase.as_ref().map(|p| Arc::clone(&p.busy_ns)),
clip.is_some(),
cancel,
&unresolved_refs,
&shortened_ways,
|_sub, p2, c| {
// A gzip stream cannot be opened mid-way, so the relation sub-pass
// re-decodes from byte 0 and indexes only `<relation>` elements
// (`pass2_wants`). One extra DEFLATE decode of the file, paid only under
// a clip; a `.osm.gz` planet is not the production path (bz2 is).
let codec = GzPass2TypedCodec {
tables: c.tables,
store: Arc::clone(&store),
backend,
way_enc: c.way_enc,
rel_enc: c.rel_enc,
clip_active: c.clip_active,
unresolved_refs: Arc::clone(&unresolved_refs),
shortened_ways: Arc::clone(&shortened_ways),
sieve: c.sieve,
};
crate::xml_to_pbf::gz_gatling_run_typed(
input,
codec,
&mut Checkpointed::new(p2, cancel, "pass2"),
n_workers,
)?;
Ok(Vec::new())
},
)?;
if let Some(p) = p2_phase {
p.done(&pass2_done_fields(
out_dir,
&report,
&unresolved_refs,
&shortened_ways,
Vec::new(),
));
}
Ok((node_count, report.ways, report.rels))
}
/// Is the pass-2 blob plan in effect? `OSM_KATANA_PASS2_PLAN=0` turns it off.
///
/// **Bench seam, not a product knob.** Turning it off is never faster and never
/// changes a single output row; the only caller that wants it is
/// `benches/pbf_pass2_plan.rs`, which has to be able to run the pre-change
/// behaviour to prove the change is worth anything.
///
/// Read on every convert rather than cached: the bench runs both arms in ONE
/// process, and a `OnceLock` would silently give the second arm the first arm's
/// setting — a bench that always measures the same thing twice and reports 1.00×
/// as a result.
fn pass2_plan_enabled() -> bool {
!matches!(
std::env::var("OSM_KATANA_PASS2_PLAN").as_deref(),
Ok("0") | Ok("off") | Ok("false")
)
}
/// Resolved 2-pass `.pbf` → GeoParquet on the Gatling engine (NO rayon).
///
/// Each pass reads the file once; `split` cuts the byte stream at PBF blob frame
/// boundaries and workers zlib+protobuf-decode one blob each. Pass 1 fuses node
/// coord extraction (→ `node_coords.bin`, sorted into a `NodeStore`) and
/// nodes.parquet encoding (per-worker row groups); pass 2 resolves way/relation
/// geometry against the store. Reuses the shared [`Pass1TypedSink`] /
/// [`Pass2TypedSink`], so PBF now shares the exact coord-sort / node-store path
/// as the bz2 and gz converters. Returns `(node_count, way_count, rel_count)`.
#[allow(clippy::too_many_arguments)]
pub fn read_resolved_pbf(
input: &Path,
out_dir: &Path,
compression: &str,
tables: crate::Tables,
n_workers: usize,
pb: indicatif::ProgressBar,
backend: &dyn GpuBackend,
sink: Option<&crate::phase_log::PhaseSink>,
clip: Option<Arc<crate::Clip>>,
cancel: &crate::CancelToken,
) -> Result<(usize, usize, usize)> {
let n_workers_str = n_workers.to_string();
// ── Pass 1: node coords + node records (decode+parse per worker) ─────────
pb.set_message("pass 1/2 (coords)");
let p1_phase = sink.map(|s| s.phase("pass1", &[("workers", &n_workers_str)]));
let (node_encoder, node_writer) = if tables.nodes {
let (enc, w) = writer::parallel_node_writer(&out_dir.join("nodes.parquet"), compression)?;
(Some(enc), Some(w))
} else {
(None, None)
};
let node_acc = if tables.nodes {
Some(writer::NodeAccumulator::new())
} else {
None
};
let mut p1 = Pass1TypedSink {
coord_writer: CoordFileWriter::new(&out_dir.join("node_coords.bin"))?,
node_acc,
node_encoder: node_encoder.clone(),
node_writer,
node_count: 0,
pb: pb.clone(),
prev_stub: Vec::new(),
tables,
clip: clip.clone(),
busy_ns: p1_phase.as_ref().map(|p| Arc::clone(&p.busy_ns)),
};
// The pass-2 plan, filled in by the pass-1 workers. Pass 1 inflates and
// protobuf-parses every blob anyway; asking each parsed block whether it
// carries a way or relation group is one pointer walk, and it buys pass 2
// the right to leave the node-only blobs closed.
let plan_builder = Arc::new(crate::pbf_io::BlobPlanBuilder::new());
{
let codec = PbfPass1TypedCodec {
tables,
node_encoder,
clip: clip.clone(),
cursor: crate::xml_to_pbf::SlotCursor::new(),
plan: Arc::clone(&plan_builder),
};
crate::xml_to_pbf::pbf_gatling_run_typed(
input,
codec,
&mut Checkpointed::new(&mut p1, cancel, "pass1"),
n_workers,
)?;
}
let plan = Arc::new(
Arc::try_unwrap(plan_builder)
.map(crate::pbf_io::BlobPlanBuilder::finish)
// A live clone can only mean a worker outlived the engine, which it
// cannot; the fallback is the safe direction anyway (open everything).
.unwrap_or_else(|_| crate::pbf_io::BlobPlan::open_everything()),
);
// The bench's regression arm. A speed claim needs the slower arm on the same
// bytes, the same machine and the same minute, and this is the only way to
// get the pre-plan behaviour out of the shipped code rather than out of a
// remembered number. Off unless explicitly set; see `benches/pbf_pass2_plan.rs`.
let plan = if pass2_plan_enabled() {
plan
} else {
Arc::new(crate::pbf_io::BlobPlan::open_everything())
};
let (coord_writer, node_count) = p1.finish()?;
let coord_payload_bytes = coord_writer.coord_payload_bytes();
if let Some(p) = p1_phase {
let mut f = pass1_done_fields(out_dir, node_count, &coord_writer);
f.push(("blobs", plan.seen_len().to_string()));
f.push(("blobs_with_way_or_rel", plan.wanted_len().to_string()));
p.done(&f);
}
pb.set_message("sorting node coords …");
let sort_phase = sink.map(|s| s.phase("sort_and_stree", &[("nodes", &node_count.to_string())]));
// The k-way merge inside `finalize` is the ONE step with no per-work-unit
// checkpoint, so it is bracketed instead: cancelled before, we never enter it;
// cancelled during, we stop the moment it returns and pass 2 never opens a
// writer. Measured 0.862 s of a 2 664.653 s Sweden cut — 0.03 % of the wall.
cancel.check("sort_and_stree", 0)?;
let store = Arc::new(coord_writer.finalize()?);
if let Some(p) = sort_phase {
p.done(&sort_and_stree_done_fields(out_dir, coord_payload_bytes));
}
cancel.check("pass2", 0)?;
// ── Pass 2: ways + relations resolved (decode+parse per worker) ──────────
pb.set_message(format!("pass 2/2 ({} nodes)", store.len()));
let p2_phase = sink.map(|s| {
s.phase(
"pass2",
&[
("workers", &n_workers_str),
("store_nodes", &store.len().to_string()),
],
)
});
let opened = Arc::new(AtomicU64::new(0));
let skipped = Arc::new(AtomicU64::new(0));
let unresolved_refs = Arc::new(AtomicU64::new(0));
let shortened_ways = Arc::new(AtomicU64::new(0));
// The relation clip prefilter's two sub-passes are `run_pass2`'s; what is PBF's
// is where the relation sub-pass opens the file: at the frame of the first blob
// pass 1 saw a relation group in, so it inflates none of the way blobs the way
// sub-pass already opened. `None` means pass 1 saw no relation blob (or the plan
// is empty): start at 0 and let the plan open everything, as before.
let report = run_pass2(
out_dir,
compression,
&store,
backend,
tables,
&pb,
p2_phase.as_ref().map(|p| Arc::clone(&p.busy_ns)),
clip.is_some(),
cancel,
&unresolved_refs,
&shortened_ways,
|sub, p2, c| {
let start = match sub {
SubPass::Relations => plan.first_rel_blob().unwrap_or(0),
SubPass::Whole | SubPass::Ways => 0,
};
let codec = PbfPass2TypedCodec {
tables: c.tables,
store: Arc::clone(&store),
way_enc: c.way_enc,
rel_enc: c.rel_enc,
clip_active: c.clip_active,
cursor: crate::xml_to_pbf::SlotCursor::at(start),
plan: Arc::clone(&plan),
opened: Arc::clone(&opened),
skipped: Arc::clone(&skipped),
unresolved_refs: Arc::clone(&unresolved_refs),
shortened_ways: Arc::clone(&shortened_ways),
sieve: c.sieve,
};
crate::xml_to_pbf::pbf_gatling_run_typed_from(
input,
codec,
&mut Checkpointed::new(p2, cancel, "pass2"),
n_workers,
start,
)?;
Ok(match sub {
SubPass::Relations => vec![
("rel_pass_start_offset", start.to_string()),
("rel_blobs_planned", plan.rel_len().to_string()),
],
SubPass::Whole | SubPass::Ways => Vec::new(),
})
},
)?;
The file goes on; the console shows all of it.