Async GeoParquet writer (#587)

Rust only for now

Cargo.toml

@@ -24,7 +24,7 @@ geos = ["dep:geos"]
 geozero = ["dep:geozero"]
 ipc_compression = ["arrow-ipc/lz4", "arrow-ipc/zstd"]
 parquet = ["dep:parquet"]
-parquet_async = ["parquet", "parquet/async", "dep:futures"]
+parquet_async = ["parquet", "parquet/async", "dep:futures", "dep:tokio"]
 parquet_compression = [
   "parquet/snap",
   "parquet/brotli",
@@ -88,6 +88,7 @@ sqlx = { version = "0.7", optional = true, default-features = false, features =
   "tls-rustls",
 ] }
 thiserror = "1"
+tokio = { version = "1", default-features = false, optional = true }
 
 
 [dev-dependencies]

src/io/parquet/mod.rs

@@ -52,3 +52,5 @@ pub use reader::{read_geoparquet_async, ParquetDataset, ParquetFile, ParquetRead
 pub use writer::{
     write_geoparquet, GeoParquetWriter, GeoParquetWriterEncoding, GeoParquetWriterOptions,
 };
+#[cfg(feature = "parquet_async")]
+pub use writer::{write_geoparquet_async, GeoParquetWriterAsync};

src/io/parquet/writer/async.rs

@@ -0,0 +1,67 @@
+use crate::error::Result;
+use crate::io::parquet::writer::encode::encode_record_batch;
+use crate::io::parquet::writer::metadata::GeoParquetMetadataBuilder;
+use crate::io::parquet::writer::options::GeoParquetWriterOptions;
+use crate::table::GeoTable;
+use arrow_array::RecordBatch;
+use arrow_schema::Schema;
+use parquet::arrow::AsyncArrowWriter;
+use parquet::file::metadata::KeyValue;
+use tokio::io::AsyncWrite;
+
+pub async fn write_geoparquet_async<W: AsyncWrite + Unpin + Send>(
+    table: &mut GeoTable,
+    writer: W,
+    options: &GeoParquetWriterOptions,
+) -> Result<()> {
+    let mut parquet_writer = GeoParquetWriterAsync::try_new(writer, table.schema(), options)?;
+
+    for batch in table.batches() {
+        parquet_writer.write_batch(batch).await?;
+    }
+
+    parquet_writer.finish().await?;
+    Ok(())
+}
+
+pub struct GeoParquetWriterAsync<W: AsyncWrite + Unpin + Send> {
+    writer: AsyncArrowWriter<W>,
+    metadata_builder: GeoParquetMetadataBuilder,
+}
+
+impl<W: AsyncWrite + Unpin + Send> GeoParquetWriterAsync<W> {
+    pub fn try_new(writer: W, schema: &Schema, options: &GeoParquetWriterOptions) -> Result<Self> {
+        let metadata_builder = GeoParquetMetadataBuilder::try_new(schema, options)?;
+
+        let writer = AsyncArrowWriter::try_new(
+            writer,
+            metadata_builder.output_schema.clone(),
+            options.writer_properties.clone(),
+        )?;
+
+        Ok(Self {
+            writer,
+            metadata_builder,
+        })
+    }
+
+    pub async fn write_batch(&mut self, batch: &RecordBatch) -> Result<()> {
+        let encoded_batch = encode_record_batch(batch, &mut self.metadata_builder)?;
+        self.writer.write(&encoded_batch).await?;
+        Ok(())
+    }
+
+    pub fn writer(&self) -> &AsyncArrowWriter<W> {
+        &self.writer
+    }
+
+    pub async fn finish(mut self) -> Result<()> {
+        if let Some(geo_meta) = self.metadata_builder.finish() {
+            let kv_metadata = KeyValue::new("geo".to_string(), serde_json::to_string(&geo_meta)?);
+            self.writer.append_key_value_metadata(kv_metadata);
+        }
+
+        self.writer.close().await?;
+        Ok(())
+    }
+}

src/io/parquet/writer/mod.rs

@@ -1,7 +1,10 @@
+mod r#async;
 mod encode;
 mod metadata;
 mod options;
 mod sync;
 
 pub use options::{GeoParquetWriterEncoding, GeoParquetWriterOptions};
+#[cfg(feature = "parquet_async")]
+pub use r#async::{write_geoparquet_async, GeoParquetWriterAsync};
 pub use sync::{write_geoparquet, GeoParquetWriter};

src/io/parquet/writer/sync.rs

@@ -52,6 +52,10 @@ impl<W: Write + Send> GeoParquetWriter<W> {
         Ok(())
     }
 
+    pub fn writer(&self) -> &ArrowWriter<W> {
+        &self.writer
+    }
+
     pub fn finish(mut self) -> Result<()> {
         if let Some(geo_meta) = self.metadata_builder.finish() {
             let kv_metadata = KeyValue::new("geo".to_string(), serde_json::to_string(&geo_meta)?);