use anyhow::{Context, Result};
use arrow::array::{Array, Float64Array, Int64Array, StringArray, TimestampSecondArray, TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray};
use arrow::datatypes::DataType;
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
use std::path::{Path, PathBuf};
use stt_core::projection::tile_coords_to_lonlat;
use stt_core::timestamp::{normalize_timestamp_to_ms, TimestampUnit};
use stt_core::types::{BoundingBox, TimeRange};
#[derive(Debug, Clone)]
pub enum DataSource {
GeoParquet {
path: PathBuf,
time_field: String,
time_format: String,
},
SttArchive {
path: PathBuf,
},
}
impl DataSource {
pub fn display_name(&self) -> String {
match self {
DataSource::GeoParquet { path, .. } => {
path.file_name()
.map(|n| n.to_string_lossy().to_string())
.unwrap_or_else(|| "unknown".to_string())
}
DataSource::SttArchive { path } => {
path.file_name()
.map(|n| n.to_string_lossy().to_string())
.unwrap_or_else(|| "unknown".to_string())
}
}
}
}
#[derive(Debug, Clone)]
pub struct AnalyzableFeature {
pub lon: f64,
pub lat: f64,
pub timestamp: u64,
pub geometry_type: GeometryType,
pub vertex_count: usize,
pub estimated_size: usize,
pub property_count: usize,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum GeometryType {
Point,
LineString,
Polygon,
MultiPoint,
MultiLineString,
MultiPolygon,
Unknown,
}
impl std::fmt::Display for GeometryType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
GeometryType::Point => write!(f, "Point"),
GeometryType::LineString => write!(f, "LineString"),
GeometryType::Polygon => write!(f, "Polygon"),
GeometryType::MultiPoint => write!(f, "MultiPoint"),
GeometryType::MultiLineString => write!(f, "MultiLineString"),
GeometryType::MultiPolygon => write!(f, "MultiPolygon"),
GeometryType::Unknown => write!(f, "Unknown"),
}
}
}
#[derive(Debug)]
pub struct LoadedData {
pub features: Vec<AnalyzableFeature>,
pub bounds: BoundingBox,
pub time_range: TimeRange,
}
pub fn load_data(source: &DataSource) -> Result<LoadedData> {
match source {
DataSource::GeoParquet { path, time_field, time_format } => {
load_geoparquet(path, time_field, time_format)
}
DataSource::SttArchive { path } => {
load_stt_archive(path)
}
}
}
fn load_geoparquet(path: &Path, time_field: &str, time_format: &str) -> Result<LoadedData> {
use indicatif::{ProgressBar, ProgressStyle};
let pb = ProgressBar::new_spinner();
pb.set_style(
ProgressStyle::default_spinner()
.template("{spinner:.green} {msg}")
.unwrap(),
);
pb.set_message("Loading GeoParquet file...");
let file = std::fs::File::open(path).context("Failed to open GeoParquet file")?;
let builder = ParquetRecordBatchReaderBuilder::try_new(file)?;
let schema = builder.schema().clone();
let geom_col_name = find_geometry_column(&schema)?;
let time_col_idx = schema.fields().iter().position(|f| f.name() == time_field)
.ok_or_else(|| anyhow::anyhow!("Time field '{}' not found", time_field))?;
let reader = builder.build()?;
let mut features = Vec::new();
let mut min_lon = f64::MAX;
let mut max_lon = f64::MIN;
let mut min_lat = f64::MAX;
let mut max_lat = f64::MIN;
let mut min_time = u64::MAX;
let mut max_time = u64::MIN;
for batch_result in reader {
let batch = batch_result.context("Failed to read Parquet batch")?;
let geometries = extract_geometries_from_batch(&batch, &geom_col_name)?;
let timestamps = extract_timestamps_from_batch(&batch, time_col_idx, time_format)?;
let property_count = schema.fields().len() - 2;
for i in 0..batch.num_rows() {
let (geom_type, vertex_count, lon, lat) = geometries.get(i)
.cloned()
.unwrap_or((GeometryType::Unknown, 0, 0.0, 0.0));
let timestamp = timestamps.get(i).copied().unwrap_or(0);
min_lon = min_lon.min(lon);
max_lon = max_lon.max(lon);
min_lat = min_lat.min(lat);
max_lat = max_lat.max(lat);
min_time = min_time.min(timestamp);
max_time = max_time.max(timestamp);
let estimated_size = 100 + (vertex_count * 16) + (property_count * 20);
features.push(AnalyzableFeature {
lon,
lat,
timestamp,
geometry_type: geom_type,
vertex_count,
estimated_size,
property_count,
});
}
if features.len() % 100_000 == 0 {
pb.set_message(format!("Loaded {} features...", features.len()));
}
}
pb.finish_with_message(format!("Loaded {} features", features.len()));
Ok(LoadedData {
features,
bounds: BoundingBox::new(min_lon, min_lat, max_lon, max_lat),
time_range: TimeRange::new(min_time, max_time),
})
}
fn load_stt_archive(path: &Path) -> Result<LoadedData> {
use indicatif::{ProgressBar, ProgressStyle};
use stt_core::ArchiveReader;
let pb = ProgressBar::new_spinner();
pb.set_style(
ProgressStyle::default_spinner()
.template("{spinner:.green} {msg}")
.unwrap(),
);
pb.set_message("Loading STT archive...");
let reader = ArchiveReader::open(path)?;
let metadata = reader.metadata();
let entries = reader.entries();
let mut features = Vec::new();
for entry in entries {
anyhow::ensure!(
entry.zoom < 32,
"corrupt archive: tile entry {}/{}/{} has zoom {} (must be < 32)",
entry.zoom,
entry.x,
entry.y,
entry.zoom,
);
let center = tile_coords_to_lonlat(1, 1, entry.zoom, entry.x, entry.y, 2);
let (lon, lat) = (center.x(), center.y());
features.push(AnalyzableFeature {
lon,
lat,
timestamp: entry.time_start.max(0) as u64,
geometry_type: GeometryType::Point, vertex_count: 1,
estimated_size: entry.length as usize,
property_count: 0,
});
}
pb.finish_with_message(format!("Loaded {} tile entries", features.len()));
let b = metadata.bounds;
let bounds = BoundingBox::new(b.min_lon, b.min_lat, b.max_lon, b.max_lat);
let tr = metadata.time_range;
let time_range = TimeRange::new(tr.start, tr.end);
Ok(LoadedData {
features,
bounds,
time_range,
})
}
fn find_geometry_column(schema: &arrow::datatypes::Schema) -> Result<String> {
let common_names = ["geometry", "geom", "wkb_geometry", "the_geom", "shape"];
for name in common_names {
if schema.field_with_name(name).is_ok() {
return Ok(name.to_string());
}
}
for field in schema.fields() {
if matches!(field.data_type(), DataType::Binary | DataType::LargeBinary) {
return Ok(field.name().clone());
}
}
for field in schema.fields() {
if matches!(field.data_type(), DataType::Struct(_)) {
return Ok(field.name().clone());
}
}
let has_lon = schema.field_with_name("lon").is_ok()
|| schema.field_with_name("longitude").is_ok()
|| schema.field_with_name("x").is_ok();
let has_lat = schema.field_with_name("lat").is_ok()
|| schema.field_with_name("latitude").is_ok()
|| schema.field_with_name("y").is_ok();
if has_lon && has_lat {
return Ok("__lon_lat__".to_string());
}
anyhow::bail!("Could not find geometry column in Parquet schema")
}
fn extract_geometries_from_batch(
batch: &arrow::record_batch::RecordBatch,
geom_col_name: &str,
) -> Result<Vec<(GeometryType, usize, f64, f64)>> {
let mut results = Vec::with_capacity(batch.num_rows());
if geom_col_name == "__lon_lat__" {
let lon_col = batch.column_by_name("lon")
.or_else(|| batch.column_by_name("longitude"))
.or_else(|| batch.column_by_name("x"));
let lat_col = batch.column_by_name("lat")
.or_else(|| batch.column_by_name("latitude"))
.or_else(|| batch.column_by_name("y"));
if let (Some(lon), Some(lat)) = (lon_col, lat_col) {
if let (Some(lon_arr), Some(lat_arr)) = (
lon.as_any().downcast_ref::<Float64Array>(),
lat.as_any().downcast_ref::<Float64Array>(),
) {
for i in 0..batch.num_rows() {
if lon_arr.is_valid(i) && lat_arr.is_valid(i) {
results.push((GeometryType::Point, 1, lon_arr.value(i), lat_arr.value(i)));
} else {
results.push((GeometryType::Unknown, 0, 0.0, 0.0));
}
}
return Ok(results);
}
}
anyhow::bail!("Expected lon/lat columns but could not read them");
}
let geom_col = batch.column_by_name(geom_col_name)
.ok_or_else(|| anyhow::anyhow!("Geometry column '{}' not found", geom_col_name))?;
if let Some(struct_array) = geom_col.as_any().downcast_ref::<arrow::array::StructArray>() {
let x_col = struct_array.column_by_name("x")
.or_else(|| struct_array.column_by_name("longitude"))
.or_else(|| struct_array.column_by_name("lon"));
let y_col = struct_array.column_by_name("y")
.or_else(|| struct_array.column_by_name("latitude"))
.or_else(|| struct_array.column_by_name("lat"));
if let (Some(x), Some(y)) = (x_col, y_col) {
if let (Some(x_arr), Some(y_arr)) = (
x.as_any().downcast_ref::<Float64Array>(),
y.as_any().downcast_ref::<Float64Array>(),
) {
for i in 0..batch.num_rows() {
if x_arr.is_valid(i) && y_arr.is_valid(i) {
results.push((GeometryType::Point, 1, x_arr.value(i), y_arr.value(i)));
} else {
results.push((GeometryType::Unknown, 0, 0.0, 0.0));
}
}
return Ok(results);
}
}
}
if let Some(binary_array) = geom_col.as_any().downcast_ref::<arrow::array::BinaryArray>() {
for i in 0..batch.num_rows() {
if binary_array.is_valid(i) {
let wkb = binary_array.value(i);
if let Some((geom_type, vertex_count, lon, lat)) = parse_wkb_info(wkb) {
results.push((geom_type, vertex_count, lon, lat));
} else {
results.push((GeometryType::Unknown, 0, 0.0, 0.0));
}
} else {
results.push((GeometryType::Unknown, 0, 0.0, 0.0));
}
}
return Ok(results);
}
let lon_col = batch.column_by_name("lon")
.or_else(|| batch.column_by_name("longitude"))
.or_else(|| batch.column_by_name("x"));
let lat_col = batch.column_by_name("lat")
.or_else(|| batch.column_by_name("latitude"))
.or_else(|| batch.column_by_name("y"));
if let (Some(lon), Some(lat)) = (lon_col, lat_col) {
if let (Some(lon_arr), Some(lat_arr)) = (
lon.as_any().downcast_ref::<Float64Array>(),
lat.as_any().downcast_ref::<Float64Array>(),
) {
for i in 0..batch.num_rows() {
if lon_arr.is_valid(i) && lat_arr.is_valid(i) {
results.push((GeometryType::Point, 1, lon_arr.value(i), lat_arr.value(i)));
} else {
results.push((GeometryType::Unknown, 0, 0.0, 0.0));
}
}
return Ok(results);
}
}
anyhow::bail!("Could not extract geometries from column '{}'", geom_col_name)
}
fn extract_timestamps_from_batch(
batch: &arrow::record_batch::RecordBatch,
col_idx: usize,
time_format: &str,
) -> Result<Vec<u64>> {
let column = batch.column(col_idx);
let mut timestamps = Vec::with_capacity(batch.num_rows());
macro_rules! push_ts_column {
($arr:expr, $unit:expr) => {{
for i in 0..batch.num_rows() {
if $arr.is_valid(i) {
timestamps.push(normalize_timestamp_to_ms(i, $arr.value(i), $unit)?);
} else {
timestamps.push(0);
}
}
return Ok(timestamps);
}};
}
if let Some(ts_array) = column.as_any().downcast_ref::<TimestampSecondArray>() {
push_ts_column!(ts_array, TimestampUnit::Second);
}
if let Some(ts_array) = column.as_any().downcast_ref::<TimestampMillisecondArray>() {
push_ts_column!(ts_array, TimestampUnit::Millisecond);
}
if let Some(ts_array) = column.as_any().downcast_ref::<TimestampMicrosecondArray>() {
push_ts_column!(ts_array, TimestampUnit::Microsecond);
}
if let Some(ts_array) = column.as_any().downcast_ref::<TimestampNanosecondArray>() {
push_ts_column!(ts_array, TimestampUnit::Nanosecond);
}
if let Some(int_array) = column.as_any().downcast_ref::<Int64Array>() {
let unit = match time_format {
"unix-sec" => TimestampUnit::Second,
_ => TimestampUnit::Millisecond,
};
for i in 0..batch.num_rows() {
if int_array.is_valid(i) {
timestamps.push(normalize_timestamp_to_ms(i, int_array.value(i), unit)?);
} else {
timestamps.push(0);
}
}
return Ok(timestamps);
}
if let Some(str_array) = column.as_any().downcast_ref::<StringArray>() {
for i in 0..batch.num_rows() {
if str_array.is_valid(i) {
let s = str_array.value(i);
let ts = parse_iso8601(s).unwrap_or(0);
timestamps.push(ts);
} else {
timestamps.push(0);
}
}
return Ok(timestamps);
}
anyhow::bail!("Unsupported timestamp column type")
}
fn parse_iso8601(s: &str) -> Result<u64> {
use chrono::{DateTime, NaiveDateTime};
if let Ok(dt) = s.parse::<DateTime<chrono::Utc>>() {
return Ok(dt.timestamp_millis() as u64);
}
if let Ok(dt) = NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S") {
return Ok(dt.and_utc().timestamp_millis() as u64);
}
if let Ok(date) = chrono::NaiveDate::parse_from_str(s, "%Y-%m-%d") {
let dt = date.and_hms_opt(0, 0, 0).unwrap().and_utc();
return Ok(dt.timestamp_millis() as u64);
}
anyhow::bail!("Failed to parse timestamp: {}", s)
}
const WKB_POINT: u32 = 1;
const WKB_LINESTRING: u32 = 2;
const WKB_POLYGON: u32 = 3;
const WKB_MULTIPOINT: u32 = 4;
const WKB_MULTILINESTRING: u32 = 5;
const WKB_MULTIPOLYGON: u32 = 6;
fn parse_wkb_info(wkb: &[u8]) -> Option<(GeometryType, usize, f64, f64)> {
if wkb.len() < 5 {
return None;
}
let little_endian = wkb[0] == 1;
let geom_type = if little_endian {
u32::from_le_bytes([wkb[1], wkb[2], wkb[3], wkb[4]]) % 1000
} else {
u32::from_be_bytes([wkb[1], wkb[2], wkb[3], wkb[4]]) % 1000
};
match geom_type {
WKB_POINT => parse_wkb_point_info(wkb, little_endian),
WKB_LINESTRING => parse_wkb_linestring_info(wkb, little_endian),
WKB_POLYGON => parse_wkb_polygon_info(wkb, little_endian),
WKB_MULTIPOINT => parse_wkb_multipoint_info(wkb, little_endian),
WKB_MULTILINESTRING => parse_wkb_multilinestring_info(wkb, little_endian),
WKB_MULTIPOLYGON => parse_wkb_multipolygon_info(wkb, little_endian),
_ => None,
}
}
fn read_f64(wkb: &[u8], offset: usize, little_endian: bool) -> Option<f64> {
if offset + 8 > wkb.len() {
return None;
}
let bytes: [u8; 8] = wkb[offset..offset+8].try_into().ok()?;
Some(if little_endian {
f64::from_le_bytes(bytes)
} else {
f64::from_be_bytes(bytes)
})
}
fn read_u32(wkb: &[u8], offset: usize, little_endian: bool) -> Option<u32> {
if offset + 4 > wkb.len() {
return None;
}
let bytes: [u8; 4] = wkb[offset..offset+4].try_into().ok()?;
Some(if little_endian {
u32::from_le_bytes(bytes)
} else {
u32::from_be_bytes(bytes)
})
}
fn parse_wkb_point_info(wkb: &[u8], little_endian: bool) -> Option<(GeometryType, usize, f64, f64)> {
let x = read_f64(wkb, 5, little_endian)?;
let y = read_f64(wkb, 13, little_endian)?;
Some((GeometryType::Point, 1, x, y))
}
fn parse_wkb_linestring_info(wkb: &[u8], little_endian: bool) -> Option<(GeometryType, usize, f64, f64)> {
let num_points = read_u32(wkb, 5, little_endian)? as usize;
if num_points == 0 {
return None;
}
let mut sum_x = 0.0;
let mut sum_y = 0.0;
let mut offset = 9;
for _ in 0..num_points {
let x = read_f64(wkb, offset, little_endian)?;
let y = read_f64(wkb, offset + 8, little_endian)?;
sum_x += x;
sum_y += y;
offset += 16;
}
Some((GeometryType::LineString, num_points, sum_x / num_points as f64, sum_y / num_points as f64))
}
fn parse_wkb_polygon_info(wkb: &[u8], little_endian: bool) -> Option<(GeometryType, usize, f64, f64)> {
let num_rings = read_u32(wkb, 5, little_endian)? as usize;
if num_rings == 0 {
return None;
}
let mut total_points = 0usize;
let mut sum_x = 0.0;
let mut sum_y = 0.0;
let mut offset = 9;
for ring_idx in 0..num_rings {
let num_points = read_u32(wkb, offset, little_endian)? as usize;
offset += 4;
for _ in 0..num_points {
let x = read_f64(wkb, offset, little_endian)?;
let y = read_f64(wkb, offset + 8, little_endian)?;
if ring_idx == 0 {
sum_x += x;
sum_y += y;
total_points += 1;
}
offset += 16;
}
}
let vertex_count: usize = total_points;
let centroid_x = if total_points > 0 { sum_x / total_points as f64 } else { 0.0 };
let centroid_y = if total_points > 0 { sum_y / total_points as f64 } else { 0.0 };
Some((GeometryType::Polygon, vertex_count, centroid_x, centroid_y))
}
fn parse_wkb_multipoint_info(wkb: &[u8], little_endian: bool) -> Option<(GeometryType, usize, f64, f64)> {
let num_geoms = read_u32(wkb, 5, little_endian)? as usize;
if num_geoms == 0 {
return None;
}
let mut sum_x = 0.0;
let mut sum_y = 0.0;
let mut offset = 9;
for _ in 0..num_geoms {
offset += 5; let x = read_f64(wkb, offset, little_endian)?;
let y = read_f64(wkb, offset + 8, little_endian)?;
sum_x += x;
sum_y += y;
offset += 16;
}
Some((GeometryType::MultiPoint, num_geoms, sum_x / num_geoms as f64, sum_y / num_geoms as f64))
}
fn parse_wkb_multilinestring_info(wkb: &[u8], little_endian: bool) -> Option<(GeometryType, usize, f64, f64)> {
let num_geoms = read_u32(wkb, 5, little_endian)? as usize;
if num_geoms == 0 {
return None;
}
let mut total_points = 0usize;
let mut sum_x = 0.0;
let mut sum_y = 0.0;
let mut offset = 9;
for _ in 0..num_geoms {
offset += 5; let num_points = read_u32(wkb, offset, little_endian)? as usize;
offset += 4;
for _ in 0..num_points {
let x = read_f64(wkb, offset, little_endian)?;
let y = read_f64(wkb, offset + 8, little_endian)?;
sum_x += x;
sum_y += y;
total_points += 1;
offset += 16;
}
}
let centroid_x = if total_points > 0 { sum_x / total_points as f64 } else { 0.0 };
let centroid_y = if total_points > 0 { sum_y / total_points as f64 } else { 0.0 };
Some((GeometryType::MultiLineString, total_points, centroid_x, centroid_y))
}
fn parse_wkb_multipolygon_info(wkb: &[u8], little_endian: bool) -> Option<(GeometryType, usize, f64, f64)> {
let num_geoms = read_u32(wkb, 5, little_endian)? as usize;
if num_geoms == 0 {
return None;
}
let mut total_points = 0usize;
let mut sum_x = 0.0;
let mut sum_y = 0.0;
let mut offset = 9;
for poly_idx in 0..num_geoms {
offset += 5; let num_rings = read_u32(wkb, offset, little_endian)? as usize;
offset += 4;
for ring_idx in 0..num_rings {
let num_points = read_u32(wkb, offset, little_endian)? as usize;
offset += 4;
for _ in 0..num_points {
let x = read_f64(wkb, offset, little_endian)?;
let y = read_f64(wkb, offset + 8, little_endian)?;
if poly_idx == 0 && ring_idx == 0 {
sum_x += x;
sum_y += y;
total_points += 1;
}
offset += 16;
}
}
}
let centroid_x = if total_points > 0 { sum_x / total_points as f64 } else { 0.0 };
let centroid_y = if total_points > 0 { sum_y / total_points as f64 } else { 0.0 };
Some((GeometryType::MultiPolygon, total_points, centroid_x, centroid_y))
}