a few things

- add integration test for vortex-file
- fix bug in PathStrategy descend
- return Writer back from the write pathway

Signed-off-by: Andrew Duffy <[email protected]>

Author: a10y

vortex-file/src/counting.rs

@@ -11,8 +11,8 @@ use vortex_io::VortexWrite;
 
 /// A wrapper around an `VortexWrite` that counts the number of bytes written.
 pub(crate) struct CountingVortexWrite<W> {
-    inner: W,
-    bytes_written: Arc<AtomicU64>,
+    pub(crate) inner: W,
+    pub(crate) bytes_written: Arc<AtomicU64>,
 }
 
 impl<W: VortexWrite> CountingVortexWrite<W> {

vortex-file/src/writer.rs

@@ -128,7 +128,7 @@ impl VortexWriteOptions {
         self,
         write: W,
         stream: S,
-    ) -> VortexResult<WriteSummary> {
+    ) -> VortexResult<(WriteSummary, W)> {
         self.write_internal(write, ArrayStreamExt::boxed(stream))
             .await
     }
@@ -137,7 +137,7 @@ impl VortexWriteOptions {
         self,
         mut write: W,
         stream: SendableArrayStream,
-    ) -> VortexResult<WriteSummary> {
+    ) -> VortexResult<(WriteSummary, W)> {
         // Set up a Context to capture the encodings used in the file.
         let ctx = ArrayContext::empty();
         let dtype = stream.dtype().clone();
@@ -216,10 +216,13 @@ impl VortexWriteOptions {
 
         write.flush().await?;
 
-        Ok(WriteSummary {
-            footer,
-            size: position,
-        })
+        Ok((
+            WriteSummary {
+                footer,
+                size: position,
+            },
+            write,
+        ))
     }
 
     /// Create a push-based [`Writer`] that can be used to incrementally write arrays to the file.
@@ -233,7 +236,11 @@ impl VortexWriteOptions {
         let write = CountingVortexWrite::new(write);
         let bytes_written = write.counter();
         let strategy = self.strategy.clone();
-        let future = self.write(write, arrays).boxed_local().fuse();
+        let future = self
+            .write(write, arrays)
+            .map(move |result| result.map(|(summary, _writer)| summary))
+            .boxed_local()
+            .fuse();
 
         Writer {
             arrays: Some(arrays_send),
@@ -370,11 +377,13 @@ impl<'rt, B: BlockingRuntime> BlockingWrite<'rt, B> {
         write: W,
         iter: impl ArrayIterator + Send + 'static,
     ) -> VortexResult<WriteSummary> {
-        self.runtime.block_on(async move {
-            self.options
-                .write(BlockingWriteAdapter(write), iter.into_array_stream())
-                .await
-        })
+        self.runtime
+            .block_on(async move {
+                self.options
+                    .write(BlockingWriteAdapter(write), iter.into_array_stream())
+                    .await
+            })
+            .map(|(summary, _)| summary)
     }
 
     pub fn writer<'w, W: Write + Unpin + 'w>(

vortex-file/tests/test_write_table.rs

@@ -0,0 +1,113 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+#![allow(clippy::tests_outside_test_module)]
+
+use std::sync::Arc;
+use std::sync::LazyLock;
+
+use futures::StreamExt;
+use futures::pin_mut;
+use vortex_array::Array;
+use vortex_array::IntoArray;
+use vortex_array::ToCanonical;
+use vortex_array::arrays::PrimitiveArray;
+use vortex_array::arrays::StructArray;
+use vortex_array::expr::session::ExprSession;
+use vortex_array::session::ArraySession;
+use vortex_array::validity::Validity;
+use vortex_buffer::ByteBuffer;
+use vortex_dtype::FieldNames;
+use vortex_dtype::field_path;
+use vortex_file::OpenOptionsSessionExt;
+use vortex_file::WriteOptionsSessionExt;
+use vortex_io::session::RuntimeSession;
+use vortex_layout::layouts::compressed::CompressingStrategy;
+use vortex_layout::layouts::flat::writer::FlatLayoutStrategy;
+use vortex_layout::layouts::path::PathStrategy;
+use vortex_layout::session::LayoutSession;
+use vortex_metrics::VortexMetrics;
+use vortex_session::VortexSession;
+
+static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+    let session = VortexSession::empty()
+        .with::<VortexMetrics>()
+        .with::<ArraySession>()
+        .with::<LayoutSession>()
+        .with::<ExprSession>()
+        .with::<RuntimeSession>();
+
+    vortex_file::register_default_encodings(&session);
+
+    session
+});
+
+#[tokio::test]
+async fn test_file_roundtrip() {
+    // Create a simple roundtrip
+    let nums = PrimitiveArray::from_iter((0..1024).cycle().take(16_384)).into_array();
+
+    let a_array = StructArray::new(
+        FieldNames::from(["raw", "compressed"]),
+        vec![nums.clone(), nums.clone()],
+        16_384,
+        Validity::NonNullable,
+    )
+    .into_array();
+
+    let b_array = PrimitiveArray::from_iter((1024..2048).cycle().take(16_384)).into_array();
+
+    let data = StructArray::new(
+        FieldNames::from(["a", "b"]),
+        vec![a_array, b_array],
+        16_384,
+        Validity::NonNullable,
+    )
+    .into_array();
+
+    // Create a writer which by default uses the BtrBlocks compressor, but for a.raw column it will
+    // leave it uncompressed.
+    let default_strategy = Arc::new(CompressingStrategy::new_btrblocks(
+        FlatLayoutStrategy::default(),
+        false,
+    ));
+
+    let writer = Arc::new(
+        PathStrategy::new(Arc::new(FlatLayoutStrategy::default()), default_strategy)
+            .set_field_writer(field_path!(a.raw), Arc::new(FlatLayoutStrategy::default())),
+    );
+
+    let bytes = Vec::new();
+    let (_, bytes) = SESSION
+        .write_options()
+        .with_strategy(writer)
+        .write(bytes, data.to_array_stream())
+        .await
+        .expect("write");
+
+    let bytes = ByteBuffer::from(bytes);
+    let vxf = SESSION.open_options().open(bytes).await.expect("open");
+
+    // Read the data back
+    let stream = vxf
+        .scan()
+        .expect("scan")
+        .into_stream()
+        .expect("into_stream");
+
+    pin_mut!(stream);
+
+    while let Some(next) = stream.next().await {
+        let next = next.expect("next");
+        let next = next.to_struct();
+        let a = next.field_by_name("a").unwrap().to_struct();
+
+        let raw = a.field_by_name("raw").unwrap();
+        let compressed = a.field_by_name("compressed").unwrap();
+
+        assert!(raw.is_canonical());
+        assert!(!compressed.is_canonical());
+
+        assert!(raw.nbytes() > compressed.nbytes());
+    }
+}

vortex-layout/src/layouts/path.rs

@@ -133,8 +133,15 @@ impl PathStrategy {
     fn descend(&self, field: &Field) -> Self {
         // Start with the existing set of overrides, then only retain the ones that contain
         // the current field
-        let mut new_writers = self.leaf_writers.clone();
-        new_writers.retain(|k, _| k.starts_with_field(field));
+        let mut new_writers = HashMap::with_capacity(self.leaf_writers.len());
+
+        for (field_path, strategy) in &self.leaf_writers {
+            if field_path.starts_with_field(field)
+                && let Some(subpath) = field_path.clone().step_into()
+            {
+                new_writers.insert(subpath, strategy.clone());
+            }
+        }
 
         Self {
             leaf_writers: new_writers,
@@ -256,7 +263,6 @@ impl LayoutStrategy for PathStrategy {
             .zip_eq(column_names)
             .enumerate()
             .map(move |(index, ((dtype, recv), name))| {
-                println!("PathStrategy visiting {name}");
                 let column_stream =
                     SequentialStreamAdapter::new(dtype.clone(), recv.into_stream().boxed())
                         .sendable();