use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use arrow_array::builder::{
BinaryBuilder, BooleanBuilder, Date32Builder, Float64Builder, Int64Builder, StringBuilder,
TimestampMicrosecondBuilder,
};
use arrow_array::cast::AsArray;
use arrow_array::{Array, ArrayRef, RecordBatch, RecordBatchReader};
use arrow_schema::{ArrowError, DataType, Field, Schema, SchemaRef, TimeUnit};
use geopackage_core::datetime::{Date, DateTime};
use geopackage_core::gpb;
use geopackage_core::ident::quote;
use geopackage_core::types::{ColumnType, GeometryType};
use rusqlite::Connection;
use rusqlite::functions::{Aggregate, Context, FunctionFlags};
use rusqlite::limits::Limit;
use rusqlite::types::ValueRef;
use crate::schema::{Column, GeometryColumn};
use crate::value::DateTimeParsing;
use crate::{Error, Layer, Result};
pub const DEFAULT_BATCH_SIZE: usize = 65_536;
pub const DEFAULT_MAX_BATCH_BYTES: usize = i32::MAX as usize;
#[must_use]
pub fn default_max_batch_bytes() -> usize {
static RESOLVED: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
*RESOLVED.get_or_init(|| {
let mut system = sysinfo::System::new();
system.refresh_memory();
let quarter = usize::try_from(system.total_memory() / 4).unwrap_or(usize::MAX);
if quarter == 0 {
DEFAULT_MAX_BATCH_BYTES
} else {
quarter.min(DEFAULT_MAX_BATCH_BYTES)
}
})
}
const EXTENSION_NAME_KEY: &str = "ARROW:extension:name";
const EXTENSION_METADATA_KEY: &str = "ARROW:extension:metadata";
const GEOARROW_WKB: &str = "geoarrow.wkb";
const DATETIME_UNIT: TimeUnit = TimeUnit::Microsecond;
const KEY_ALIAS: &str = "__gpkg_key";
const DEFAULT_MAX_THREADS: usize = 4;
impl Layer<'_> {
pub fn arrow_schema(&self) -> Result<SchemaRef> {
let geometry = self.geometry_column();
let fields: Vec<Field> = self
.schema()
.columns
.iter()
.map(|column| field_for(column, geometry))
.collect();
Ok(Arc::new(Schema::new(fields)))
}
}
fn field_for(column: &Column, geometry: Option<&GeometryColumn>) -> Field {
let is_geometry = geometry.is_some_and(|g| g.column_name == column.name);
if is_geometry {
let srs_id = geometry.map_or(0, |g| g.srs_id);
return geometry_field(&column.name, column.not_null, srs_id);
}
let nullable = !column.not_null && !column.is_primary_key();
Field::new(&column.name, data_type_for(column), nullable)
}
fn data_type_for(column: &Column) -> DataType {
let Some(declared) = &column.column_type else {
return affinity_type(&column.declared_type);
};
match declared {
ColumnType::Boolean => DataType::Boolean,
ColumnType::TinyInt
| ColumnType::SmallInt
| ColumnType::MediumInt
| ColumnType::Integer => DataType::Int64,
ColumnType::Float | ColumnType::Double => DataType::Float64,
ColumnType::Text(_) => DataType::Utf8,
ColumnType::Blob(_) => DataType::Binary,
ColumnType::Date => DataType::Date32,
ColumnType::DateTime => DataType::Timestamp(DATETIME_UNIT, Some("UTC".into())),
ColumnType::Geometry(_) => DataType::Binary,
_ => affinity_type(&column.declared_type),
}
}
fn affinity_type(declared: &str) -> DataType {
let declared = declared.to_ascii_uppercase();
let has = |needle: &str| declared.contains(needle);
if has("INT") {
DataType::Int64
} else if has("CHAR") || has("CLOB") || has("TEXT") {
DataType::Utf8
} else if has("BLOB") || declared.is_empty() {
DataType::Binary
} else {
DataType::Float64
}
}
fn geometry_field(name: &str, not_null: bool, srs_id: i32) -> Field {
let mut metadata = HashMap::new();
metadata.insert(EXTENSION_NAME_KEY.to_owned(), GEOARROW_WKB.to_owned());
metadata.insert(EXTENSION_METADATA_KEY.to_owned(), crs_metadata(srs_id));
Field::new(name, DataType::Binary, !not_null).with_metadata(metadata)
}
fn crs_metadata(srs_id: i32) -> String {
if srs_id <= 0 {
return "{}".to_owned();
}
epsg_utils::epsg_to_projjson(srs_id).map_or_else(
|_| format!(r#"{{"crs":"EPSG:{srs_id}","crs_type":"authority_code"}}"#),
|projjson| format!(r#"{{"crs":{projjson},"crs_type":"projjson"}}"#),
)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn the_default_ceiling_never_exceeds_the_offset_limit() {
let resolved = default_max_batch_bytes();
assert!(resolved > 0, "a zero ceiling would put one row in a batch");
assert!(
resolved <= DEFAULT_MAX_BATCH_BYTES,
"i32 offsets cannot address {resolved} bytes"
);
assert_eq!(resolved, default_max_batch_bytes());
}
#[test]
fn a_caller_cannot_raise_the_ceiling_past_the_offset_limit() {
let options = ArrowReadOptions::default().with_max_batch_bytes(usize::MAX);
assert_eq!(options.max_batch_bytes, DEFAULT_MAX_BATCH_BYTES);
}
#[test]
fn epsg_code_reads_the_crs_id_not_a_nested_one() {
let metadata = crs_metadata(4326);
assert!(
metadata.contains(r#""code":6422"#),
"the trap this guards against has moved, update the test: {metadata}"
);
assert_eq!(epsg_code(&metadata), Some(4326));
}
#[test]
fn epsg_code_reads_the_authority_code_form() {
assert_eq!(
epsg_code(r#"{"crs":"EPSG:27700","crs_type":"authority_code"}"#),
Some(27700)
);
}
#[test]
fn epsg_code_declines_what_it_cannot_identify() {
assert_eq!(epsg_code("{}"), None);
assert_eq!(epsg_code("not json"), None);
assert_eq!(
epsg_code(r#"{"crs":{"id":{"authority":"ESRI","code":104305}}}"#),
None
);
}
#[test]
fn a_layer_srs_round_trips_through_the_metadata() {
for code in [4326, 27700, 32630, 4979] {
assert_eq!(epsg_code(&crs_metadata(code)), Some(code), "code {code}");
}
}
#[test]
fn affinity_follows_sqlite_rules() {
assert_eq!(affinity_type("VARCHAR(20)"), DataType::Utf8);
assert_eq!(affinity_type("NVARCHAR(100)"), DataType::Utf8);
assert_eq!(affinity_type("CLOB"), DataType::Utf8);
assert_eq!(affinity_type("BIGINT"), DataType::Int64);
assert_eq!(affinity_type("UNSIGNED BIG INT"), DataType::Int64);
assert_eq!(affinity_type(""), DataType::Binary);
assert_eq!(affinity_type("DOUBLE PRECISION"), DataType::Float64);
assert_eq!(affinity_type("NUMERIC"), DataType::Float64);
assert_eq!(affinity_type("DECIMAL(10,5)"), DataType::Float64);
}
#[test]
fn int_wins_over_the_text_family() {
assert_eq!(affinity_type("INTCHAR"), DataType::Int64);
}
#[test]
fn char_family_wins_over_blob() {
assert_eq!(affinity_type("TEXTBLOB"), DataType::Utf8);
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub struct ArrowReadOptions {
pub batch_size: usize,
pub max_batch_bytes: usize,
pub threads: usize,
}
impl Default for ArrowReadOptions {
fn default() -> Self {
Self {
batch_size: DEFAULT_BATCH_SIZE,
max_batch_bytes: default_max_batch_bytes(),
threads: 0,
}
}
}
impl ArrowReadOptions {
pub fn with_batch_size(batch_size: usize) -> Self {
Self {
batch_size: batch_size.max(1),
..Self::default()
}
}
#[must_use]
pub fn with_max_batch_bytes(mut self, max_batch_bytes: usize) -> Self {
self.max_batch_bytes = max_batch_bytes.clamp(1, DEFAULT_MAX_BATCH_BYTES);
self
}
#[must_use]
pub fn with_threads(mut self, threads: usize) -> Self {
self.threads = threads;
self
}
fn resolved_threads(self) -> usize {
if self.threads > 0 {
return self.threads;
}
std::thread::available_parallelism()
.map(std::num::NonZeroUsize::get)
.unwrap_or(1)
.min(DEFAULT_MAX_THREADS)
}
}
impl Layer<'_> {
pub fn read_arrow(&self, options: ArrowReadOptions) -> Result<ArrowBatches<'_>> {
let sequential = self.read_arrow_sequential(options)?;
if options.resolved_threads() < 2 {
return Ok(sequential);
}
let Some(parallel) = self.parallel_source(options)? else {
return Ok(sequential);
};
Ok(ArrowBatches {
schema: sequential.schema,
source: BatchSource::Parallel(parallel),
})
}
fn read_arrow_sequential(&self, options: ArrowReadOptions) -> Result<ArrowBatches<'_>> {
let schema = self.arrow_schema()?;
let key = match self.primary_key_column() {
Some(pk) => quote(pk)?,
None => "rowid".to_owned(),
};
let key_field = self.primary_key_column().and_then(|pk| {
schema
.fields()
.iter()
.position(|field| *field.name() == *pk)
});
let mut selected = String::new();
for field in schema.fields() {
if !selected.is_empty() {
selected.push(',');
}
selected.push_str("e(field.name())?);
}
let (row_columns, aggregate_arguments) = match key_field {
Some(_) => (selected.clone(), selected),
None => (
format!("{key} AS \"{KEY_ALIAS}\",{selected}"),
format!("\"{KEY_ALIAS}\",{selected}"),
),
};
let table = quote(self.table_name())?;
let rows_sql =
format!("SELECT {row_columns} FROM {table} WHERE {key} >= ?1 ORDER BY {key} LIMIT ?2");
let sql = rows_sql.clone();
let geometry_index = self.geometry_column().and_then(|geom| {
schema
.fields()
.iter()
.position(|field| *field.name() == geom.column_name)
});
let conn = self.gpkg().connection();
let datetime = self.conversion_options().datetime;
let names: Vec<String> = schema
.fields()
.iter()
.map(|field| field.name().clone())
.collect();
let arg_count =
i32::try_from(names.len() + usize::from(key_field.is_none())).unwrap_or(i32::MAX);
let aggregate = if arg_count <= conn.limit(Limit::SQLITE_LIMIT_FUNCTION_ARG)? {
Some(AggregateState::register(
conn,
arg_count,
BatchFiller {
names: names.clone(),
types: schema
.fields()
.iter()
.map(|field| field.data_type().clone())
.collect(),
key_argument: key_field.unwrap_or(0),
field_offset: usize::from(key_field.is_none()),
geometry_index,
datetime,
capacity: options.batch_size.clamp(1, DEFAULT_BATCH_SIZE),
max_bytes: options.max_batch_bytes.clamp(1, DEFAULT_MAX_BATCH_BYTES),
output: Arc::new(Mutex::new(None)),
failure: Arc::new(Mutex::new(None)),
},
)?)
} else {
None
};
let aggregate_sql = aggregate
.as_ref()
.map(|state| {
format!(
"SELECT {}({aggregate_arguments}) FROM ({rows_sql})",
state.name
)
})
.unwrap_or_default();
Ok(ArrowBatches {
schema: Arc::clone(&schema),
source: BatchSource::Sequential(SequentialBatches {
conn,
schema,
sql,
aggregate_sql,
key_field,
geometry_index,
names,
datetime,
batch_size: options.batch_size.max(1),
max_batch_bytes: options.max_batch_bytes.clamp(1, DEFAULT_MAX_BATCH_BYTES),
last_batch_rows: 0,
next_key: i64::MIN,
exhausted: false,
aggregate,
}),
})
}
fn parallel_source(&self, options: ArrowReadOptions) -> Result<Option<ParallelBatches>> {
let Some(path) = database_path(self.gpkg().connection())? else {
return Ok(None);
};
let Some(key) = self.primary_key_column() else {
return Ok(None);
};
let Some(span) = dense_key_span(self.gpkg().connection(), self.table_name(), key)? else {
return Ok(None);
};
let batch_size = options.batch_size.max(1);
let rows = span.1.saturating_sub(span.0).saturating_add(1);
if rows < i64::try_from(batch_size.saturating_mul(2)).unwrap_or(i64::MAX) {
return Ok(None);
}
Ok(Some(ParallelBatches::spawn(
path,
self.table_name().to_owned(),
self.conversion_options(),
span,
batch_size,
options.max_batch_bytes.clamp(1, DEFAULT_MAX_BATCH_BYTES),
options.resolved_threads(),
)))
}
}
fn database_path(conn: &Connection) -> Result<Option<std::path::PathBuf>> {
let file: String = conn.query_row(
"SELECT file FROM pragma_database_list WHERE name = 'main'",
[],
|row| row.get(0),
)?;
if file.is_empty() {
return Ok(None);
}
Ok(Some(std::path::PathBuf::from(file)))
}
fn dense_key_span(conn: &Connection, table: &str, key: &str) -> Result<Option<(i64, i64)>> {
let sql = format!(
"SELECT min({key}), max({key}), count(*) FROM {table}",
key = quote(key)?,
table = quote(table)?
);
let (min, max, count): (Option<i64>, Option<i64>, i64) =
conn.query_row(&sql, [], |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)))?;
let (Some(min), Some(max)) = (min, max) else {
return Ok(None);
};
let span = max.checked_sub(min).and_then(|d| d.checked_add(1));
if span != Some(count) {
return Ok(None);
}
Ok(Some((min, max)))
}
struct ParallelBatches {
receivers: Vec<std::sync::mpsc::Receiver<std::result::Result<WorkerMessage, ArrowError>>>,
workers: Vec<std::thread::JoinHandle<()>>,
turn: usize,
done: bool,
}
impl ParallelBatches {
fn spawn(
path: std::path::PathBuf,
table: String,
conversion: crate::ConversionOptions,
(first, last): (i64, i64),
batch_size: usize,
max_batch_bytes: usize,
threads: usize,
) -> Self {
let mut receivers = Vec::with_capacity(threads);
let mut workers = Vec::with_capacity(threads);
for worker in 0..threads {
let (tx, rx) = std::sync::mpsc::sync_channel(1);
let path = path.clone();
let table = table.clone();
let handle = std::thread::spawn(move || {
run_worker(
&path,
&table,
conversion,
first,
last,
batch_size,
max_batch_bytes,
threads,
worker,
&tx,
);
});
receivers.push(rx);
workers.push(handle);
}
Self {
receivers,
workers,
turn: 0,
done: false,
}
}
}
impl Iterator for ParallelBatches {
type Item = std::result::Result<RecordBatch, ArrowError>;
fn next(&mut self) -> Option<Self::Item> {
if self.done {
return None;
}
loop {
match self.receivers.get(self.turn)?.recv() {
Ok(Ok(WorkerMessage::Batch(batch))) => return Some(Ok(batch)),
Ok(Ok(WorkerMessage::WindowEnd)) => {
self.turn = (self.turn + 1) % self.receivers.len().max(1);
}
Ok(Err(error)) => {
self.done = true;
return Some(Err(error));
}
Err(_) => {
self.done = true;
return None;
}
}
}
}
}
impl Drop for ParallelBatches {
fn drop(&mut self) {
self.receivers.clear();
for worker in self.workers.drain(..) {
drop(worker.join());
}
}
}
#[expect(
clippy::too_many_arguments,
reason = "a worker's whole context, passed once at spawn; a struct would be used by this call site alone"
)]
fn run_worker(
path: &std::path::Path,
table: &str,
conversion: crate::ConversionOptions,
first: i64,
last: i64,
batch_size: usize,
max_batch_bytes: usize,
threads: usize,
worker: usize,
tx: &std::sync::mpsc::SyncSender<std::result::Result<WorkerMessage, ArrowError>>,
) {
let send_error = |error: Error| {
drop(tx.send(Err(ArrowError::ExternalError(Box::new(error)))));
};
let gpkg = match crate::GeoPackage::open_read_only(path) {
Ok(gpkg) => gpkg,
Err(error) => return send_error(error),
};
let layer = match gpkg.layer(table) {
Ok(layer) => layer.with_conversion_options(conversion),
Err(error) => return send_error(error),
};
let options = ArrowReadOptions::with_batch_size(batch_size)
.with_threads(1)
.with_max_batch_bytes(max_batch_bytes);
let mut batches = match layer.read_arrow(options) {
Ok(batches) => batches,
Err(error) => return send_error(error),
};
let BatchSource::Sequential(source) = &mut batches.source else {
return;
};
let stride = match i64::try_from(batch_size.saturating_mul(threads)) {
Ok(stride) if stride > 0 => stride,
_ => return,
};
let start = match i64::try_from(batch_size.saturating_mul(worker)) {
Ok(offset) => match first.checked_add(offset) {
Some(start) => start,
None => return,
},
Err(_) => return,
};
let mut key = start;
while key <= last {
let mut remaining = batch_size;
let mut at = key;
while remaining > 0 {
match source.read_batch_at(at, remaining) {
Ok(Some(batch)) => {
let rows = source.last_batch_rows;
if tx.send(Ok(WorkerMessage::Batch(batch))).is_err() {
return; }
if rows == 0 {
break;
}
remaining -= rows.min(remaining);
at = source.next_key;
}
Ok(None) => return,
Err(error) => return send_error(error),
}
}
if tx.send(Ok(WorkerMessage::WindowEnd)).is_err() {
return; }
match key.checked_add(stride) {
Some(next) => key = next,
None => return,
}
}
}
enum WorkerMessage {
Batch(RecordBatch),
WindowEnd,
}
struct AggregateState {
name: String,
arg_count: i32,
output: Arc<Mutex<Option<FilledBatch>>>,
failure: Arc<Mutex<Option<Error>>>,
}
impl AggregateState {
fn register(conn: &Connection, arg_count: i32, filler: BatchFiller) -> Result<Self> {
static NEXT: AtomicU64 = AtomicU64::new(0);
let name = format!(
"geopackage_fill_arrow_{}",
NEXT.fetch_add(1, Ordering::Relaxed)
);
let output = Arc::clone(&filler.output);
let failure = Arc::clone(&filler.failure);
conn.create_aggregate_function(
name.as_str(),
arg_count,
FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
filler,
)?;
Ok(Self {
name,
arg_count,
output,
failure,
})
}
}
struct BatchFiller {
names: Vec<String>,
types: Vec<DataType>,
key_argument: usize,
field_offset: usize,
geometry_index: Option<usize>,
datetime: DateTimeParsing,
capacity: usize,
max_bytes: usize,
output: Arc<Mutex<Option<FilledBatch>>>,
failure: Arc<Mutex<Option<Error>>>,
}
struct FilledBatch {
builders: Vec<ColumnBuilder>,
rows: usize,
last_key: Option<i64>,
bytes: usize,
truncated: bool,
}
impl Aggregate<FilledBatch, i64> for BatchFiller {
fn init(&self, _: &mut Context<'_>) -> rusqlite::Result<FilledBatch> {
let mut builders = Vec::with_capacity(self.types.len());
for (index, data_type) in self.types.iter().enumerate() {
let is_geometry = Some(index) == self.geometry_index;
builders.push(
ColumnBuilder::new(data_type, is_geometry, self.capacity)
.map_err(|e| rusqlite::Error::UserFunctionError(Box::new(e)))?,
);
}
Ok(FilledBatch {
builders,
rows: 0,
last_key: None,
bytes: 0,
truncated: false,
})
}
fn step(&self, ctx: &mut Context<'_>, acc: &mut FilledBatch) -> rusqlite::Result<()> {
if acc.truncated {
return Ok(());
}
let geometry_bytes = self.geometry_index.map_or(0, |index| {
match ctx.get_raw(index + self.field_offset) {
ValueRef::Blob(blob) => blob.len(),
_ => 0,
}
});
if acc.rows > 0 && acc.bytes.saturating_add(geometry_bytes) > self.max_bytes {
acc.truncated = true;
return Ok(());
}
if let ValueRef::Integer(key) = ctx.get_raw(self.key_argument) {
acc.last_key = Some(key);
}
for (index, builder) in acc.builders.iter_mut().enumerate() {
let value = ctx.get_raw(index + self.field_offset);
if let Err(error) = builder.append(&self.names, index, value, self.datetime) {
if let Ok(mut slot) = self.failure.lock() {
*slot = Some(error);
}
return Err(rusqlite::Error::UserFunctionError(
"geopackage: columnar read failed".into(),
));
}
}
acc.rows += 1;
acc.bytes = acc.bytes.saturating_add(geometry_bytes);
Ok(())
}
fn finalize(&self, _: &mut Context<'_>, acc: Option<FilledBatch>) -> rusqlite::Result<i64> {
let rows = acc.as_ref().map_or(0, |batch| batch.rows);
if let Some(batch) = acc
&& let Ok(mut slot) = self.output.lock()
{
*slot = Some(batch);
}
Ok(i64::try_from(rows).unwrap_or(i64::MAX))
}
}
struct SequentialBatches<'a> {
conn: &'a Connection,
schema: SchemaRef,
sql: String,
aggregate_sql: String,
key_field: Option<usize>,
geometry_index: Option<usize>,
names: Vec<String>,
datetime: DateTimeParsing,
max_batch_bytes: usize,
last_batch_rows: usize,
batch_size: usize,
next_key: i64,
exhausted: bool,
aggregate: Option<AggregateState>,
}
impl Drop for SequentialBatches<'_> {
fn drop(&mut self) {
if let Some(state) = &self.aggregate {
drop(
self.conn
.remove_function(state.name.as_str(), state.arg_count),
);
}
}
}
impl SequentialBatches<'_> {
fn read_batch_at(&mut self, key: i64, limit: usize) -> Result<Option<RecordBatch>> {
self.next_key = key;
self.exhausted = false;
let full = self.batch_size;
self.batch_size = limit.max(1);
let batch = self.next_batch();
self.batch_size = full;
batch
}
fn next_batch(&mut self) -> Result<Option<RecordBatch>> {
if self.aggregate.is_some() {
return self.next_batch_aggregate();
}
self.next_batch_direct()
}
fn next_batch_aggregate(&mut self) -> Result<Option<RecordBatch>> {
let queried = self.conn.query_row(
&self.aggregate_sql,
rusqlite::params![
self.next_key,
i64::try_from(self.batch_size).unwrap_or(i64::MAX)
],
|row| row.get::<_, i64>(0),
);
if let Some(state) = &self.aggregate
&& let Ok(mut slot) = state.failure.lock()
&& let Some(error) = slot.take()
{
self.exhausted = true;
return Err(error);
}
let rows_read = queried?;
let filled = self
.aggregate
.as_ref()
.and_then(|state| state.output.lock().ok().and_then(|mut slot| slot.take()));
let Some(filled) = filled else {
self.exhausted = true;
return Ok(None);
};
let rows_read = usize::try_from(rows_read).unwrap_or(0);
if rows_read == 0 {
self.exhausted = true;
return Ok(None);
}
let rows_appended = filled.rows;
self.last_batch_rows = rows_appended;
self.advance(filled.last_key, rows_appended, filled.truncated);
let arrays: Vec<ArrayRef> = filled
.builders
.into_iter()
.map(ColumnBuilder::finish)
.collect();
Ok(Some(RecordBatch::try_new(
Arc::clone(&self.schema),
arrays,
)?))
}
fn advance(&mut self, last_key: Option<i64>, rows_read: usize, truncated: bool) {
match last_key.and_then(|key| key.checked_add(1)) {
Some(next) => self.next_key = next,
None => self.exhausted = true,
}
if rows_read < self.batch_size && !truncated {
self.exhausted = true;
}
}
fn next_batch_direct(&mut self) -> Result<Option<RecordBatch>> {
let capacity = self.batch_size.min(DEFAULT_BATCH_SIZE);
let mut builders: Vec<ColumnBuilder> = self
.schema
.fields()
.iter()
.enumerate()
.map(|(index, field)| {
ColumnBuilder::new(
field.data_type(),
Some(index) == self.geometry_index,
capacity,
)
})
.collect::<Result<_>>()?;
let mut rows_read = 0usize;
let mut last_key = None;
let mut bytes = 0usize;
let mut truncated = false;
{
let mut stmt = self.conn.prepare_cached(&self.sql)?;
let mut rows = stmt.query(rusqlite::params![
self.next_key,
i64::try_from(self.batch_size).unwrap_or(i64::MAX)
])?;
let offset = usize::from(self.key_field.is_none());
while let Some(row) = rows.next()? {
let geometry_bytes = match self.geometry_index {
Some(index) => match row.get_ref(index + offset)? {
ValueRef::Blob(blob) => blob.len(),
_ => 0,
},
None => 0,
};
if rows_read > 0 && bytes.saturating_add(geometry_bytes) > self.max_batch_bytes {
truncated = true;
break;
}
last_key = Some(row.get::<_, i64>(self.key_field.unwrap_or(0))?);
for (index, builder) in builders.iter_mut().enumerate() {
builder.append(
&self.names,
index,
row.get_ref(index + offset)?,
self.datetime,
)?;
}
rows_read += 1;
bytes = bytes.saturating_add(geometry_bytes);
}
}
if rows_read == 0 {
self.exhausted = true;
return Ok(None);
}
self.last_batch_rows = rows_read;
self.advance(last_key, rows_read, truncated);
let arrays: Vec<ArrayRef> = builders.into_iter().map(ColumnBuilder::finish).collect();
Ok(Some(RecordBatch::try_new(
Arc::clone(&self.schema),
arrays,
)?))
}
}
impl Iterator for SequentialBatches<'_> {
type Item = std::result::Result<RecordBatch, ArrowError>;
fn next(&mut self) -> Option<Self::Item> {
if self.exhausted {
return None;
}
match self.next_batch() {
Ok(Some(batch)) => Some(Ok(batch)),
Ok(None) => None,
Err(error) => {
self.exhausted = true;
Some(Err(ArrowError::ExternalError(Box::new(error))))
}
}
}
}
pub struct ArrowBatches<'a> {
schema: SchemaRef,
source: BatchSource<'a>,
}
enum BatchSource<'a> {
Sequential(SequentialBatches<'a>),
Parallel(ParallelBatches),
}
impl Iterator for ArrowBatches<'_> {
type Item = std::result::Result<RecordBatch, ArrowError>;
fn next(&mut self) -> Option<Self::Item> {
match &mut self.source {
BatchSource::Sequential(batches) => batches.next(),
BatchSource::Parallel(batches) => batches.next(),
}
}
}
impl RecordBatchReader for ArrowBatches<'_> {
fn schema(&self) -> SchemaRef {
Arc::clone(&self.schema)
}
}
enum ColumnBuilder {
Boolean(BooleanBuilder),
Int64(Int64Builder),
Float64(Float64Builder),
Utf8(StringBuilder),
Binary(BinaryBuilder),
Date32(Date32Builder),
Timestamp(TimestampMicrosecondBuilder),
Geometry(BinaryBuilder),
}
impl ColumnBuilder {
fn new(data_type: &DataType, is_geometry: bool, capacity: usize) -> Result<Self> {
const GEOMETRY_BYTES: usize = 64;
const VALUE_BYTES: usize = 16;
if is_geometry {
return Ok(Self::Geometry(BinaryBuilder::with_capacity(
capacity,
capacity * GEOMETRY_BYTES,
)));
}
Ok(match data_type {
DataType::Boolean => Self::Boolean(BooleanBuilder::with_capacity(capacity)),
DataType::Int64 => Self::Int64(Int64Builder::with_capacity(capacity)),
DataType::Float64 => Self::Float64(Float64Builder::with_capacity(capacity)),
DataType::Utf8 => Self::Utf8(StringBuilder::with_capacity(
capacity,
capacity * VALUE_BYTES,
)),
DataType::Binary => Self::Binary(BinaryBuilder::with_capacity(
capacity,
capacity * VALUE_BYTES,
)),
DataType::Date32 => Self::Date32(Date32Builder::with_capacity(capacity)),
DataType::Timestamp(_, _) => {
Self::Timestamp(TimestampMicrosecondBuilder::with_capacity(capacity))
}
other => {
return Err(Error::UnsupportedArrowType {
data_type: other.to_string(),
});
}
})
}
fn append(
&mut self,
names: &[String],
index: usize,
value: ValueRef<'_>,
datetime: DateTimeParsing,
) -> Result<()> {
if let ValueRef::Null = value {
self.append_null();
return Ok(());
}
match (self, value) {
(Self::Boolean(builder), ValueRef::Integer(int)) => builder.append_value(int != 0),
(Self::Int64(builder), ValueRef::Integer(int)) => builder.append_value(int),
(Self::Float64(builder), ValueRef::Real(real)) => builder.append_value(real),
(Self::Float64(builder), ValueRef::Integer(int)) => builder.append_value(int as f64),
(Self::Utf8(builder), ValueRef::Text(bytes)) => builder.append_value(text(bytes)?),
(Self::Binary(builder), ValueRef::Blob(bytes)) => builder.append_value(bytes),
(Self::Date32(builder), ValueRef::Text(bytes)) => {
let text = text(bytes)?;
let date = Date::parse(text).map_err(|source| Error::InvalidDateTimeValue {
column: column_name(names, index),
text: text.to_owned(),
source,
})?;
builder.append_value(date.days_since_epoch());
}
(Self::Timestamp(builder), ValueRef::Text(bytes)) => {
let text = text(bytes)?;
let parsed = match datetime {
DateTimeParsing::Strict => DateTime::parse_strict(text),
DateTimeParsing::Lenient => DateTime::parse_lenient(text),
};
let stamp = parsed.map_err(|source| Error::InvalidDateTimeValue {
column: column_name(names, index),
text: text.to_owned(),
source,
})?;
let micros =
stamp
.micros_since_epoch()
.map_err(|source| Error::InvalidDateTimeValue {
column: column_name(names, index),
text: text.to_owned(),
source,
})?;
builder.append_value(micros);
}
(Self::Geometry(builder), ValueRef::Blob(blob)) => {
let offset = gpb::body_offset(blob).map_err(geopackage_core::Error::from)?;
builder.append_value(blob.get(offset..).unwrap_or_default());
}
(builder, other) => {
return Err(Error::ArrowValueMismatch {
column: column_name(names, index),
expected: builder.type_name(),
found: storage_class(other),
});
}
}
Ok(())
}
fn append_null(&mut self) {
match self {
Self::Boolean(builder) => builder.append_null(),
Self::Int64(builder) => builder.append_null(),
Self::Float64(builder) => builder.append_null(),
Self::Utf8(builder) => builder.append_null(),
Self::Binary(builder) | Self::Geometry(builder) => builder.append_null(),
Self::Date32(builder) => builder.append_null(),
Self::Timestamp(builder) => builder.append_null(),
}
}
fn type_name(&self) -> &'static str {
match self {
Self::Boolean(_) => "Boolean",
Self::Int64(_) => "Int64",
Self::Float64(_) => "Float64",
Self::Utf8(_) => "Utf8",
Self::Binary(_) => "Binary",
Self::Date32(_) => "Date32",
Self::Timestamp(_) => "Timestamp",
Self::Geometry(_) => "Binary (geoarrow.wkb)",
}
}
fn finish(mut self) -> ArrayRef {
match &mut self {
Self::Boolean(builder) => Arc::new(builder.finish()),
Self::Int64(builder) => Arc::new(builder.finish()),
Self::Float64(builder) => Arc::new(builder.finish()),
Self::Utf8(builder) => Arc::new(builder.finish()),
Self::Binary(builder) | Self::Geometry(builder) => Arc::new(builder.finish()),
Self::Date32(builder) => Arc::new(builder.finish()),
Self::Timestamp(builder) => Arc::new(builder.finish().with_timezone(std::sync::Arc::<
str,
>::from(
"UTC"
))),
}
}
}
fn column_name(names: &[String], index: usize) -> String {
names.get(index).cloned().unwrap_or_default()
}
fn text(bytes: &[u8]) -> Result<&str> {
Ok(std::str::from_utf8(bytes).map_err(rusqlite::Error::from)?)
}
fn storage_class(value: ValueRef<'_>) -> &'static str {
match value {
ValueRef::Null => "NULL",
ValueRef::Integer(_) => "INTEGER",
ValueRef::Real(_) => "REAL",
ValueRef::Text(_) => "TEXT",
ValueRef::Blob(_) => "BLOB",
}
}
struct RowLayout {
fid: Option<usize>,
geometry: Option<usize>,
values: Vec<Option<usize>>,
}
struct ArrowRow {
batch: Arc<RecordBatch>,
layout: Arc<RowLayout>,
row: usize,
}
enum ArrowRowResult {
Row(ArrowRow),
Failed(Error),
}
impl crate::writer::WritableRow for ArrowRowResult {
fn write(self, writer: &mut crate::FeatureWriter<'_>) -> Result<(i64, Option<[f64; 4]>)> {
match self {
Self::Row(row) => row.write(writer),
Self::Failed(error) => Err(error),
}
}
}
impl crate::writer::WritableRow for ArrowRow {
fn write(self, writer: &mut crate::FeatureWriter<'_>) -> Result<(i64, Option<[f64; 4]>)> {
let fid = match self
.layout
.fid
.and_then(|index| self.batch.columns().get(index))
{
Some(column) => read_i64(column, self.row)?,
None => None,
};
let mut values = Vec::with_capacity(self.layout.values.len());
for (position, index) in self.layout.values.iter().enumerate() {
let bound = match index.and_then(|index| self.batch.columns().get(index)) {
Some(column) => bind_value(column, self.row, position, &self.batch)?,
None => rusqlite::types::ToSqlOutput::Borrowed(rusqlite::types::ValueRef::Null),
};
values.push(bound);
}
let geometry = self
.layout
.geometry
.and_then(|index| self.batch.columns().get(index));
match geometry {
Some(column) if !column.is_null(self.row) => {
let wkb = binary_at(column, self.row)?;
writer.insert_wkb_bound(fid, wkb, &values)
}
_ => writer.insert_row_bound(fid, &values).map(|fid| (fid, None)),
}
}
}
impl Layer<'_> {
pub fn write_arrow<R>(&self, batches: R, batch_size: usize) -> Result<Vec<i64>>
where
R: IntoIterator<Item = std::result::Result<RecordBatch, ArrowError>>,
{
self.write_arrow_with(batches, batch_size, crate::BulkIndexOptions::default())
}
pub fn write_arrow_with<R>(
&self,
batches: R,
batch_size: usize,
options: crate::BulkIndexOptions,
) -> Result<Vec<i64>>
where
R: IntoIterator<Item = std::result::Result<RecordBatch, ArrowError>>,
{
let geometry_column = self.geometry_column().map(|g| g.column_name.clone());
let value_columns: Vec<String> = self
.value_columns()
.iter()
.filter(|column| Some(column.name.as_str()) != self.primary_key_column())
.map(|column| column.name.clone())
.collect();
let primary_key = self.primary_key_column().map(str::to_owned);
let rows = batches.into_iter().flat_map(move |batch| {
let taken = batch.map_err(Error::Arrow).and_then(|batch| {
rows_of(
&batch,
primary_key.as_deref(),
geometry_column.as_deref(),
&value_columns,
)
});
match taken {
Ok(rows) => rows.into_iter().map(ArrowRowResult::Row).collect(),
Err(error) => vec![ArrowRowResult::Failed(error)],
}
});
self.write_all_impl(rows, batch_size, options, crate::bulk::no_fault)
}
}
fn rows_of(
batch: &RecordBatch,
primary_key: Option<&str>,
geometry: Option<&str>,
value_columns: &[String],
) -> Result<Vec<ArrowRow>> {
let schema = batch.schema();
for field in schema.fields() {
let known = Some(field.name().as_str()) == primary_key
|| Some(field.name().as_str()) == geometry
|| value_columns.iter().any(|name| name == field.name());
if !known {
return Err(Error::NoSuchColumn {
table_name: String::new(),
column_name: field.name().clone(),
});
}
}
let index_of = |name: &str| schema.fields().iter().position(|f| f.name() == name);
let layout = Arc::new(RowLayout {
fid: primary_key.and_then(index_of),
geometry: geometry.and_then(index_of),
values: value_columns.iter().map(|name| index_of(name)).collect(),
});
let batch = Arc::new(batch.clone());
Ok((0..batch.num_rows())
.map(|row| ArrowRow {
batch: Arc::clone(&batch),
layout: Arc::clone(&layout),
row,
})
.collect())
}
fn binary_at(column: &ArrayRef, row: usize) -> Result<&[u8]> {
if let Some(binary) = column.as_binary_opt::<i32>() {
return Ok(binary.value(row));
}
if let Some(binary) = column.as_binary_opt::<i64>() {
return Ok(binary.value(row));
}
Err(Error::ArrowValueMismatch {
column: String::new(),
expected: "Binary or LargeBinary",
found: "another Arrow type",
})
}
fn read_i64(column: &ArrayRef, row: usize) -> Result<Option<i64>> {
if column.is_null(row) {
return Ok(None);
}
let values = column
.as_primitive_opt::<arrow_array::types::Int64Type>()
.ok_or_else(|| Error::ArrowValueMismatch {
column: String::new(),
expected: "Int64",
found: "other",
})?;
Ok(Some(values.value(row)))
}
fn bind_value<'a>(
column: &'a ArrayRef,
row: usize,
position: usize,
batch: &RecordBatch,
) -> Result<rusqlite::types::ToSqlOutput<'a>> {
use arrow_array::types::{
Date32Type, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type,
TimestampMicrosecondType, TimestampMillisecondType,
};
use rusqlite::types::{ToSqlOutput, Value as SqlV, ValueRef};
let borrowed = |value: ValueRef<'a>| Ok(ToSqlOutput::Borrowed(value));
let owned = |value: SqlV| Ok(ToSqlOutput::Owned(value));
if column.is_null(row) {
return borrowed(ValueRef::Null);
}
let name = || {
batch
.schema()
.fields()
.get(position)
.map(|field| field.name().clone())
.unwrap_or_default()
};
let mismatch = |expected: &'static str| Error::ArrowValueMismatch {
column: name(),
expected,
found: "an array of another type",
};
let out_of_range = |source| Error::InvalidDateTimeValue {
column: name(),
text: "an Arrow date or timestamp outside the representable range".to_owned(),
source,
};
match column.data_type() {
DataType::Boolean => owned(SqlV::Integer(i64::from(
column
.as_boolean_opt()
.ok_or_else(|| mismatch("Boolean"))?
.value(row),
))),
DataType::Int8 => owned(SqlV::Integer(i64::from(
column
.as_primitive_opt::<Int8Type>()
.ok_or_else(|| mismatch("Int8"))?
.value(row),
))),
DataType::Int16 => owned(SqlV::Integer(i64::from(
column
.as_primitive_opt::<Int16Type>()
.ok_or_else(|| mismatch("Int16"))?
.value(row),
))),
DataType::Int32 => owned(SqlV::Integer(i64::from(
column
.as_primitive_opt::<Int32Type>()
.ok_or_else(|| mismatch("Int32"))?
.value(row),
))),
DataType::Int64 => borrowed(ValueRef::Integer(
column
.as_primitive_opt::<Int64Type>()
.ok_or_else(|| mismatch("Int64"))?
.value(row),
)),
DataType::Float32 => owned(SqlV::Real(f64::from(
column
.as_primitive_opt::<Float32Type>()
.ok_or_else(|| mismatch("Float32"))?
.value(row),
))),
DataType::Float64 => borrowed(ValueRef::Real(
column
.as_primitive_opt::<Float64Type>()
.ok_or_else(|| mismatch("Float64"))?
.value(row),
)),
DataType::Utf8 => borrowed(ValueRef::Text(
column
.as_string_opt::<i32>()
.ok_or_else(|| mismatch("Utf8"))?
.value(row)
.as_bytes(),
)),
DataType::LargeUtf8 => borrowed(ValueRef::Text(
column
.as_string_opt::<i64>()
.ok_or_else(|| mismatch("LargeUtf8"))?
.value(row)
.as_bytes(),
)),
DataType::Binary => borrowed(ValueRef::Blob(
column
.as_binary_opt::<i32>()
.ok_or_else(|| mismatch("Binary"))?
.value(row),
)),
DataType::LargeBinary => borrowed(ValueRef::Blob(
column
.as_binary_opt::<i64>()
.ok_or_else(|| mismatch("LargeBinary"))?
.value(row),
)),
DataType::Date32 => owned(SqlV::Text(
Date::from_days_since_epoch(
column
.as_primitive_opt::<Date32Type>()
.ok_or_else(|| mismatch("Date32"))?
.value(row),
)
.map_err(out_of_range)?
.to_string(),
)),
DataType::Timestamp(TimeUnit::Microsecond, _) => owned(SqlV::Text(
DateTime::from_micros_since_epoch(
column
.as_primitive_opt::<TimestampMicrosecondType>()
.ok_or_else(|| mismatch("Timestamp"))?
.value(row),
)
.map_err(out_of_range)?
.to_string(),
)),
DataType::Timestamp(TimeUnit::Millisecond, _) => owned(SqlV::Text(
DateTime::from_micros_since_epoch(
column
.as_primitive_opt::<TimestampMillisecondType>()
.ok_or_else(|| mismatch("Timestamp"))?
.value(row)
.saturating_mul(1_000),
)
.map_err(out_of_range)?
.to_string(),
)),
other => Err(Error::UnsupportedArrowType {
data_type: other.to_string(),
}),
}
}
impl crate::TableSchemaBuilder {
pub fn from_arrow_schema(self, schema: &Schema) -> Result<Self> {
let mut builder = self;
for field in schema.fields() {
if *field.name() == builder.primary_key_name() {
continue;
}
if field.metadata().get(EXTENSION_NAME_KEY).map(String::as_str) == Some(GEOARROW_WKB) {
let srs_id = field
.metadata()
.get(EXTENSION_METADATA_KEY)
.and_then(|json| epsg_code(json))
.unwrap_or(0);
builder = builder.geometry(
crate::GeometrySpec::new(GeometryType::Geometry, srs_id)
.column_name(field.name()),
);
continue;
}
let column_type = column_type_for(field.data_type())?;
let mut column = crate::ColumnSpec::new(field.name(), column_type);
if !field.is_nullable() {
column = column.not_null();
}
builder = builder.column(column);
}
Ok(builder)
}
}
fn column_type_for(data_type: &DataType) -> Result<ColumnType> {
Ok(match data_type {
DataType::Boolean => ColumnType::Boolean,
DataType::Int8 | DataType::UInt8 => ColumnType::TinyInt,
DataType::Int16 | DataType::UInt16 => ColumnType::SmallInt,
DataType::Int32 | DataType::UInt32 => ColumnType::MediumInt,
DataType::Int64 | DataType::UInt64 => ColumnType::Integer,
DataType::Float32 => ColumnType::Float,
DataType::Float64 => ColumnType::Double,
DataType::Utf8 | DataType::LargeUtf8 => ColumnType::Text(None),
DataType::Binary | DataType::LargeBinary => ColumnType::Blob(None),
DataType::Date32 => ColumnType::Date,
DataType::Timestamp(_, _) => ColumnType::DateTime,
other => {
return Err(Error::UnsupportedArrowType {
data_type: other.to_string(),
});
}
})
}
fn epsg_code(metadata: &str) -> Option<i32> {
let value: serde_json::Value = serde_json::from_str(metadata).ok()?;
let crs = value.get("crs")?;
if let Some(id) = crs.get("id")
&& id
.get("authority")
.and_then(serde_json::Value::as_str)
.is_some_and(|a| a.eq_ignore_ascii_case("EPSG"))
{
return id.get("code")?.as_i64()?.try_into().ok();
}
let code = crs.as_str()?.strip_prefix("EPSG:")?;
code.parse().ok()
}