use geopackage_core::geometry::GpbGeometry;
use geopackage_core::gpb;
use geopackage_core::ident::quote;
use geopackage_core::triggers;
use rusqlite::Connection;
use rusqlite::types::ValueRef;
use crate::Result;
use crate::packed::{self, NodeSink};
pub const DEFAULT_BULK_THRESHOLD: usize = 10_000;
pub const DEFAULT_FILL_FACTOR: f64 = 1.0;
#[derive(Debug, Clone, Copy)]
#[non_exhaustive]
pub struct BulkIndexOptions {
pub bulk_threshold: usize,
pub structural_check: StructuralCheck,
pub fill_factor: f64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
#[non_exhaustive]
pub enum StructuralCheck {
#[default]
RtreeOnly,
FullDatabase,
}
impl Default for BulkIndexOptions {
fn default() -> Self {
Self {
bulk_threshold: DEFAULT_BULK_THRESHOLD,
structural_check: StructuralCheck::RtreeOnly,
fill_factor: DEFAULT_FILL_FACTOR,
}
}
}
impl BulkIndexOptions {
pub fn with_threshold(bulk_threshold: usize) -> Self {
Self {
bulk_threshold,
..Self::default()
}
}
pub fn always_bulk() -> Self {
Self::with_threshold(0)
}
pub fn never_bulk() -> Self {
Self::with_threshold(usize::MAX)
}
#[must_use]
pub fn with_structural_check(mut self, structural_check: StructuralCheck) -> Self {
self.structural_check = structural_check;
self
}
#[must_use]
pub fn with_fill_factor(mut self, fill_factor: f64) -> Self {
self.fill_factor = fill_factor;
self
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum BuildPath {
Triggered,
Bulk,
TriggeredFallback,
}
pub(crate) type TestFault = fn(&Connection, &str) -> Result<()>;
pub(crate) fn no_fault(_: &Connection, _: &str) -> Result<()> {
Ok(())
}
struct ShadowTables<'c> {
conn: &'c Connection,
node_sql: String,
rowid_sql: String,
parent_sql: String,
rowid_buffer: Vec<(i64, i64)>,
}
impl ShadowTables<'_> {
fn flush(&mut self) -> Result<()> {
self.rowid_buffer.sort_unstable_by_key(|&(rowid, _)| rowid);
let mut stmt = self.conn.prepare_cached(&self.rowid_sql)?;
for &(rowid, nodeno) in &self.rowid_buffer {
stmt.execute(rusqlite::params![rowid, nodeno])?;
}
Ok(())
}
}
impl NodeSink for ShadowTables<'_> {
fn node(&mut self, nodeno: i64, blob: &[u8]) -> Result<()> {
self.conn
.prepare_cached(&self.node_sql)?
.execute(rusqlite::params![nodeno, blob])?;
Ok(())
}
fn rowid(&mut self, rowid: i64, nodeno: i64) -> Result<()> {
self.rowid_buffer.push((rowid, nodeno));
Ok(())
}
fn parent(&mut self, nodeno: i64, parentnode: i64) -> Result<()> {
self.conn
.prepare_cached(&self.parent_sql)?
.execute(rusqlite::params![nodeno, parentnode])?;
Ok(())
}
}
fn node_size(conn: &Connection, rtree: &str) -> Result<usize> {
let size: i64 = conn.query_row(
&format!(
"SELECT length(data) FROM {} WHERE nodeno = 1",
quote(&format!("{rtree}_node"))?
),
[],
|r| r.get(0),
)?;
Ok(usize::try_from(size).unwrap_or(0))
}
fn write_packed(
conn: &Connection,
rtree: &str,
entries: &[(i64, [f64; 4])],
node_size: usize,
fill_factor: f64,
) -> Result<()> {
let node_table = quote(&format!("{rtree}_node"))?;
let rowid_table = quote(&format!("{rtree}_rowid"))?;
let parent_table = quote(&format!("{rtree}_parent"))?;
conn.execute_batch(&format!(
"DELETE FROM {node_table}; DELETE FROM {rowid_table}; DELETE FROM {parent_table};"
))?;
let mut sink = ShadowTables {
conn,
node_sql: format!("INSERT INTO {node_table} VALUES (?1, ?2)"),
rowid_sql: format!("INSERT INTO {rowid_table} VALUES (?1, ?2)"),
parent_sql: format!("INSERT INTO {parent_table} VALUES (?1, ?2)"),
rowid_buffer: Vec::with_capacity(entries.len()),
};
packed::pack_into(entries, node_size, fill_factor, &mut sink)?;
sink.flush()
}
pub(crate) fn table_row_count(conn: &Connection, table: &str) -> Result<usize> {
let count: i64 = conn.query_row(
&format!("SELECT count(*) FROM {}", quote(table)?),
[],
|r| r.get(0),
)?;
Ok(usize::try_from(count).unwrap_or(usize::MAX))
}
fn envelope_of(blob: &[u8]) -> std::result::Result<Option<[f64; 4]>, geopackage_core::Error> {
let (header, _) = gpb::parse_header(blob)?;
if header.empty {
return Ok(None);
}
let Some(body_bounds) = GpbGeometry::parse(blob)?.xy_envelope() else {
return Ok(None);
};
let bounds = match header.envelope.xy_bounds() {
Some((min_x, max_x, min_y, max_y)) => [min_x, max_x, min_y, max_y],
None => body_bounds,
};
Ok(Some(bounds))
}
fn accumulate_envelopes(
conn: &Connection,
table: &str,
geom: &str,
pk: &str,
) -> Result<Vec<(i64, [f64; 4])>> {
let (t, c, i) = (quote(table)?, quote(geom)?, quote(pk)?);
let sql = format!("SELECT {i}, {c} FROM {t} WHERE {c} NOT NULL");
let mut stmt = conn.prepare(&sql)?;
let mut rows = stmt.query([])?;
let mut out = Vec::with_capacity(table_row_count(conn, table)?);
while let Some(row) = rows.next()? {
let id: i64 = row.get(0)?;
let blob = match row.get_ref(1)? {
ValueRef::Blob(b) => b,
_ => {
return Err(rusqlite::Error::InvalidColumnType(
1,
geom.to_string(),
rusqlite::types::Type::Blob,
)
.into());
}
};
if let Some(bounds) = envelope_of(blob)? {
out.push((id, bounds));
}
}
Ok(out)
}
fn gate(
conn: &Connection,
rtree: &str,
mut expected: Vec<(i64, [f64; 4])>,
structural_check: StructuralCheck,
) -> Result<bool> {
let quoted = quote(rtree)?;
let count: i64 = conn.query_row(&format!("SELECT count(*) FROM {quoted}"), [], |r| r.get(0))?;
if usize::try_from(count).unwrap_or(usize::MAX) != expected.len() {
return Ok(false);
}
expected.sort_unstable_by_key(|&(id, _)| id);
if expected
.iter()
.zip(expected.iter().skip(1))
.any(|(a, b)| a.0 == b.0)
{
return Ok(false);
}
let mut matched = vec![false; expected.len()];
let mut hits = 0usize;
let mut stmt = conn.prepare(&format!("SELECT id, minx, maxx, miny, maxy FROM {quoted}"))?;
let mut rows = stmt.query([])?;
while let Some(row) = rows.next()? {
let id: i64 = row.get(0)?;
let (s_min_x, s_max_x, s_min_y, s_max_y): (f64, f64, f64, f64) =
(row.get(1)?, row.get(2)?, row.get(3)?, row.get(4)?);
let Ok(at) = expected.binary_search_by_key(&id, |&(id, _)| id) else {
return Ok(false);
};
let (Some(seen), Some(&(_, [min_x, max_x, min_y, max_y]))) =
(matched.get_mut(at), expected.get(at))
else {
return Ok(false);
};
if std::mem::replace(seen, true) {
return Ok(false);
}
hits += 1;
if !(s_min_x <= min_x && s_max_x >= max_x && s_min_y <= min_y && s_max_y >= max_y) {
return Ok(false);
}
}
if hits != expected.len() {
return Ok(false);
}
let rtree_report: String = conn.query_row("SELECT rtreecheck(?1)", [rtree], |r| r.get(0))?;
if rtree_report != "ok" {
return Ok(false);
}
if structural_check == StructuralCheck::FullDatabase {
let integrity: String = conn.query_row("PRAGMA integrity_check", [], |r| r.get(0))?;
return Ok(integrity == "ok");
}
Ok(true)
}
#[expect(
clippy::too_many_arguments,
reason = "internal build entry point threading the whole build context; a parameter struct would be used by these two call sites alone"
)]
pub(crate) fn fill_index<F>(
conn: &Connection,
table: &str,
geom: &str,
pk: &str,
rtree: &str,
options: BulkIndexOptions,
precomputed: Option<Vec<(i64, [f64; 4])>>,
fault: TestFault,
after: F,
) -> Result<BuildPath>
where
F: FnOnce(&Connection) -> Result<()>,
{
let tx = conn.unchecked_transaction()?;
let path = fill_index_in_transaction(
&tx,
table,
geom,
pk,
rtree,
options,
precomputed,
fault,
after,
)?;
tx.commit()?;
Ok(path)
}
#[expect(
clippy::too_many_arguments,
reason = "internal build entry point threading the whole build context; a parameter struct would be used by these two call sites alone"
)]
pub(crate) fn fill_index_in_transaction<F>(
conn: &Connection,
table: &str,
geom: &str,
pk: &str,
rtree: &str,
options: BulkIndexOptions,
precomputed: Option<Vec<(i64, [f64; 4])>>,
fault: TestFault,
after: F,
) -> Result<BuildPath>
where
F: FnOnce(&Connection) -> Result<()>,
{
let accumulated = match precomputed {
Some(entries) => entries,
None => accumulate_envelopes(conn, table, geom, pk)?,
};
let quoted_rtree = quote(rtree)?;
let create_vtab = triggers::create_rtree_table_sql(table, geom)?;
conn.execute_batch(&format!("DROP TABLE IF EXISTS {quoted_rtree}"))?;
conn.execute_batch(&create_vtab)?;
let node_size = node_size(conn, rtree)?;
write_packed(conn, rtree, &accumulated, node_size, options.fill_factor)?;
fault(conn, rtree)?;
let path = if gate(conn, rtree, accumulated, options.structural_check)? {
BuildPath::Bulk
} else {
conn.execute_batch(&format!("DROP TABLE {quoted_rtree}"))?;
conn.execute_batch(&create_vtab)?;
conn.execute_batch(&triggers::populate_rtree_sql(table, geom, pk)?)?;
BuildPath::TriggeredFallback
};
after(conn)?;
Ok(path)
}