You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@sedona.apache.org by pr...@apache.org on 2026/05/31 07:08:45 UTC
(sedona-spatialbench) 02/06: feat: add CLI args and sequential raster runner
This is an automated email from the ASF dual-hosted git repository.
prantogg pushed a commit to branch pranav/feature/raster-benchmark
in repository https://gitbox.apache.org/repos/asf/sedona-spatialbench.git
commit fc0e6ab74e10f009f322e73dfcff213983ded5e6
Author: Pranav Toggi <pr...@gmail.com>
AuthorDate: Sat May 30 23:48:54 2026 -0700
feat: add CLI args and sequential raster runner
Wire --raster-format cog and --max-footprints CLI args into
spatialbench-cli. Add sequential raster_runner that generates COGs
for all footprints × scenes and collects ManifestEntry metadata.
---
spatialbench-cli/Cargo.toml | 1 +
spatialbench-cli/src/main.rs | 52 ++++++++++++++++++++
spatialbench-cli/src/raster_runner.rs | 93 +++++++++++++++++++++++++++++++++++
3 files changed, 146 insertions(+)
diff --git a/spatialbench-cli/Cargo.toml b/spatialbench-cli/Cargo.toml
index cebd31e..ae7ffd1 100644
--- a/spatialbench-cli/Cargo.toml
+++ b/spatialbench-cli/Cargo.toml
@@ -34,6 +34,7 @@ parquet = "56"
clap = { version = "4.5.32", features = ["derive"] }
spatialbench = { path = "../spatialbench", version = "0.3.0"}
spatialbench-arrow = { path = "../spatialbench-arrow", version = "0.3.0" }
+spatialbench-raster = { path = "../spatialbench-raster", version = "0.3.0", features = ["cog-writer"] }
tokio = { version = "1.44.1", features = ["full"]}
futures = "0.3.31"
num_cpus = "1.0"
diff --git a/spatialbench-cli/src/main.rs b/spatialbench-cli/src/main.rs
index 431720b..e32049a 100644
--- a/spatialbench-cli/src/main.rs
+++ b/spatialbench-cli/src/main.rs
@@ -27,6 +27,7 @@ mod generate;
mod output_plan;
mod parquet;
mod plan;
+mod raster_runner;
mod runner;
mod s3_writer;
mod spatial_config_file;
@@ -47,6 +48,9 @@ use log::{debug, info, LevelFilter};
use spatialbench::distribution::Distributions;
use spatialbench::spatial::overrides::{set_overrides, SpatialOverrides};
use spatialbench::text::TextPool;
+use spatialbench_raster::cog::CogConfig;
+use spatialbench_raster::footprint::FootprintGrid;
+use spatialbench_raster::scaling::scaling_tier;
use std::fmt::Display;
use std::fs::{self, File};
use std::io::{self, BufWriter, Stdout, Write};
@@ -137,6 +141,14 @@ struct Cli {
/// Typical values range from 10MB to 100MB.
#[arg(long, default_value_t = DEFAULT_PARQUET_ROW_GROUP_BYTES)]
parquet_row_group_bytes: i64,
+
+ /// Raster output format (e.g., cog). When set, generates raster COGs.
+ #[arg(long, value_enum)]
+ raster_format: Option<RasterFormat>,
+
+ /// Maximum number of footprints to generate (dev flag for fast iteration).
+ #[arg(long)]
+ max_footprints: Option<u32>,
}
#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
@@ -234,6 +246,11 @@ enum OutputFormat {
Parquet,
}
+#[derive(Debug, Copy, Clone, PartialEq, Eq, ValueEnum)]
+enum RasterFormat {
+ Cog,
+}
+
#[tokio::main]
async fn main() -> io::Result<()> {
// Parse command line arguments
@@ -363,6 +380,41 @@ impl Cli {
// Run
let runner = runner::PlanRunner::new(output_plans, self.num_threads);
runner.run().await?;
+
+ // Raster generation
+ if let Some(RasterFormat::Cog) = self.raster_format {
+ let sf = self.scale_factor as u32;
+ let tier = scaling_tier(sf).map_err(|e| {
+ io::Error::new(io::ErrorKind::InvalidInput, format!("raster scaling: {e}"))
+ })?;
+ let cog_config = CogConfig::default();
+ let grid = FootprintGrid::new(
+ spatialbench::spatial::ContinentAffines::default(),
+ cog_config.raster,
+ self.max_footprints,
+ );
+ let footprints = grid.generate(tier);
+
+ info!(
+ "{} footprints, {} COGs/footprint, {} total COGs",
+ footprints.len(),
+ tier.scenes_per_footprint,
+ footprints.len() as u64 * tier.scenes_per_footprint as u64
+ );
+
+ let raster_dir = self.output_dir.join("raster");
+ let manifest = raster_runner::run_raster(
+ &footprints,
+ tier,
+ &cog_config,
+ &raster_dir,
+ self.num_threads,
+ )
+ .await?;
+
+ info!("generated {} manifest entries", manifest.len());
+ }
+
info!("Generation complete!");
Ok(())
}
diff --git a/spatialbench-cli/src/raster_runner.rs b/spatialbench-cli/src/raster_runner.rs
new file mode 100644
index 0000000..0ca0f94
--- /dev/null
+++ b/spatialbench-cli/src/raster_runner.rs
@@ -0,0 +1,93 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Sequential raster COG generation runner.
+//!
+//! Generates COGs for all footprints × scenes, collecting [`ManifestEntry`]
+//! metadata for downstream STAC catalog generation.
+
+use spatialbench_raster::cog::{write_cog, CogConfig};
+use spatialbench_raster::footprint::Footprint;
+use spatialbench_raster::scaling::ScalingTier;
+
+use log::info;
+
+use std::io;
+use std::path::Path;
+
+/// Metadata for a single generated COG, collected for STAC catalog writing.
+///
+/// Does not store a `PathBuf` — the path is deterministic from
+/// `(footprint_id, cog_id)` as `pile/{fp:05}/{cog:04}.tif` and
+/// reconstructed when needed, avoiding heap allocations.
+#[derive(Debug, Clone, Copy)]
+#[allow(dead_code)] // Fields used in commit 4 (STAC writer)
+pub struct ManifestEntry {
+ /// Footprint ID (matches [`Footprint::id`]).
+ pub footprint_id: u32,
+ /// COG scene ID within this footprint.
+ pub cog_id: u32,
+ /// Bounding box in EPSG:4326: [west, south, east, north].
+ pub bbox_4326: [f64; 4],
+ /// EPSG code for the COG's CRS.
+ pub epsg: u32,
+}
+
+/// Generate COGs for all footprints and return manifest entries.
+///
+/// This is the sequential version — commit 2 upgrades to parallel.
+pub async fn run_raster(
+ footprints: &[Footprint],
+ tier: &ScalingTier,
+ cog_config: &CogConfig,
+ output_dir: &Path,
+ _num_threads: usize,
+) -> io::Result<Vec<ManifestEntry>> {
+ let pile_dir = output_dir.join("pile");
+ let total_cogs = footprints.len() as u64 * tier.scenes_per_footprint as u64;
+ let mut manifest = Vec::with_capacity(total_cogs as usize);
+
+ // Pre-create all footprint directories
+ for fp in footprints {
+ let fp_dir = pile_dir.join(format!("{:05}", fp.id));
+ std::fs::create_dir_all(&fp_dir)?;
+ }
+
+ let mut count = 0u64;
+ for fp in footprints {
+ let fp_dir = pile_dir.join(format!("{:05}", fp.id));
+ for cog_id in 0..tier.scenes_per_footprint {
+ let path = fp_dir.join(format!("{:04}.tif", cog_id));
+ write_cog(cog_config, fp, cog_id, &path)?;
+
+ manifest.push(ManifestEntry {
+ footprint_id: fp.id,
+ cog_id,
+ bbox_4326: fp.bbox_4326,
+ epsg: fp.epsg,
+ });
+
+ count += 1;
+ if count.is_multiple_of(100) {
+ info!("generated {count}/{total_cogs} COGs");
+ }
+ }
+ }
+
+ info!("generated {count}/{total_cogs} COGs (complete)");
+ Ok(manifest)
+}