feat(cli,core): allow compression (and other options)

Author: gadomski

cli/Cargo.toml

@@ -14,14 +14,21 @@ categories = ["science", "data-structures"]
 default = ["gdal", "geoparquet", "pgstac"]
 duckdb = ["dep:stac-duckdb"]
 gdal = ["stac/gdal"]
-geoparquet = ["dep:bytes", "stac/geoparquet-compression"]
+geoparquet = [
+    "dep:bytes",
+    "dep:geoarrow",
+    "stac/geoparquet-compression",
+    "parquet",
+]
 pgstac = ["stac-server/pgstac"]
 python = ["dep:pyo3", "pgstac", "duckdb", "stac-duckdb/bundled", "geoparquet"]
 
 [dependencies]
 axum = "0.7"
 bytes = { version = "1", optional = true }
 clap = { version = "4", features = ["derive"] }
+geoarrow = { git = "https://github.com/geoarrow/geoarrow-rs", rev = "6b877486bf98f280bc04b589eb7ce25b20e629f0", optional = true }
+parquet = { version = "52", optional = true }
 pyo3 = { version = "0.22", optional = true }
 reqwest = "0.12"
 serde = "1"

cli/src/args.rs

@@ -19,6 +19,13 @@ pub struct Args {
     #[arg(short, long, value_enum)]
     pub output_format: Option<Format>,
 
+    /// The type of geoparquet compression to use.
+    ///
+    /// Possible values: uncompressed (default), snappy, gzip(level), lzo, brotli(level), lz4, zstd(level), lz4_raw
+    #[arg(long, value_enum)]
+    #[cfg(feature = "geoparquet")]
+    pub geoparquet_compression: Option<parquet::basic::Compression>,
+
     /// The subcommand to run.
     #[command(subcommand)]
     pub subcommand: Subcommand,

cli/src/lib.rs

@@ -67,6 +67,8 @@ pub type Result<T> = std::result::Result<T, Error>;
 ///     compact: false,
 ///     input_format: None,
 ///     output_format: None,
+///     #[cfg(feature = "geoparquet")]
+///     geoparquet_compression: None,
 ///     subcommand: Subcommand::Sort(sort_args),
 /// };
 /// # tokio_test::block_on(async {
@@ -82,6 +84,8 @@ pub async fn run(args: Args) -> Result<()> {
         compact: args.compact,
         input_format,
         output_format,
+        #[cfg(feature = "geoparquet")]
+        geoparquet_compression: args.geoparquet_compression,
         writer,
         buffer: 100,
     };

cli/src/runner.rs

@@ -16,6 +16,10 @@ where
     /// The output format.
     pub output_format: Format,
 
+    /// Geoparquet compression.
+    #[cfg(feature = "geoparquet")]
+    pub geoparquet_compression: Option<parquet::basic::Compression>,
+
     /// The output writeable stream.
     pub writer: W,
 
@@ -49,7 +53,19 @@ where
                 #[cfg(feature = "geoparquet")]
                 Format::Parquet => {
                     if let Some(value) = value.to_stac() {
-                        stac::geoparquet::to_writer(&mut self.writer, value)?;
+                        let mut options = geoarrow::io::parquet::GeoParquetWriterOptions::default();
+                        if let Some(compression) = self.geoparquet_compression {
+                            let writer_properites =
+                                parquet::file::properties::WriterProperties::builder()
+                                    .set_compression(compression)
+                                    .build();
+                            options.writer_properties = Some(writer_properites);
+                        }
+                        stac::geoparquet::to_writer_with_options(
+                            &mut self.writer,
+                            value,
+                            &options,
+                        )?;
                     } else {
                         writeln!(self.writer, "{}", value)?;
                     }

core/src/geoparquet.rs

@@ -20,12 +20,12 @@ pub fn has_extension(href: &str) -> bool {
 }
 
 #[cfg(feature = "geoparquet")]
-pub use has_feature::{from_reader, to_writer};
+pub use has_feature::{from_reader, to_writer, to_writer_with_options};
 
 #[cfg(feature = "geoparquet")]
 mod has_feature {
     use crate::{Error, ItemCollection, Result, Value};
-    use geoarrow::io::parquet::GeoParquetRecordBatchReaderBuilder;
+    use geoarrow::io::parquet::{GeoParquetRecordBatchReaderBuilder, GeoParquetWriterOptions};
     use parquet::file::reader::ChunkReader;
     use std::io::Write;
 
@@ -46,6 +46,35 @@ mod has_feature {
     /// stac::geoparquet::to_writer(&mut cursor, item.into()).unwrap();
     /// ```
     pub fn to_writer<W>(writer: W, value: Value) -> Result<()>
+    where
+        W: Write + Send,
+    {
+        to_writer_with_options(writer, value, &Default::default())
+    }
+
+    /// Writes a [Value] to a [std::io::Write] as
+    /// [stac-geoparquet](https://github.com/stac-utils/stac-geoparquet) with the provided options.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use std::io::Cursor;
+    /// use stac::Item;
+    /// use geoarrow::io::parquet::GeoParquetWriterOptions;
+    /// use parquet::{basic::Compression, file::properties::WriterProperties};
+    ///
+    /// let item: Item = stac::read("examples/simple-item.json").unwrap();
+    /// let mut cursor = Cursor::new(Vec::new());
+    /// let mut options = GeoParquetWriterOptions::default();
+    /// let writer_properties = WriterProperties::builder().set_compression(Compression::SNAPPY).build();
+    /// options.writer_properties = Some(writer_properties);
+    /// stac::geoparquet::to_writer_with_options(&mut cursor, item.into(), &options).unwrap();
+    /// ```
+    pub fn to_writer_with_options<W>(
+        writer: W,
+        value: Value,
+        options: &GeoParquetWriterOptions,
+    ) -> Result<()>
     where
         W: Write + Send,
     {
@@ -55,7 +84,7 @@ mod has_feature {
                 geoarrow::io::parquet::write_geoparquet(
                     table.into_record_batch_reader(),
                     writer,
-                    &Default::default(),
+                    options,
                 )
                 .map_err(Error::from)
             }