fix: don't consume the geoparquet writer when finishing (#842)

Author: gadomski

crates/core/src/error.rs

@@ -5,21 +5,6 @@ use thiserror::Error;
 #[derive(Error, Debug)]
 #[non_exhaustive]
 pub enum Error {
-    /// [arrow_schema::ArrowError]
-    #[error(transparent)]
-    #[cfg(feature = "geoarrow")]
-    Arrow(#[from] arrow_schema::ArrowError),
-
-    /// The schema of two sets of items don't match.
-    #[cfg(feature = "geoarrow")]
-    #[error("Arrow schema mismatch")]
-    ArrowSchemaMismatch,
-
-    /// The arrow table is empty
-    #[cfg(feature = "geoarrow")]
-    #[error("Empty arrow table")]
-    EmptyArrowTable,
-
     /// [chrono::ParseError]
     #[error(transparent)]
     ChronoParse(#[from] chrono::ParseError),
@@ -28,19 +13,10 @@ pub enum Error {
     #[error("{0} is not enabled")]
     FeatureNotEnabled(&'static str),
 
-    /// [geoarrow_schema::error::GeoArrowError]
-    #[error(transparent)]
-    #[cfg(feature = "geoarrow")]
-    GeoArrow(#[from] geoarrow_schema::error::GeoArrowError),
-
     /// [geojson::Error]
     #[error(transparent)]
     Geojson(#[from] Box<geojson::Error>),
 
-    /// [std::io::Error]
-    #[error(transparent)]
-    Io(#[from] std::io::Error),
-
     /// Returned when a STAC object has the wrong type field.
     #[error("incorrect type: expected={expected}, actual={actual}")]
     IncorrectType {
@@ -64,32 +40,26 @@ pub enum Error {
     #[error("invalid datetime: {0}")]
     InvalidDatetime(String),
 
+    /// [std::io::Error]
+    #[error(transparent)]
+    Io(#[from] std::io::Error),
+
     /// Returned when there is not a required field on a STAC object
     #[error("no \"{0}\" field in the JSON object")]
     MissingField(&'static str),
 
-    /// No geoparquet metadata in a stac-geoparquet file.
-    #[error("no geoparquet metadata")]
-    #[cfg(feature = "geoparquet")]
-    MissingGeoparquetMetadata,
+    /// There is not an href, when an href is required.
+    #[error("no href")]
+    NoHref,
 
     /// There are no items, when items are required.
     #[error("no items")]
     NoItems,
 
-    /// There is not an href, when an href is required.
-    #[error("no href")]
-    NoHref,
-
     /// This is not a JSON object.
     #[error("json value is not an object")]
     NotAnObject(serde_json::Value),
 
-    /// [parquet::errors::ParquetError]
-    #[error(transparent)]
-    #[cfg(feature = "geoparquet")]
-    Parquet(#[from] parquet::errors::ParquetError),
-
     /// [serde_json::Error]
     #[error(transparent)]
     SerdeJson(#[from] serde_json::Error),
@@ -114,8 +84,43 @@ pub enum Error {
     #[error(transparent)]
     UrlParse(#[from] url::ParseError),
 
+    /// [arrow_schema::ArrowError]
+    #[error(transparent)]
+    #[cfg(feature = "geoarrow")]
+    Arrow(#[from] arrow_schema::ArrowError),
+
+    /// The schema of two sets of items don't match.
+    #[cfg(feature = "geoarrow")]
+    #[error("Arrow schema mismatch")]
+    ArrowSchemaMismatch,
+
+    /// The arrow table is empty
+    #[cfg(feature = "geoarrow")]
+    #[error("Empty arrow table")]
+    EmptyArrowTable,
+
+    /// [geoarrow_schema::error::GeoArrowError]
+    #[error(transparent)]
+    #[cfg(feature = "geoarrow")]
+    GeoArrow(#[from] geoarrow_schema::error::GeoArrowError),
+
     /// [wkb::error::WkbError]
     #[error(transparent)]
     #[cfg(feature = "geoarrow")]
     Wkb(#[from] wkb::error::WkbError),
+
+    /// The geoparquet writer has been closed.
+    #[error("The geoparquet writer has already been closed")]
+    #[cfg(feature = "geoparquet")]
+    ClosedGeoparquetWriter,
+
+    /// No geoparquet metadata in a stac-geoparquet file.
+    #[error("no geoparquet metadata")]
+    #[cfg(feature = "geoparquet")]
+    MissingGeoparquetMetadata,
+
+    /// [parquet::errors::ParquetError]
+    #[error(transparent)]
+    #[cfg(feature = "geoparquet")]
+    Parquet(#[from] parquet::errors::ParquetError),
 }

crates/core/src/geoparquet.rs

@@ -100,7 +100,7 @@ where
     WriterBuilder::new(writer)
         .compression(compression)
         .build(item_collection.into().items)
-        .and_then(|writer| writer.finish())
+        .and_then(|mut writer| writer.finish())
 }
 
 /// Builder for a stac-geoparquet writer.
@@ -115,7 +115,9 @@ pub struct WriterBuilder<W: Write + Send> {
 #[allow(missing_debug_implementations)]
 pub struct Writer<W: Write + Send> {
     geoarrow_encoder: Encoder,
-    encoder: GeoParquetRecordBatchEncoder,
+    // We make this an option so we can consume it during write but keep write
+    // as only requiring a mutable reference.
+    encoder: Option<GeoParquetRecordBatchEncoder>,
     arrow_writer: ArrowWriter<W>,
 }
 
@@ -191,7 +193,7 @@ impl<W: Write + Send> WriterBuilder<W> {
     ///
     /// let item: Item = stac::read("examples/simple-item.json").unwrap();
     /// let cursor = Cursor::new(Vec::new());
-    /// let writer = WriterBuilder::new(cursor).build(vec![item]).unwrap();
+    /// let mut writer = WriterBuilder::new(cursor).build(vec![item]).unwrap();
     /// writer.finish().unwrap();
     /// ```
     pub fn build(self, items: Vec<Item>) -> Result<Writer<W>> {
@@ -222,13 +224,15 @@ impl<W: Write + Send> Writer<W> {
         arrow_writer.write(&record_batch)?;
         Ok(Writer {
             geoarrow_encoder,
-            encoder,
+            encoder: Some(encoder),
             arrow_writer,
         })
     }
 
     /// Writes more items to this writer.
     ///
+    /// It's an error to write after `finish` has been called.
+    ///
     /// # Examples
     ///
     /// ```
@@ -244,13 +248,19 @@ impl<W: Write + Send> Writer<W> {
     /// ```
     pub fn write(&mut self, items: Vec<Item>) -> Result<()> {
         let record_batch = self.geoarrow_encoder.encode(items)?;
-        let record_batch = self.encoder.encode_record_batch(&record_batch)?;
+        let record_batch = if let Some(encoder) = self.encoder.as_mut() {
+            encoder.encode_record_batch(&record_batch)?
+        } else {
+            return Err(Error::ClosedGeoparquetWriter);
+        };
         self.arrow_writer.write(&record_batch)?;
         Ok(())
     }
 
     /// Finishes writing.
     ///
+    /// It's an error to call finish twice.
+    ///
     /// # Examples
     ///
     /// ```
@@ -259,12 +269,16 @@ impl<W: Write + Send> Writer<W> {
     ///
     /// let item: Item = stac::read("examples/simple-item.json").unwrap();
     /// let cursor = Cursor::new(Vec::new());
-    /// let writer = WriterBuilder::new(cursor).build(vec![item]).unwrap();
+    /// let mut writer = WriterBuilder::new(cursor).build(vec![item]).unwrap();
     /// writer.finish().unwrap();
     /// ```
-    pub fn finish(mut self) -> Result<()> {
-        self.arrow_writer
-            .append_key_value_metadata(self.encoder.into_keyvalue()?);
+    pub fn finish(&mut self) -> Result<()> {
+        if let Some(encoder) = self.encoder.take() {
+            self.arrow_writer
+                .append_key_value_metadata(encoder.into_keyvalue()?);
+        } else {
+            return Err(Error::ClosedGeoparquetWriter);
+        }
         self.arrow_writer.append_key_value_metadata(KeyValue::new(
             VERSION_KEY.to_string(),
             Some(VERSION.to_string()),