use crate::interfaces::cli::exec::stages::{stage, stage_context};
use crate::osm_pbf_file::data_opts;
use crate::osm_tag::tag_policy;
use crate::database::table;
use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
use std::io::Write;
use std::time::Instant;
fn resolve_buffer_bytes(buffer_limit_in_mb: Option<u64>) -> usize {
const MIB: u64 = 1024 * 1024;
let mb = buffer_limit_in_mb.unwrap_or_else(|| {
let mut sys = sysinfo::System::new();
sys.refresh_memory();
let total_bytes = sys.total_memory();
let target_bytes = total_bytes * 60 / 100;
let mb = target_bytes / MIB;
println!(
"\x1b[1;36mbuffer-limit-in-mb\x1b[0m default = {mb} MiB (60% of {} MiB total ram)",
total_bytes / MIB
);
mb
});
(mb * MIB) as usize
}
fn fmt_count(n: usize) -> String {
if n >= 1_000_000 {
format!("{:.1}M", n as f64 / 1_000_000.0)
} else if n >= 1_000 {
format!("{:.1}K", n as f64 / 1_000.0)
} else {
n.to_string()
}
}
fn fmt_bytes(n: usize) -> String {
const KIB: f64 = 1024.0;
const MIB: f64 = KIB * 1024.0;
const GIB: f64 = MIB * 1024.0;
let f = n as f64;
if f >= GIB {
format!("{:.2}GiB", f / GIB)
} else if f >= MIB {
format!("{:.1}MiB", f / MIB)
} else if f >= KIB {
format!("{:.0}KiB", f / KIB)
} else {
format!("{n}B")
}
}
#[allow(clippy::too_many_arguments)]
pub fn command_handler_extract_osm_pbf_data(
data_path: &str,
threads: &u8,
inputs: &[String],
include_relations: bool,
include_ways: bool,
include_nodes: bool,
ignore_info: bool,
tags_include_list: Option<String>,
tags_ignore_list: Option<String>,
recreate: bool,
buffer_limit_in_mb: Option<u64>,
) {
let buffer_bytes = resolve_buffer_bytes(buffer_limit_in_mb);
let conn = crate::database::store::open_osm_data_write(data_path);
for (i, input) in inputs.iter().enumerate() {
let Some(resolved) = super::resolve_input(&conn, data_path, i, input) else {
continue;
};
if recreate {
crate::database::store::destroy_source(&conn, resolved.id, false);
}
let file = crate::osm_pbf_file::osm_pbf_file::open(Some(&conn), data_path);
stage::extract_osm_pbf_data.require(&stage_context {
osm_data: Some(&conn),
file: Some(&resolved),
bases: vec![],
});
let osm_pbf_file_path = &resolved.path;
let fname = &resolved.name;
println!("\x1b[1;32mfile\x1b[0m {osm_pbf_file_path}");
print!("\x1b[1;32mloading\x1b[0m blob-chunks index...");
let _ = std::io::stdout().flush();
let chunk_count = crate::osm_pbf_file::blob_index::count_by_file_id(&conn, resolved.id);
println!(" done ({chunk_count} chunks)");
let decoder_threads = threads.saturating_sub(1).max(1);
let multi = MultiProgress::with_draw_target(crate::interfaces::cli::progress_draw_target());
let decoder_bar = multi.add(ProgressBar::new_spinner());
decoder_bar.set_style(
ProgressStyle::with_template(
"{prefix:.bold.green} {msg:<80} [{bar:20.green/white}] {percent:>3}% {pos:>6}/{len:<6} chunks {per_sec} eta {eta}",
)
.unwrap()
.progress_chars("=> "),
);
decoder_bar.set_prefix(format!("decode {decoder_threads:>2}t"));
decoder_bar.set_message(fname.clone());
decoder_bar.enable_steady_tick(std::time::Duration::from_millis(100));
let writer_bar = multi.add(ProgressBar::new(0));
writer_bar.set_style(
ProgressStyle::with_template(
"{prefix:.bold.cyan} {msg:<80} [{bar:20.cyan/blue}] {percent:>3}% {human_pos:>6}/{human_len:<6} rows {per_sec} eta {eta}",
)
.unwrap()
.progress_chars("=> "),
);
writer_bar.set_prefix(format!("writer {:>2}t", 1));
writer_bar.set_message("nodes: 0 ways: 0 rel: 0 flushes: 0 bytes: 0B");
writer_bar.enable_steady_tick(std::time::Duration::from_millis(100));
let start = Instant::now();
let tags_include = tags_include_list
.as_deref()
.map(|s| s.split(',').map(|t| t.trim().to_string()).collect());
let tags_ignore = tags_ignore_list
.as_deref()
.map(|s| s.split(',').map(|t| t.trim().to_string()).collect());
let decoder_bar_cb = decoder_bar.clone();
let writer_bar_cb = writer_bar.clone();
let write_conn = crate::database::store::open_osm_data(data_path);
let counts = file
.extract_osm_data(
osm_pbf_file_path,
write_conn,
data_opts {
include_nodes,
include_ways,
include_relations,
ignore_info,
tags: tag_policy {
include: tags_include,
ignore: tags_ignore,
},
buffer_bytes,
},
*threads,
move |p| {
if decoder_bar_cb.length().is_none() {
decoder_bar_cb.set_length(p.total_chunks as u64);
}
if !decoder_bar_cb.is_finished() {
decoder_bar_cb.set_position(p.chunks_done as u64);
decoder_bar_cb.set_message(format!(
"nodes: {:>8} ways: {:>7} rel: {:>6}",
fmt_count(p.node_count),
fmt_count(p.way_count),
fmt_count(p.relation_count),
));
if p.chunks_done >= p.total_chunks {
decoder_bar_cb.finish();
}
}
let rows_decoded = p.node_count + p.way_count + p.relation_count;
let rows_written = p.nodes_written + p.ways_written + p.relations_written;
writer_bar_cb.set_length(rows_decoded as u64);
writer_bar_cb.set_position(rows_written as u64);
writer_bar_cb.set_message(format!(
"nodes: {:>8} ways: {:>7} rel: {:>6} flushes: {:>4} bytes: {:>9}",
fmt_count(p.nodes_written),
fmt_count(p.ways_written),
fmt_count(p.relations_written),
p.flushes_done,
fmt_bytes(p.bytes_flushed),
));
},
)
.expect("the blob chunk index was checked before decoding");
decoder_bar.finish();
writer_bar.finish();
let elapsed = start.elapsed().as_secs_f64();
println!(
"\x1b[1;32mextracted\x1b[0m {fname} nodes: {} ways: {} relations: {} in {elapsed:.1}s ×{threads} threads",
fmt_count(counts.nodes),
fmt_count(counts.ways),
fmt_count(counts.relations),
);
}
index_elements(&conn);
}
fn index_elements(conn: &rusqlite::Connection) {
let bar = ProgressBar::new_spinner();
bar.set_draw_target(crate::interfaces::cli::progress_draw_target());
bar.set_style(ProgressStyle::with_template("{prefix:.bold.green} {msg} {spinner}").unwrap());
bar.set_prefix("indexing");
bar.set_message("osm_nodes");
bar.enable_steady_tick(std::time::Duration::from_millis(100));
let start = Instant::now();
crate::osm_node::osm_nodes::create_indexes(conn);
bar.disable_steady_tick();
println!(
"\x1b[1;32mindexed\x1b[0m osm_nodes in {:.1}s",
start.elapsed().as_secs_f64()
);
bar.set_message("osm_ways");
bar.enable_steady_tick(std::time::Duration::from_millis(100));
let start = Instant::now();
crate::osm_way::osm_ways::create_indexes(conn);
bar.disable_steady_tick();
println!(
"\x1b[1;32mindexed\x1b[0m osm_ways in {:.1}s",
start.elapsed().as_secs_f64()
);
bar.set_message("osm_relations");
bar.enable_steady_tick(std::time::Duration::from_millis(100));
let start = Instant::now();
crate::osm_relation::osm_relations::create_indexes(conn);
bar.finish_and_clear();
println!(
"\x1b[1;32mindexed\x1b[0m osm_relations in {:.1}s",
start.elapsed().as_secs_f64()
);
}