Fix windows issue and other touchups

Signed-off-by: Adam Gutglick <[email protected]>

Author: AdamGS

vortex-datafusion/src/persistent/mod.rs

@@ -66,7 +66,6 @@ mod tests {
     use vortex::array::arrays::VarBinArray;
     use vortex::array::validity::Validity;
     use vortex::buffer::buffer;
-    use vortex::error::vortex_err;
     use vortex::file::WriteOptionsSessionExt;
     use vortex::session::VortexSession;
 
@@ -75,10 +74,8 @@ mod tests {
     use crate::persistent::register_vortex_format_factory;
 
     #[rstest]
-    #[case(Some(1))]
-    #[case(None)]
     #[tokio::test]
-    async fn query_file(#[case] limit: Option<usize>) -> anyhow::Result<()> {
+    async fn test_query_file(#[values(Some(1), None)] limit: Option<usize>) -> anyhow::Result<()> {
         let session = VortexSession::default();
         let temp_dir = tempdir()?;
         let strings = ChunkedArray::from_iter([
@@ -100,28 +97,26 @@ mod tests {
             Validity::NonNullable,
         )?;
 
-        let filepath = temp_dir.path().join("data.vortex");
-
         let mut f = OpenOptions::new()
             .write(true)
-            .create(true)
-            .truncate(true)
-            .open(&filepath)
+            .create_new(true)
+            .open(temp_dir.path().join("data.vortex"))
             .await?;
 
-        session
+        let summary = session
             .write_options()
             .write(&mut f, st.to_array_stream())
             .await?;
 
+        assert_eq!(summary.row_count(), 8);
+
+        // For reasons I don't understand, this is suddenly important on windows
+        drop(f);
+
         let ctx = SessionContext::default();
         let format = Arc::new(VortexFormat::new(session));
-        let table_url = ListingTableUrl::parse(
-            temp_dir
-                .path()
-                .to_str()
-                .ok_or_else(|| vortex_err!("Path is not valid UTF-8"))?,
-        )?;
+
+        let table_url = ListingTableUrl::parse(temp_dir.path().to_str().unwrap())?;
         assert!(table_url.is_collection());
 
         let config = ListingTableConfig::new(table_url)

vortex-datafusion/src/persistent/opener.rs

@@ -335,7 +335,7 @@ impl FileOpener for VortexOpener {
                 })
                 .try_flatten()
                 .map(move |batch| {
-                    if projector.projection().column_indices().is_empty() {
+                    if projector.projection().as_ref().is_empty() {
                         batch
                     } else {
                         batch.and_then(|b| projector.project_batch(&b))

vortex-datafusion/src/persistent/sink.rs

@@ -3,13 +3,12 @@
 
 use std::any::Any;
 use std::sync::Arc;
-use std::sync::atomic::AtomicU64;
-use std::sync::atomic::Ordering;
 
 use arrow_schema::SchemaRef;
 use async_trait::async_trait;
 use datafusion_common::DataFusionError;
 use datafusion_common::Result as DFResult;
+use datafusion_common::exec_datafusion_err;
 use datafusion_common_runtime::JoinSet;
 use datafusion_common_runtime::SpawnedTask;
 use datafusion_datasource::file_sink_config::FileSink;
@@ -33,6 +32,7 @@ use vortex::dtype::DType;
 use vortex::dtype::arrow::FromArrowType;
 use vortex::error::VortexResult;
 use vortex::file::WriteOptionsSessionExt;
+use vortex::file::WriteSummary;
 use vortex::io::ObjectStoreWriter;
 use vortex::io::VortexWrite;
 use vortex::session::VortexSession;
@@ -108,58 +108,52 @@ impl FileSink for VortexSink {
         mut file_stream_rx: DemuxedStreamReceiver,
         object_store: Arc<dyn ObjectStore>,
     ) -> DFResult<u64> {
-        // This is a hack
-        let row_counter = Arc::new(AtomicU64::new(0));
-
-        let mut file_write_tasks: JoinSet<DFResult<Path>> = JoinSet::new();
+        let mut file_write_tasks: JoinSet<DFResult<(Path, WriteSummary)>> = JoinSet::new();
 
         // TODO(adamg):
         // 1. We can probably be better at signaling how much memory we're consuming (potentially when reading too), see ParquetSink::spawn_writer_tasks_and_join.
         while let Some((path, rx)) = file_stream_rx.recv().await {
             let session = self.session.clone();
-            let row_counter = row_counter.clone();
             let object_store = object_store.clone();
             let writer_schema = get_writer_schema(&self.config);
             let dtype = DType::from_arrow(writer_schema);
 
             // We need to spawn work because there's a dependency between the different files. If one file has too many batches buffered,
             // the demux task might deadlock itself.
             file_write_tasks.spawn(async move {
-                let stream = ReceiverStream::new(rx).map(move |rb| {
-                    row_counter.fetch_add(rb.num_rows() as u64, Ordering::Relaxed);
-                    VortexResult::Ok(ArrayRef::from_arrow(rb, false))
-                });
+                let stream = ReceiverStream::new(rx)
+                    .map(move |rb| VortexResult::Ok(ArrayRef::from_arrow(rb, false)));
 
                 let stream_adapter = ArrayStreamAdapter::new(dtype, stream);
 
-                let mut sink = ObjectStoreWriter::new(object_store.clone(), &path)
+                let mut object_writer = ObjectStoreWriter::new(object_store, &path)
                     .await
-                    .map_err(|e| {
-                        DataFusionError::Execution(format!(
-                            "Failed to create ObjectStoreWriter: {e}"
-                        ))
-                    })?;
+                    .map_err(|e| exec_datafusion_err!("Failed to create ObjectStoreWriter: {e}"))?;
 
-                session
+                let summary = session
                     .write_options()
-                    .write(&mut sink, stream_adapter)
+                    .write(&mut object_writer, stream_adapter)
                     .await
-                    .map_err(|e| {
-                        DataFusionError::Execution(format!("Failed to write Vortex file: {e}"))
-                    })?;
+                    .map_err(|e| exec_datafusion_err!("Failed to write Vortex file: {e}"))?;
 
-                sink.shutdown().await.map_err(|e| {
-                    DataFusionError::Execution(format!("Failed to shutdown Vortex writer: {e}"))
-                })?;
+                object_writer
+                    .shutdown()
+                    .await
+                    .map_err(|e| exec_datafusion_err!("Failed to shutdown Vortex writer: {e}"))?;
 
-                Ok(path)
+                Ok((path, summary))
             });
         }
 
+        let mut row_count = 0;
+
         while let Some(result) = file_write_tasks.join_next().await {
             match result {
-                Ok(path) => {
-                    let path = path?;
+                Ok(r) => {
+                    let (path, summary) = r?;
+
+                    row_count += summary.row_count();
+
                     tracing::info!(path = %path, "Successfully written file");
                 }
                 Err(e) => {
@@ -177,7 +171,7 @@ impl FileSink for VortexSink {
             .await
             .map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??;
 
-        Ok(row_counter.load(Ordering::SeqCst))
+        Ok(row_count)
     }
 }