Streaming Writes implementation

src/common/mod.rs

@@ -2,3 +2,6 @@ pub mod properties;
 
 #[cfg(feature = "async")]
 pub mod fetch;
+
+#[cfg(feature = "async")]
+pub mod stream;

src/common/stream.rs

@@ -0,0 +1,31 @@
+use futures::AsyncWrite;
+
+pub struct WrappedWritableStream<'writer> {
+    pub stream: wasm_streams::writable::IntoAsyncWrite<'writer>,
+}
+
+impl<'writer> AsyncWrite for WrappedWritableStream<'writer> {
+    fn poll_write(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+        buf: &[u8],
+    ) -> std::task::Poll<std::io::Result<usize>> {
+        AsyncWrite::poll_write(std::pin::Pin::new(&mut self.get_mut().stream), cx, buf)
+    }
+
+    fn poll_flush(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<std::io::Result<()>> {
+        AsyncWrite::poll_flush(std::pin::Pin::new(&mut self.get_mut().stream), cx)
+    }
+
+    fn poll_close(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<std::io::Result<()>> {
+        AsyncWrite::poll_close(std::pin::Pin::new(&mut self.get_mut().stream), cx)
+    }
+}
+
+unsafe impl<'writer> Send for WrappedWritableStream<'writer> {}

src/lib.rs

@@ -17,6 +17,9 @@ pub mod writer;
 #[cfg(feature = "writer")]
 pub mod writer_properties;
 
+#[cfg(all(feature = "writer", feature = "async"))]
+pub mod writer_async;
+
 // When the `wee_alloc` feature is enabled, use `wee_alloc` as the global
 // allocator.
 /*#[cfg(feature = "wee_alloc")]

src/wasm.rs

@@ -249,3 +249,24 @@ pub async fn read_parquet_stream(
     });
     Ok(wasm_streams::ReadableStream::from_stream(stream).into_raw())
 }
+
+#[wasm_bindgen(js_name = "transformParquetStream")]
+#[cfg(all(feature = "writer", feature = "async"))]
+pub fn transform_parquet_stream(
+    stream: wasm_streams::readable::sys::ReadableStream,
+    writer_properties: Option<crate::writer_properties::WriterProperties>,
+) -> WasmResult<wasm_streams::readable::sys::ReadableStream> {
+    use futures::StreamExt;
+    use wasm_bindgen::convert::TryFromJsValue;
+    let batches = wasm_streams::ReadableStream::from_raw(stream)
+        .into_stream()
+        .map(|maybe_chunk| {
+            let chunk = maybe_chunk.unwrap();
+            arrow_wasm::RecordBatch::try_from_js_value(chunk).unwrap()
+        });
+    let output_stream = super::writer_async::transform_parquet_stream(
+        batches,
+        writer_properties.unwrap_or_default(),
+    );
+    Ok(output_stream.unwrap())
+}

src/writer_async.rs

@@ -0,0 +1,38 @@
+use crate::common::stream::WrappedWritableStream;
+use crate::error::Result;
+use async_compat::CompatExt;
+use futures::StreamExt;
+use parquet::arrow::async_writer::AsyncArrowWriter;
+use wasm_bindgen_futures::spawn_local;
+
+pub fn transform_parquet_stream(
+    batches: impl futures::Stream<Item = arrow_wasm::RecordBatch> + 'static,
+    writer_properties: crate::writer_properties::WriterProperties,
+) -> Result<wasm_streams::readable::sys::ReadableStream> {
+    let options = Some(writer_properties.into());
+    // let encoding = writer_properties.get_encoding();
+
+    let (writable_stream, output_stream) = {
+        let raw_stream = wasm_streams::transform::sys::TransformStream::new().unwrap();
+        let raw_writable = raw_stream.writable();
+        let inner_writer = wasm_streams::WritableStream::from_raw(raw_writable).into_async_write();
+        let writable_stream = WrappedWritableStream {
+            stream: inner_writer,
+        };
+        (writable_stream, raw_stream.readable())
+    };
+    spawn_local::<_>(async move {
+        let mut adapted_stream = batches.peekable();
+        let mut pinned_stream = std::pin::pin!(adapted_stream);
+        let first_batch = pinned_stream.as_mut().peek().await.unwrap();
+        let schema = first_batch.schema().into_inner();
+        // Need to create an encoding for each column
+        let mut writer =
+            AsyncArrowWriter::try_new(writable_stream.compat(), schema, options).unwrap();
+        while let Some(batch) = pinned_stream.next().await {
+            let _ = writer.write(&batch.into()).await;
+        }
+        let _ = writer.close().await;
+    });
+    Ok(output_stream)
+}

tests/js/read-write.test.ts

@@ -1,7 +1,7 @@
 import * as wasm from "../../pkg/node/parquet_wasm";
 import { readFileSync } from "fs";
 import { tableFromIPC, tableToIPC } from "apache-arrow";
-import { testArrowTablesEqual, readExpectedArrowData } from "./utils";
+import { testArrowTablesEqual, readExpectedArrowData, temporaryServer } from "./utils";
 import { describe, it, expect } from "vitest";
 
 // Path from repo root
@@ -89,3 +89,21 @@ it("reads empty file", async (t) => {
   expect(table.numCols).toStrictEqual(0);
   // console.log("empty table schema", table.schema);
 });
+
+it("read stream-write stream-read stream round trip (no writer properties provided)", async (t) => {
+  const server = await temporaryServer();
+  const listeningPort = server.addresses()[0].port;
+  const rootUrl = `http://localhost:${listeningPort}`;
+
+  const expectedTable = readExpectedArrowData();
+
+  const url = `${rootUrl}/1-partition-brotli.parquet`;
+  const originalStream = await wasm.readParquetStream(url);
+
+  const stream = await wasm.transformParquetStream(originalStream);
+  const accumulatedBuffer = new Uint8Array(await new Response(stream).arrayBuffer());
+  const roundtripTable = tableFromIPC(wasm.readParquet(accumulatedBuffer).intoIPCStream());
+
+  testArrowTablesEqual(expectedTable, roundtripTable);
+  await server.close();
+})