feat: multi-line json streaming

crates/datasources/Cargo.toml

@@ -75,6 +75,7 @@ scylla = { version = "0.12.0" }
 glob = "0.3.1"
 indexmap = "2.2.5"
 async-sqlite = "0.2.2"
+json-stream = "0.1.1"
 
 # SSH tunnels
 [target.'cfg(any(target_os = "linux", target_os = "macos"))'.dependencies]

crates/datasources/src/json/errors.rs

@@ -1,3 +1,4 @@
+use datafusion::error::DataFusionError;
 use datafusion_ext::errors::ExtensionError;
 
 use crate::object_store::errors::ObjectStoreSourceError;
@@ -13,21 +14,41 @@ pub enum JsonError {
     #[error("no objects found {0}")]
     NotFound(String),
 
+    #[error("sending data already in progress")]
+    SendAlreadyInProgress,
+
     #[error(transparent)]
     ObjectStoreSource(#[from] ObjectStoreSourceError),
 
     #[error(transparent)]
     ObjectStore(#[from] object_store::Error),
 
+    #[error(transparent)]
+    ObjectStorePath(#[from] object_store::path::Error),
+
     #[error(transparent)]
     Arrow(#[from] datafusion::arrow::error::ArrowError),
 
     #[error(transparent)]
     Datafusion(#[from] datafusion::error::DataFusionError),
+
+    #[error(transparent)]
+    ChannelSend(#[from] futures::channel::mpsc::SendError),
+
+    #[error(transparent)]
+    ChannelRecv(#[from] futures::channel::mpsc::TryRecvError),
 }
 
 impl From<JsonError> for ExtensionError {
     fn from(e: JsonError) -> Self {
         ExtensionError::String(e.to_string())
     }
 }
+
+impl From<JsonError> for DataFusionError {
+    fn from(e: JsonError) -> Self {
+        DataFusionError::External(Box::new(e))
+    }
+}
+
+pub type Result<T, E = JsonError> = std::result::Result<T, E>;

crates/datasources/src/json/stream.rs

@@ -1,97 +1,88 @@
 use std::pin::Pin;
-use std::sync::{Arc, Mutex};
+use std::sync::Arc;
 use std::task::{Context, Poll};
 
 use datafusion::arrow::datatypes::{Schema, SchemaRef};
 use datafusion::arrow::json::ReaderBuilder;
 use datafusion::arrow::record_batch::RecordBatch;
-use datafusion::error::DataFusionError;
+use datafusion::error::{DataFusionError, Result as DFResult};
 use datafusion::execution::TaskContext;
 use datafusion::physical_plan::streaming::PartitionStream;
 use datafusion::physical_plan::{RecordBatchStream, SendableRecordBatchStream};
-use futures::{Stream, StreamExt};
+use futures::channel::mpsc;
+use futures::{SinkExt, Stream, StreamExt, TryStreamExt};
+// use json_stream::JsonStream;
 use object_store::{ObjectMeta, ObjectStore};
-use serde_json::{Map, Value};
+use serde_json::{Deserializer, Map, Value};
 
-use crate::json::errors::JsonError;
-use crate::json::table::push_unwind_json_values;
+use crate::json::errors::{JsonError, Result};
 
-pub type SendableCheckedRecordBatchStream =
-    Pin<Box<dyn Stream<Item = Result<RecordBatch, DataFusionError>> + Send>>;
-
-pub struct JsonStream {
+pub struct JsonRecordBatchStream {
     schema: Arc<Schema>,
-    stream: SendableCheckedRecordBatchStream,
+    // this is the same as a sendable recordbatch stream, but declared
+    // separtley so we can have isomorphic values using adapters.
+    stream: Pin<Box<dyn Stream<Item = DFResult<RecordBatch>> + Send>>,
 }
 
-impl Stream for JsonStream {
-    type Item = Result<RecordBatch, DataFusionError>;
+impl Stream for JsonRecordBatchStream {
+    type Item = DFResult<RecordBatch>;
 
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         self.stream.poll_next_unpin(cx)
     }
 }
 
-impl RecordBatchStream for JsonStream {
+impl RecordBatchStream for JsonRecordBatchStream {
     fn schema(&self) -> SchemaRef {
         self.schema.clone()
     }
 }
 
-pub struct JsonPartitionStream {
+/// WrappedPartition wraps a vector of serde_json documents as a
+/// Partition as one of DataFusion's streaming table. Well all of
+pub struct WrappedPartition {
     schema: Arc<Schema>,
-    stream: Mutex<Option<SendableCheckedRecordBatchStream>>,
+    stream: Vec<Map<String, Value>>,
 }
 
-impl PartitionStream for JsonPartitionStream {
+impl PartitionStream for WrappedPartition {
     fn schema(&self) -> &SchemaRef {
         &self.schema
     }
 
     fn execute(&self, _ctx: Arc<TaskContext>) -> SendableRecordBatchStream {
-        let partition = self
-            .stream
-            .lock()
-            .unwrap()
-            .take()
-            .expect("stream can only be used once")
-            .boxed();
-
-        Box::pin(JsonStream {
-            schema: self.schema.clone(),
-            stream: partition,
-        })
+        Box::pin(JsonStreamHandler::new(
+            self.schema.clone(),
+            futures::stream::iter(self.stream.clone().into_iter().map(Ok)).boxed(),
+        ))
     }
 }
 
-impl JsonPartitionStream {
+impl WrappedPartition {
     pub fn new(schema: Arc<Schema>, chunk: Vec<Map<String, Value>>) -> Self {
-        let stream_schema = schema.clone();
-        let stream = futures::stream::iter(chunk)
-            .chunks(1000)
-            .map(move |objs| {
-                let mut decoder = ReaderBuilder::new(stream_schema.to_owned()).build_decoder()?;
-                decoder
-                    .serialize(&objs)
-                    .map_err(|e| DataFusionError::External(Box::new(e)))?;
-                Ok(decoder.flush()?.unwrap())
-            })
-            .boxed();
-
         Self {
             schema: schema.clone(),
-            stream: Mutex::new(Some(stream)),
+            stream: chunk,
         }
     }
 }
 
-pub(crate) struct LazyJsonPartitionStream {
+/// ObjectStorePartition holds a reference to the object store and the
+/// object metadata and represents a partition that is read only when
+/// the partition is executed.
+pub(crate) struct ObjectStorePartition {
     schema: Arc<Schema>,
     store: Arc<dyn ObjectStore>,
     obj: ObjectMeta,
 }
 
-impl PartitionStream for LazyJsonPartitionStream {
+impl ObjectStorePartition {
+    pub fn new(schema: Arc<Schema>, store: Arc<dyn ObjectStore>, obj: ObjectMeta) -> Self {
+        Self { schema, store, obj }
+    }
+}
+
+impl PartitionStream for ObjectStorePartition {
     fn schema(&self) -> &SchemaRef {
         &self.schema
     }
@@ -101,45 +92,147 @@ impl PartitionStream for LazyJsonPartitionStream {
         let store = self.store.clone();
         let obj = self.obj.clone();
 
-        Box::pin(JsonStream {
+        Box::pin(JsonRecordBatchStream {
             schema: self.schema.clone(),
             stream: futures::stream::once(async move {
-                futures::stream::iter(match Self::build(stream_schema, store, obj).await {
-                    Ok(batches) => batches,
-                    Err(e) => vec![Err(DataFusionError::External(Box::new(e)))],
-                })
+                match JsonStreamHandler::setup_read_stream(store, obj).await {
+                    Ok(st) => JsonStreamHandler::new(stream_schema, st),
+                    Err(e) => JsonStreamHandler::wrap_error(e),
+                }
             })
             .flatten()
             .boxed(),
         })
     }
 }
 
-impl LazyJsonPartitionStream {
-    pub fn new(schema: Arc<Schema>, store: Arc<dyn ObjectStore>, obj: ObjectMeta) -> Self {
-        Self { schema, store, obj }
+/// JsonObjectStream represents a sequence of "json documents" in an
+/// intermediate format produced by serde_json.
+type JsonObjectStream = Pin<Box<dyn Stream<Item = Result<Map<String, Value>>> + Send>>;
+
+/// JsonStreamHandler is the basis of all stream handling, converting
+/// streams of serde_json objects to RecordBatches, including from
+/// object store and from iterators of values.
+///
+/// These are used by the PartitionStream implementations (above)
+/// which interface into the table function and providers.
+struct JsonStreamHandler {
+    schema: Arc<Schema>,
+    buf: Pin<Box<dyn Stream<Item = DFResult<RecordBatch>> + Send>>,
+    worker: tokio::task::JoinHandle<DFResult<()>>,
+}
+
+impl RecordBatchStream for JsonStreamHandler {
+    fn schema(&self) -> SchemaRef {
+        self.schema.clone()
     }
+}
+
+impl Stream for JsonStreamHandler {
+    type Item = DFResult<RecordBatch>;
 
-    async fn build(
-        schema: Arc<Schema>,
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        match self.buf.poll_next_unpin(cx) {
+            Poll::Pending => Poll::Pending,
+            Poll::Ready(None) => {
+                if self.worker.is_finished() {
+                    Poll::Pending
+                } else {
+                    Poll::Ready(None)
+                }
+            }
+            Poll::Ready(Some(val)) => Poll::Ready(Some(val)),
+        }
+    }
+}
+
+impl JsonStreamHandler {
+    fn new(schema: Arc<Schema>, stream: JsonObjectStream) -> Self {
+        let stream_schema = schema.clone();
+        let (mut sink, buf) = mpsc::channel(2048);
+
+        Self {
+            schema,
+            buf: buf
+                .chunks(1024)
+                .map(move |chunk| {
+                    let mut decoder =
+                        ReaderBuilder::new(stream_schema.to_owned()).build_decoder()?;
+                    decoder
+                        .serialize(&chunk)
+                        .map_err(|e| DataFusionError::External(Box::new(e)))?;
+                    Ok(decoder.flush()?.unwrap())
+                })
+                .boxed(),
+            worker: tokio::spawn(async move {
+                let mut stream = stream;
+                while let Some(row) = stream.next().await {
+                    sink.send(row?).await.map_err(JsonError::from)?;
+                }
+                Ok(())
+            }),
+        }
+    }
+
+    fn wrap_error(err: JsonError) -> Self {
+        Self {
+            schema: Arc::new(Schema::empty()),
+            buf: futures::stream::empty().boxed(),
+            worker: tokio::task::spawn_local(async move { Err(err.into()) }),
+        }
+    }
+
+    async fn setup_read_stream(
         store: Arc<dyn ObjectStore>,
         obj: ObjectMeta,
-    ) -> Result<Vec<Result<RecordBatch, DataFusionError>>, JsonError> {
-        let mut data = Vec::new();
-        push_unwind_json_values(
-            &mut data,
-            serde_json::from_slice::<Value>(&store.get(&obj.location).await?.bytes().await?),
-        )?;
-
-        Ok(data
-            .chunks(1000)
-            .map(|chunk| {
-                let mut decoder = ReaderBuilder::new(schema.clone()).build_decoder()?;
-                decoder
-                    .serialize(chunk)
-                    .map_err(|e| DataFusionError::External(Box::new(e)))?;
-                Ok(decoder.flush()?.unwrap())
-            })
-            .collect())
+    ) -> Result<JsonObjectStream> {
+        // Ok(JsonStream::<Value, _>::new(
+        //     store
+        //         .get(&obj.location)
+        //         .await?
+        //         .into_stream()
+        //         .map_err(JsonError::ObjectStore),
+        // )
+        // .flat_map(Self::unwind_json_value)
+        // .boxed())
+        Ok(Box::pin(
+            futures::stream::iter(
+                Deserializer::from_reader(std::io::BufReader::new(std::io::Cursor::new(
+                    store.get(&obj.location).await?.bytes().await?.to_vec(),
+                )))
+                .into_iter(),
+            )
+            .map_err(JsonError::SerdeJson)
+            .flat_map(Self::unwind_json_value),
+        ))
+    }
+
+    fn unwind_json_value(input: Result<Value>) -> JsonObjectStream {
+        futures::stream::iter(match input {
+            Ok(value) => match value {
+                Value::Array(vals) => {
+                    let mut out = Vec::with_capacity(vals.len());
+                    for v in vals {
+                        match v {
+                            Value::Object(doc) => out.push(Ok(doc)),
+                            Value::Null => out.push(Ok(Map::new())),
+                            _ => out.push(Err(JsonError::UnspportedType(
+                                "only objects and arrays of objects are supported",
+                            ))),
+                        }
+                    }
+                    out
+                }
+                Value::Object(doc) => vec![Ok(doc)],
+                Value::Null => vec![Ok(Map::new())],
+                _ => {
+                    vec![Err(JsonError::UnspportedType(
+                        "only objects and arrays of objects are supported",
+                    ))]
+                }
+            },
+            Err(e) => vec![Err(e)],
+        })
+        .boxed()
     }
 }