1use anyhow::{Context, Result};
6use arrow::array::{Array, Float64Array, Int64Array, StringArray, TimestampSecondArray, TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray};
7use arrow::datatypes::DataType;
8use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
9use std::path::{Path, PathBuf};
10use stt_core::projection::tile_coords_to_lonlat;
11use stt_core::timestamp::{normalize_timestamp_to_ms, TimestampUnit};
12use stt_core::types::{BoundingBox, TimeRange};
13
14#[derive(Debug, Clone)]
16pub enum DataSource {
17 GeoParquet {
18 path: PathBuf,
19 time_field: String,
20 time_format: String,
21 },
22 SttArchive {
23 path: PathBuf,
24 },
25}
26
27impl DataSource {
28 pub fn display_name(&self) -> String {
29 match self {
30 DataSource::GeoParquet { path, .. } => {
31 path.file_name()
32 .map(|n| n.to_string_lossy().to_string())
33 .unwrap_or_else(|| "unknown".to_string())
34 }
35 DataSource::SttArchive { path } => {
36 path.file_name()
37 .map(|n| n.to_string_lossy().to_string())
38 .unwrap_or_else(|| "unknown".to_string())
39 }
40 }
41 }
42}
43
44#[derive(Debug, Clone)]
46pub struct AnalyzableFeature {
47 pub lon: f64,
49 pub lat: f64,
51 pub timestamp: u64,
53 pub geometry_type: GeometryType,
55 pub vertex_count: usize,
57 pub estimated_size: usize,
59 pub property_count: usize,
61}
62
63#[derive(Debug, Clone, Copy, PartialEq, Eq)]
65pub enum GeometryType {
66 Point,
67 LineString,
68 Polygon,
69 MultiPoint,
70 MultiLineString,
71 MultiPolygon,
72 Unknown,
73}
74
75impl std::fmt::Display for GeometryType {
76 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
77 match self {
78 GeometryType::Point => write!(f, "Point"),
79 GeometryType::LineString => write!(f, "LineString"),
80 GeometryType::Polygon => write!(f, "Polygon"),
81 GeometryType::MultiPoint => write!(f, "MultiPoint"),
82 GeometryType::MultiLineString => write!(f, "MultiLineString"),
83 GeometryType::MultiPolygon => write!(f, "MultiPolygon"),
84 GeometryType::Unknown => write!(f, "Unknown"),
85 }
86 }
87}
88
89#[derive(Debug)]
91pub struct LoadedData {
92 pub features: Vec<AnalyzableFeature>,
93 pub bounds: BoundingBox,
94 pub time_range: TimeRange,
95}
96
97pub fn load_data(source: &DataSource) -> Result<LoadedData> {
99 match source {
100 DataSource::GeoParquet { path, time_field, time_format } => {
101 load_geoparquet(path, time_field, time_format)
102 }
103 DataSource::SttArchive { path } => {
104 load_stt_archive(path)
105 }
106 }
107}
108
109fn load_geoparquet(path: &Path, time_field: &str, time_format: &str) -> Result<LoadedData> {
111 use indicatif::{ProgressBar, ProgressStyle};
112
113 let pb = ProgressBar::new_spinner();
114 pb.set_style(
115 ProgressStyle::default_spinner()
116 .template("{spinner:.green} {msg}")
117 .unwrap(),
118 );
119 pb.set_message("Loading GeoParquet file...");
120
121 let file = std::fs::File::open(path).context("Failed to open GeoParquet file")?;
122 let builder = ParquetRecordBatchReaderBuilder::try_new(file)?;
123 let schema = builder.schema().clone();
124
125 let geom_col_name = find_geometry_column(&schema)?;
127 let time_col_idx = schema.fields().iter().position(|f| f.name() == time_field)
128 .ok_or_else(|| anyhow::anyhow!("Time field '{}' not found", time_field))?;
129
130 let reader = builder.build()?;
131
132 let mut features = Vec::new();
133 let mut min_lon = f64::MAX;
134 let mut max_lon = f64::MIN;
135 let mut min_lat = f64::MAX;
136 let mut max_lat = f64::MIN;
137 let mut min_time = u64::MAX;
138 let mut max_time = u64::MIN;
139
140 for batch_result in reader {
141 let batch = batch_result.context("Failed to read Parquet batch")?;
142
143 let geometries = extract_geometries_from_batch(&batch, &geom_col_name)?;
144 let timestamps = extract_timestamps_from_batch(&batch, time_col_idx, time_format)?;
145
146 let property_count = schema.fields().len() - 2; for i in 0..batch.num_rows() {
150 let (geom_type, vertex_count, lon, lat) = geometries.get(i)
151 .cloned()
152 .unwrap_or((GeometryType::Unknown, 0, 0.0, 0.0));
153 let timestamp = timestamps.get(i).copied().unwrap_or(0);
154
155 min_lon = min_lon.min(lon);
157 max_lon = max_lon.max(lon);
158 min_lat = min_lat.min(lat);
159 max_lat = max_lat.max(lat);
160 min_time = min_time.min(timestamp);
161 max_time = max_time.max(timestamp);
162
163 let estimated_size = 100 + (vertex_count * 16) + (property_count * 20);
165
166 features.push(AnalyzableFeature {
167 lon,
168 lat,
169 timestamp,
170 geometry_type: geom_type,
171 vertex_count,
172 estimated_size,
173 property_count,
174 });
175 }
176
177 if features.len() % 100_000 == 0 {
178 pb.set_message(format!("Loaded {} features...", features.len()));
179 }
180 }
181
182 pb.finish_with_message(format!("Loaded {} features", features.len()));
183
184 Ok(LoadedData {
185 features,
186 bounds: BoundingBox::new(min_lon, min_lat, max_lon, max_lat),
187 time_range: TimeRange::new(min_time, max_time),
188 })
189}
190
191fn load_stt_archive(path: &Path) -> Result<LoadedData> {
193 use indicatif::{ProgressBar, ProgressStyle};
194 use stt_core::ArchiveReader;
195
196 let pb = ProgressBar::new_spinner();
197 pb.set_style(
198 ProgressStyle::default_spinner()
199 .template("{spinner:.green} {msg}")
200 .unwrap(),
201 );
202 pb.set_message("Loading STT archive...");
203
204 let reader = ArchiveReader::open(path)?;
205 let metadata = reader.metadata();
206 let entries = reader.entries();
207
208 let mut features = Vec::new();
209
210 for entry in entries {
212 anyhow::ensure!(
213 entry.zoom < 32,
214 "corrupt archive: tile entry {}/{}/{} has zoom {} (must be < 32)",
215 entry.zoom,
216 entry.x,
217 entry.y,
218 entry.zoom,
219 );
220 let center = tile_coords_to_lonlat(1, 1, entry.zoom, entry.x, entry.y, 2);
226 let (lon, lat) = (center.x(), center.y());
227
228 features.push(AnalyzableFeature {
231 lon,
232 lat,
233 timestamp: entry.time_start.max(0) as u64,
234 geometry_type: GeometryType::Point, vertex_count: 1,
236 estimated_size: entry.length as usize,
237 property_count: 0,
238 });
239 }
240
241 pb.finish_with_message(format!("Loaded {} tile entries", features.len()));
242
243 let b = metadata.bounds;
245 let bounds = BoundingBox::new(b.min_lon, b.min_lat, b.max_lon, b.max_lat);
246
247 let tr = metadata.time_range;
249 let time_range = TimeRange::new(tr.start, tr.end);
250
251 Ok(LoadedData {
252 features,
253 bounds,
254 time_range,
255 })
256}
257
258fn find_geometry_column(schema: &arrow::datatypes::Schema) -> Result<String> {
264 let common_names = ["geometry", "geom", "wkb_geometry", "the_geom", "shape"];
265
266 for name in common_names {
267 if schema.field_with_name(name).is_ok() {
268 return Ok(name.to_string());
269 }
270 }
271
272 for field in schema.fields() {
274 if matches!(field.data_type(), DataType::Binary | DataType::LargeBinary) {
275 return Ok(field.name().clone());
276 }
277 }
278
279 for field in schema.fields() {
281 if matches!(field.data_type(), DataType::Struct(_)) {
282 return Ok(field.name().clone());
283 }
284 }
285
286 let has_lon = schema.field_with_name("lon").is_ok()
288 || schema.field_with_name("longitude").is_ok()
289 || schema.field_with_name("x").is_ok();
290 let has_lat = schema.field_with_name("lat").is_ok()
291 || schema.field_with_name("latitude").is_ok()
292 || schema.field_with_name("y").is_ok();
293
294 if has_lon && has_lat {
295 return Ok("__lon_lat__".to_string());
296 }
297
298 anyhow::bail!("Could not find geometry column in Parquet schema")
299}
300
301fn extract_geometries_from_batch(
303 batch: &arrow::record_batch::RecordBatch,
304 geom_col_name: &str,
305) -> Result<Vec<(GeometryType, usize, f64, f64)>> {
306 let mut results = Vec::with_capacity(batch.num_rows());
307
308 if geom_col_name == "__lon_lat__" {
310 let lon_col = batch.column_by_name("lon")
311 .or_else(|| batch.column_by_name("longitude"))
312 .or_else(|| batch.column_by_name("x"));
313 let lat_col = batch.column_by_name("lat")
314 .or_else(|| batch.column_by_name("latitude"))
315 .or_else(|| batch.column_by_name("y"));
316
317 if let (Some(lon), Some(lat)) = (lon_col, lat_col) {
318 if let (Some(lon_arr), Some(lat_arr)) = (
319 lon.as_any().downcast_ref::<Float64Array>(),
320 lat.as_any().downcast_ref::<Float64Array>(),
321 ) {
322 for i in 0..batch.num_rows() {
323 if lon_arr.is_valid(i) && lat_arr.is_valid(i) {
324 results.push((GeometryType::Point, 1, lon_arr.value(i), lat_arr.value(i)));
325 } else {
326 results.push((GeometryType::Unknown, 0, 0.0, 0.0));
327 }
328 }
329 return Ok(results);
330 }
331 }
332 anyhow::bail!("Expected lon/lat columns but could not read them");
333 }
334
335 let geom_col = batch.column_by_name(geom_col_name)
336 .ok_or_else(|| anyhow::anyhow!("Geometry column '{}' not found", geom_col_name))?;
337
338 if let Some(struct_array) = geom_col.as_any().downcast_ref::<arrow::array::StructArray>() {
340 let x_col = struct_array.column_by_name("x")
341 .or_else(|| struct_array.column_by_name("longitude"))
342 .or_else(|| struct_array.column_by_name("lon"));
343 let y_col = struct_array.column_by_name("y")
344 .or_else(|| struct_array.column_by_name("latitude"))
345 .or_else(|| struct_array.column_by_name("lat"));
346
347 if let (Some(x), Some(y)) = (x_col, y_col) {
348 if let (Some(x_arr), Some(y_arr)) = (
349 x.as_any().downcast_ref::<Float64Array>(),
350 y.as_any().downcast_ref::<Float64Array>(),
351 ) {
352 for i in 0..batch.num_rows() {
353 if x_arr.is_valid(i) && y_arr.is_valid(i) {
354 results.push((GeometryType::Point, 1, x_arr.value(i), y_arr.value(i)));
355 } else {
356 results.push((GeometryType::Unknown, 0, 0.0, 0.0));
357 }
358 }
359 return Ok(results);
360 }
361 }
362 }
363
364 if let Some(binary_array) = geom_col.as_any().downcast_ref::<arrow::array::BinaryArray>() {
366 for i in 0..batch.num_rows() {
367 if binary_array.is_valid(i) {
368 let wkb = binary_array.value(i);
369 if let Some((geom_type, vertex_count, lon, lat)) = parse_wkb_info(wkb) {
370 results.push((geom_type, vertex_count, lon, lat));
371 } else {
372 results.push((GeometryType::Unknown, 0, 0.0, 0.0));
373 }
374 } else {
375 results.push((GeometryType::Unknown, 0, 0.0, 0.0));
376 }
377 }
378 return Ok(results);
379 }
380
381 let lon_col = batch.column_by_name("lon")
383 .or_else(|| batch.column_by_name("longitude"))
384 .or_else(|| batch.column_by_name("x"));
385 let lat_col = batch.column_by_name("lat")
386 .or_else(|| batch.column_by_name("latitude"))
387 .or_else(|| batch.column_by_name("y"));
388
389 if let (Some(lon), Some(lat)) = (lon_col, lat_col) {
390 if let (Some(lon_arr), Some(lat_arr)) = (
391 lon.as_any().downcast_ref::<Float64Array>(),
392 lat.as_any().downcast_ref::<Float64Array>(),
393 ) {
394 for i in 0..batch.num_rows() {
395 if lon_arr.is_valid(i) && lat_arr.is_valid(i) {
396 results.push((GeometryType::Point, 1, lon_arr.value(i), lat_arr.value(i)));
397 } else {
398 results.push((GeometryType::Unknown, 0, 0.0, 0.0));
399 }
400 }
401 return Ok(results);
402 }
403 }
404
405 anyhow::bail!("Could not extract geometries from column '{}'", geom_col_name)
406}
407
408fn extract_timestamps_from_batch(
410 batch: &arrow::record_batch::RecordBatch,
411 col_idx: usize,
412 time_format: &str,
413) -> Result<Vec<u64>> {
414 let column = batch.column(col_idx);
415 let mut timestamps = Vec::with_capacity(batch.num_rows());
416
417 macro_rules! push_ts_column {
424 ($arr:expr, $unit:expr) => {{
425 for i in 0..batch.num_rows() {
426 if $arr.is_valid(i) {
427 timestamps.push(normalize_timestamp_to_ms(i, $arr.value(i), $unit)?);
428 } else {
429 timestamps.push(0);
430 }
431 }
432 return Ok(timestamps);
433 }};
434 }
435
436 if let Some(ts_array) = column.as_any().downcast_ref::<TimestampSecondArray>() {
437 push_ts_column!(ts_array, TimestampUnit::Second);
438 }
439 if let Some(ts_array) = column.as_any().downcast_ref::<TimestampMillisecondArray>() {
440 push_ts_column!(ts_array, TimestampUnit::Millisecond);
441 }
442 if let Some(ts_array) = column.as_any().downcast_ref::<TimestampMicrosecondArray>() {
443 push_ts_column!(ts_array, TimestampUnit::Microsecond);
444 }
445 if let Some(ts_array) = column.as_any().downcast_ref::<TimestampNanosecondArray>() {
446 push_ts_column!(ts_array, TimestampUnit::Nanosecond);
447 }
448
449 if let Some(int_array) = column.as_any().downcast_ref::<Int64Array>() {
451 let unit = match time_format {
452 "unix-sec" => TimestampUnit::Second,
453 _ => TimestampUnit::Millisecond,
454 };
455 for i in 0..batch.num_rows() {
456 if int_array.is_valid(i) {
457 timestamps.push(normalize_timestamp_to_ms(i, int_array.value(i), unit)?);
458 } else {
459 timestamps.push(0);
460 }
461 }
462 return Ok(timestamps);
463 }
464
465 if let Some(str_array) = column.as_any().downcast_ref::<StringArray>() {
467 for i in 0..batch.num_rows() {
468 if str_array.is_valid(i) {
469 let s = str_array.value(i);
470 let ts = parse_iso8601(s).unwrap_or(0);
471 timestamps.push(ts);
472 } else {
473 timestamps.push(0);
474 }
475 }
476 return Ok(timestamps);
477 }
478
479 anyhow::bail!("Unsupported timestamp column type")
480}
481
482fn parse_iso8601(s: &str) -> Result<u64> {
484 use chrono::{DateTime, NaiveDateTime};
485
486 if let Ok(dt) = s.parse::<DateTime<chrono::Utc>>() {
487 return Ok(dt.timestamp_millis() as u64);
488 }
489
490 if let Ok(dt) = NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S") {
491 return Ok(dt.and_utc().timestamp_millis() as u64);
492 }
493
494 if let Ok(date) = chrono::NaiveDate::parse_from_str(s, "%Y-%m-%d") {
495 let dt = date.and_hms_opt(0, 0, 0).unwrap().and_utc();
496 return Ok(dt.timestamp_millis() as u64);
497 }
498
499 anyhow::bail!("Failed to parse timestamp: {}", s)
500}
501
502const WKB_POINT: u32 = 1;
507const WKB_LINESTRING: u32 = 2;
508const WKB_POLYGON: u32 = 3;
509const WKB_MULTIPOINT: u32 = 4;
510const WKB_MULTILINESTRING: u32 = 5;
511const WKB_MULTIPOLYGON: u32 = 6;
512
513fn parse_wkb_info(wkb: &[u8]) -> Option<(GeometryType, usize, f64, f64)> {
514 if wkb.len() < 5 {
515 return None;
516 }
517
518 let little_endian = wkb[0] == 1;
519 let geom_type = if little_endian {
520 u32::from_le_bytes([wkb[1], wkb[2], wkb[3], wkb[4]]) % 1000
521 } else {
522 u32::from_be_bytes([wkb[1], wkb[2], wkb[3], wkb[4]]) % 1000
523 };
524
525 match geom_type {
526 WKB_POINT => parse_wkb_point_info(wkb, little_endian),
527 WKB_LINESTRING => parse_wkb_linestring_info(wkb, little_endian),
528 WKB_POLYGON => parse_wkb_polygon_info(wkb, little_endian),
529 WKB_MULTIPOINT => parse_wkb_multipoint_info(wkb, little_endian),
530 WKB_MULTILINESTRING => parse_wkb_multilinestring_info(wkb, little_endian),
531 WKB_MULTIPOLYGON => parse_wkb_multipolygon_info(wkb, little_endian),
532 _ => None,
533 }
534}
535
536fn read_f64(wkb: &[u8], offset: usize, little_endian: bool) -> Option<f64> {
537 if offset + 8 > wkb.len() {
538 return None;
539 }
540 let bytes: [u8; 8] = wkb[offset..offset+8].try_into().ok()?;
541 Some(if little_endian {
542 f64::from_le_bytes(bytes)
543 } else {
544 f64::from_be_bytes(bytes)
545 })
546}
547
548fn read_u32(wkb: &[u8], offset: usize, little_endian: bool) -> Option<u32> {
549 if offset + 4 > wkb.len() {
550 return None;
551 }
552 let bytes: [u8; 4] = wkb[offset..offset+4].try_into().ok()?;
553 Some(if little_endian {
554 u32::from_le_bytes(bytes)
555 } else {
556 u32::from_be_bytes(bytes)
557 })
558}
559
560fn parse_wkb_point_info(wkb: &[u8], little_endian: bool) -> Option<(GeometryType, usize, f64, f64)> {
561 let x = read_f64(wkb, 5, little_endian)?;
562 let y = read_f64(wkb, 13, little_endian)?;
563 Some((GeometryType::Point, 1, x, y))
564}
565
566fn parse_wkb_linestring_info(wkb: &[u8], little_endian: bool) -> Option<(GeometryType, usize, f64, f64)> {
567 let num_points = read_u32(wkb, 5, little_endian)? as usize;
568 if num_points == 0 {
569 return None;
570 }
571
572 let mut sum_x = 0.0;
573 let mut sum_y = 0.0;
574 let mut offset = 9;
575
576 for _ in 0..num_points {
577 let x = read_f64(wkb, offset, little_endian)?;
578 let y = read_f64(wkb, offset + 8, little_endian)?;
579 sum_x += x;
580 sum_y += y;
581 offset += 16;
582 }
583
584 Some((GeometryType::LineString, num_points, sum_x / num_points as f64, sum_y / num_points as f64))
585}
586
587fn parse_wkb_polygon_info(wkb: &[u8], little_endian: bool) -> Option<(GeometryType, usize, f64, f64)> {
588 let num_rings = read_u32(wkb, 5, little_endian)? as usize;
589 if num_rings == 0 {
590 return None;
591 }
592
593 let mut total_points = 0usize;
594 let mut sum_x = 0.0;
595 let mut sum_y = 0.0;
596 let mut offset = 9;
597
598 for ring_idx in 0..num_rings {
599 let num_points = read_u32(wkb, offset, little_endian)? as usize;
600 offset += 4;
601
602 for _ in 0..num_points {
603 let x = read_f64(wkb, offset, little_endian)?;
604 let y = read_f64(wkb, offset + 8, little_endian)?;
605
606 if ring_idx == 0 {
607 sum_x += x;
608 sum_y += y;
609 total_points += 1;
610 }
611
612 offset += 16;
613 }
614 }
615
616 let vertex_count: usize = total_points;
617 let centroid_x = if total_points > 0 { sum_x / total_points as f64 } else { 0.0 };
618 let centroid_y = if total_points > 0 { sum_y / total_points as f64 } else { 0.0 };
619
620 Some((GeometryType::Polygon, vertex_count, centroid_x, centroid_y))
621}
622
623fn parse_wkb_multipoint_info(wkb: &[u8], little_endian: bool) -> Option<(GeometryType, usize, f64, f64)> {
624 let num_geoms = read_u32(wkb, 5, little_endian)? as usize;
625 if num_geoms == 0 {
626 return None;
627 }
628
629 let mut sum_x = 0.0;
630 let mut sum_y = 0.0;
631 let mut offset = 9;
632
633 for _ in 0..num_geoms {
634 offset += 5; let x = read_f64(wkb, offset, little_endian)?;
636 let y = read_f64(wkb, offset + 8, little_endian)?;
637 sum_x += x;
638 sum_y += y;
639 offset += 16;
640 }
641
642 Some((GeometryType::MultiPoint, num_geoms, sum_x / num_geoms as f64, sum_y / num_geoms as f64))
643}
644
645fn parse_wkb_multilinestring_info(wkb: &[u8], little_endian: bool) -> Option<(GeometryType, usize, f64, f64)> {
646 let num_geoms = read_u32(wkb, 5, little_endian)? as usize;
647 if num_geoms == 0 {
648 return None;
649 }
650
651 let mut total_points = 0usize;
652 let mut sum_x = 0.0;
653 let mut sum_y = 0.0;
654 let mut offset = 9;
655
656 for _ in 0..num_geoms {
657 offset += 5; let num_points = read_u32(wkb, offset, little_endian)? as usize;
659 offset += 4;
660
661 for _ in 0..num_points {
662 let x = read_f64(wkb, offset, little_endian)?;
663 let y = read_f64(wkb, offset + 8, little_endian)?;
664 sum_x += x;
665 sum_y += y;
666 total_points += 1;
667 offset += 16;
668 }
669 }
670
671 let centroid_x = if total_points > 0 { sum_x / total_points as f64 } else { 0.0 };
672 let centroid_y = if total_points > 0 { sum_y / total_points as f64 } else { 0.0 };
673
674 Some((GeometryType::MultiLineString, total_points, centroid_x, centroid_y))
675}
676
677fn parse_wkb_multipolygon_info(wkb: &[u8], little_endian: bool) -> Option<(GeometryType, usize, f64, f64)> {
678 let num_geoms = read_u32(wkb, 5, little_endian)? as usize;
679 if num_geoms == 0 {
680 return None;
681 }
682
683 let mut total_points = 0usize;
684 let mut sum_x = 0.0;
685 let mut sum_y = 0.0;
686 let mut offset = 9;
687
688 for poly_idx in 0..num_geoms {
689 offset += 5; let num_rings = read_u32(wkb, offset, little_endian)? as usize;
691 offset += 4;
692
693 for ring_idx in 0..num_rings {
694 let num_points = read_u32(wkb, offset, little_endian)? as usize;
695 offset += 4;
696
697 for _ in 0..num_points {
698 let x = read_f64(wkb, offset, little_endian)?;
699 let y = read_f64(wkb, offset + 8, little_endian)?;
700
701 if poly_idx == 0 && ring_idx == 0 {
702 sum_x += x;
703 sum_y += y;
704 total_points += 1;
705 }
706
707 offset += 16;
708 }
709 }
710 }
711
712 let centroid_x = if total_points > 0 { sum_x / total_points as f64 } else { 0.0 };
713 let centroid_y = if total_points > 0 { sum_y / total_points as f64 } else { 0.0 };
714
715 Some((GeometryType::MultiPolygon, total_points, centroid_x, centroid_y))
716}
717