Skip to main content

vortex_file/
writer.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::io;
5use std::io::Write;
6use std::sync::Arc;
7use std::sync::atomic::AtomicU64;
8use std::sync::atomic::Ordering;
9
10use futures::FutureExt;
11use futures::StreamExt;
12use futures::TryStreamExt;
13use futures::future::Fuse;
14use futures::future::LocalBoxFuture;
15use futures::future::ready;
16use futures::pin_mut;
17use futures::select;
18use itertools::Itertools;
19use vortex_array::ArrayContext;
20use vortex_array::ArrayRef;
21use vortex_array::dtype::DType;
22use vortex_array::dtype::FieldPath;
23use vortex_array::expr::stats::Stat;
24use vortex_array::iter::ArrayIterator;
25use vortex_array::iter::ArrayIteratorExt;
26use vortex_array::stats::PRUNING_STATS;
27use vortex_array::stream::ArrayStream;
28use vortex_array::stream::ArrayStreamAdapter;
29use vortex_array::stream::ArrayStreamExt;
30use vortex_array::stream::SendableArrayStream;
31use vortex_buffer::ByteBuffer;
32use vortex_edition::EditionSessionExt;
33use vortex_error::VortexError;
34use vortex_error::VortexExpect;
35use vortex_error::VortexResult;
36use vortex_error::vortex_bail;
37use vortex_error::vortex_err;
38use vortex_io::IoBuf;
39use vortex_io::VortexWrite;
40use vortex_io::kanal_ext::KanalExt;
41use vortex_io::runtime::BlockingRuntime;
42use vortex_io::session::RuntimeSessionExt;
43use vortex_layout::BufferedBytesTracker;
44use vortex_layout::LayoutStrategy;
45use vortex_layout::LayoutWriterContext;
46use vortex_layout::layouts::file_stats::accumulate_stats;
47use vortex_layout::sequence::SequenceId;
48use vortex_layout::sequence::SequentialStreamAdapter;
49use vortex_layout::sequence::SequentialStreamExt;
50use vortex_session::SessionExt;
51use vortex_session::VortexSession;
52use vortex_session::registry::ReadContext;
53use vortex_utils::aliases::hash_map::HashMap;
54
55use crate::Footer;
56use crate::MAGIC_BYTES;
57use crate::WriteStrategyBuilder;
58use crate::counting::CountingVortexWrite;
59use crate::footer::FileStatistics;
60use crate::footer::MAX_METADATA_KEY_BYTES;
61use crate::footer::MAX_METADATA_SEGMENTS;
62use crate::segments::writer::BufferedSegmentSink;
63
64/// Configure a new writer, which can eventually be used to write an [`ArrayStream`] into a sink
65/// that implements [`VortexWrite`].
66///
67/// All write strategies are restricted to the encodings in the session's enabled editions.
68///
69/// Construct with [`WriteOptionsSessionExt::write_options`] for normal use so the writer inherits
70/// the session's runtime, array registry, and memory configuration.
71pub struct VortexWriteOptions {
72    session: VortexSession,
73    strategy: Arc<dyn LayoutStrategy>,
74    buffered_bytes: BufferedBytesTracker,
75    exclude_dtype: bool,
76    max_variable_length_statistics_size: usize,
77    file_statistics: Vec<Stat>,
78    metadata: HashMap<String, ByteBuffer>,
79}
80
81/// Extension trait for constructing [`VortexWriteOptions`] from a session.
82pub trait WriteOptionsSessionExt: SessionExt {
83    /// Create [`VortexWriteOptions`] for writing to a Vortex file.
84    fn write_options(&self) -> VortexWriteOptions {
85        VortexWriteOptions::new(self.session())
86    }
87}
88impl<S: SessionExt> WriteOptionsSessionExt for S {}
89
90impl VortexWriteOptions {
91    /// Create a new [`VortexWriteOptions`] with the given session.
92    pub fn new(session: VortexSession) -> Self {
93        let strategy = WriteStrategyBuilder::default()
94            .with_allow_encodings(session.enabled_encoding_ids().into_iter().collect())
95            .build();
96        VortexWriteOptions {
97            strategy,
98            buffered_bytes: BufferedBytesTracker::new(),
99            session,
100            exclude_dtype: false,
101            file_statistics: PRUNING_STATS.to_vec(),
102            max_variable_length_statistics_size: 64,
103            metadata: HashMap::default(),
104        }
105    }
106
107    /// Replace the default layout strategy with the provided one.
108    ///
109    /// The strategy controls repartitioning, statistics layout, compression, and leaf segment
110    /// emission. Use [`WriteStrategyBuilder`] when only a small part of the default strategy needs
111    /// customization. Replacing the strategy does not change the enabled-edition encoding policy.
112    pub fn with_strategy(mut self, strategy: Arc<dyn LayoutStrategy>) -> Self {
113        self.strategy = strategy;
114        self
115    }
116
117    /// Returns the tracker accounting for bytes that layout strategies are holding but have not
118    /// yet emitted.
119    ///
120    /// The tracker is shared with the write these options start, so it can be captured before
121    /// calling [`Self::write`] and polled while the write runs. [`Writer::buffered_bytes`] exposes
122    /// the same counter for the push-based API.
123    pub fn buffered_bytes_tracker(&self) -> BufferedBytesTracker {
124        self.buffered_bytes.clone()
125    }
126
127    /// Exclude the DType from the Vortex file. You must provide the DType to the reader.
128    // TODO(ngates): Should we store some sort of DType checksum to make sure the one passed at
129    //  read-time is sane? I guess most layouts will have some reasonable validation.
130    pub fn exclude_dtype(mut self) -> Self {
131        self.exclude_dtype = true;
132        self
133    }
134
135    /// Configure which statistics to compute at the file level.
136    ///
137    /// Pass an empty vector to omit file-level statistics.
138    pub fn with_file_statistics(mut self, file_statistics: Vec<Stat>) -> Self {
139        self.file_statistics = file_statistics;
140        self
141    }
142
143    /// Add a user-defined metadata segment (keyed, opaque bytes); a repeated key replaces the
144    /// previous value. Keys and the segment count are validated against `MAX_METADATA_*`.
145    pub fn with_metadata_segment(
146        mut self,
147        key: impl Into<String>,
148        metadata: impl Into<ByteBuffer>,
149    ) -> Self {
150        let key = key.into();
151        let metadata = metadata.into();
152        self.metadata.insert(key, metadata);
153        self
154    }
155
156    /// Add user-defined metadata segments to the file.
157    ///
158    /// If a key already exists, the previous segment for that key is replaced.
159    pub fn with_metadata_segments<I, K, B>(mut self, metadata: I) -> Self
160    where
161        I: IntoIterator<Item = (K, B)>,
162        K: Into<String>,
163        B: Into<ByteBuffer>,
164    {
165        for (key, metadata) in metadata {
166            self = self.with_metadata_segment(key, metadata);
167        }
168        self
169    }
170
171    /// Check the configured metadata segments against the `MAX_METADATA_*` limits.
172    ///
173    /// [`Self::write`] performs the same check, but only once the sink is already being written
174    /// to. Callers that accept metadata from elsewhere (FFI bindings, for example) can use this to
175    /// reject an invalid set before any bytes are produced.
176    pub fn validate_metadata(&self) -> VortexResult<()> {
177        validate_metadata_segments(&self.metadata)
178    }
179}
180
181impl VortexWriteOptions {
182    /// Drop into the blocking writer API using the given runtime.
183    ///
184    /// The returned adapter drives async writer internals on `runtime` while accepting ordinary
185    /// [`std::io::Write`] sinks and [`ArrayIterator`] inputs.
186    pub fn blocking<B: BlockingRuntime>(self, runtime: &B) -> BlockingWrite<'_, B> {
187        BlockingWrite {
188            options: self,
189            runtime,
190        }
191    }
192
193    /// Write an [`ArrayStream`] as a Vortex file.
194    ///
195    /// Note that buffers are flushed as soon as they are available with no buffering, the caller
196    /// is responsible for deciding how to configure buffering on the underlying `Write` sink.
197    ///
198    /// The set of encodings permitted in the file is snapshotted from the session's array registry
199    /// here, so encodings registered after this call are not written.
200    pub async fn write<W: VortexWrite + Unpin, S: ArrayStream + Send + 'static>(
201        self,
202        write: W,
203        stream: S,
204    ) -> VortexResult<WriteSummary> {
205        self.write_internal(write, ArrayStreamExt::boxed(stream))
206            .await
207    }
208
209    async fn write_internal<W: VortexWrite + Unpin>(
210        self,
211        mut write: W,
212        stream: SendableArrayStream,
213    ) -> VortexResult<WriteSummary> {
214        validate_metadata_segments(&self.metadata)?;
215
216        // The array context is built here, rather than when the options were constructed, so that
217        // encodings registered on the session in between are still eligible for the file.
218        let ctx = LayoutWriterContext::new(new_array_context(&self.session))
219            .with_buffered_bytes_tracker(self.buffered_bytes.clone());
220        let dtype = stream.dtype().clone();
221
222        let (mut ptr, eof) = SequenceId::root().split();
223
224        let stream = SequentialStreamAdapter::new(
225            dtype.clone(),
226            stream
227                .try_filter(|chunk| ready(!chunk.is_empty()))
228                .map(move |result| result.map(|chunk| (ptr.advance(), chunk))),
229        )
230        .sendable();
231        let (file_stats, stream) = accumulate_stats(
232            stream,
233            self.file_statistics.clone().into(),
234            self.max_variable_length_statistics_size,
235            &self.session,
236        );
237
238        // First, write the magic bytes.
239        write.write_all(ByteBuffer::copy_from(MAGIC_BYTES)).await?;
240        let mut position = MAGIC_BYTES.len() as u64;
241
242        // Create a channel to send buffers from the segment sink to the output stream.
243        let (send, recv) = kanal::bounded_async(1);
244
245        let segments = Arc::new(BufferedSegmentSink::new(send, position));
246
247        // We spawn the layout future so it is driven in the background while we write the
248        // buffer stream, so we don't need to poll it until all buffers have been drained.
249        let ctx2 = ctx.clone();
250        let session = self.session.clone();
251        let layout_fut = self.session.handle().spawn_nested(move |h| async move {
252            let session = session.with_handle(h);
253            let layout = self
254                .strategy
255                .write_stream(
256                    ctx2,
257                    Arc::<BufferedSegmentSink>::clone(&segments),
258                    stream,
259                    eof,
260                    &session,
261                )
262                .await?;
263            Ok::<_, VortexError>((layout, segments.segment_specs()))
264        });
265
266        // Flush buffers as they arrive
267        let recv_stream = recv.into_stream();
268        pin_mut!(recv_stream);
269        while let Some(buffer) = recv_stream.next().await {
270            if buffer.is_empty() {
271                continue;
272            }
273            position += buffer.len() as u64;
274            write.write_all(buffer).await?;
275        }
276
277        let (layout, segment_specs) = layout_fut.await?;
278
279        // Assemble the Footer object now that we have all the segments.
280        let statistics = if self.file_statistics.is_empty() {
281            None
282        } else {
283            Some(FileStatistics::new_with_dtype(
284                file_stats.stats_sets().into(),
285                &dtype,
286            ))
287        };
288        let mut footer = Footer::new(
289            Arc::clone(&layout),
290            segment_specs,
291            statistics,
292            ReadContext::new(ctx.array_ctx().to_ids()),
293        );
294
295        // Emit the footer buffers and EOF.
296        let (footer_buffers, metadata, approx_byte_size) = footer
297            .clone()
298            .into_serializer()
299            .with_metadata_segments(self.metadata)
300            .with_offset(position)
301            .with_exclude_dtype(self.exclude_dtype)
302            .serialize_with_metadata()?;
303        footer = footer
304            .with_metadata_segments(metadata)
305            .with_approx_byte_size(approx_byte_size);
306
307        for buffer in footer_buffers {
308            position += buffer.len() as u64;
309            write.write_all(buffer).await?;
310        }
311
312        write.flush().await?;
313
314        Ok(WriteSummary {
315            footer,
316            size: position,
317        })
318    }
319
320    /// Create a push-based [`Writer`] that can be used to incrementally write arrays to the file.
321    ///
322    /// Each pushed chunk must have dtype `dtype`. Call [`Writer::finish`] to close the input stream,
323    /// flush remaining buffers, and receive the [`WriteSummary`].
324    pub fn writer<'w, W: VortexWrite + Unpin + 'w>(self, write: W, dtype: DType) -> Writer<'w> {
325        // Create a channel for sending arrays to the layout task.
326        let (arrays_send, arrays_recv) = kanal::bounded_async(1);
327
328        let arrays =
329            ArrayStreamExt::boxed(ArrayStreamAdapter::new(dtype, arrays_recv.into_stream()));
330
331        let write = CountingVortexWrite::new(write);
332        let bytes_written = write.counter();
333        let buffered_bytes = self.buffered_bytes.clone();
334        let future = self.write(write, arrays).boxed_local().fuse();
335
336        Writer {
337            arrays: Some(arrays_send),
338            future,
339            bytes_written,
340            buffered_bytes,
341        }
342    }
343}
344
345fn new_array_context(session: &VortexSession) -> ArrayContext {
346    // NOTE(os): Setup an array context that already has all known encodings pre-populated.
347    // This is preferred for now over having an empty context here, because only the
348    // serialised array order is deterministic. The serialisation of arrays are done
349    // parallel and with an empty context they can register their encodings to the context
350    // in different order, changing the written bytes from run to run.
351    let enabled_encoding_ids = session.enabled_encoding_ids();
352    ArrayContext::new(enabled_encoding_ids.iter().cloned().sorted().collect())
353        // Only permit encodings known to the session.
354        .with_allowed_ids(enabled_encoding_ids.into_iter().collect())
355}
356
357fn validate_metadata_segments(metadata: &HashMap<String, ByteBuffer>) -> VortexResult<()> {
358    if metadata.len() > MAX_METADATA_SEGMENTS {
359        vortex_bail!(
360            "Vortex files may contain at most {} metadata segments; got {} metadata segments. Metadata keys must be non-empty and at most {} bytes",
361            MAX_METADATA_SEGMENTS,
362            metadata.len(),
363            MAX_METADATA_KEY_BYTES
364        );
365    }
366
367    for key in metadata.keys() {
368        if key.is_empty() {
369            vortex_bail!(
370                "Vortex metadata keys must be non-empty and at most {} bytes; files may contain at most {} metadata segments",
371                MAX_METADATA_KEY_BYTES,
372                MAX_METADATA_SEGMENTS
373            );
374        }
375
376        let key_bytes = key.len();
377        if key_bytes > MAX_METADATA_KEY_BYTES {
378            vortex_bail!(
379                "Vortex metadata key {key:?} is {key_bytes} bytes, but keys must be at most {} bytes; files may contain at most {} metadata segments",
380                MAX_METADATA_KEY_BYTES,
381                MAX_METADATA_SEGMENTS
382            );
383        }
384    }
385
386    Ok(())
387}
388
389/// An async API for writing Vortex files.
390pub struct Writer<'w> {
391    // The input channel for sending arrays to the writer.
392    arrays: Option<kanal::AsyncSender<VortexResult<ArrayRef>>>,
393    // The writer task that ultimately produces the footer.
394    future: Fuse<LocalBoxFuture<'w, VortexResult<WriteSummary>>>,
395    // The bytes written so far.
396    bytes_written: Arc<AtomicU64>,
397    // The buffered bytes accounting shared with the layout strategies for this write.
398    buffered_bytes: BufferedBytesTracker,
399}
400
401impl Writer<'_> {
402    /// Push a new chunk into the writer.
403    pub async fn push(&mut self, chunk: ArrayRef) -> VortexResult<()> {
404        let arrays = self.arrays.clone().vortex_expect("missing arrays sender");
405        let send_fut = async move { arrays.send(Ok(chunk)).await }.fuse();
406        pin_mut!(send_fut);
407
408        // We poll the writer future to continue writing bytes to the output, while waiting for
409        // enough room to push the next chunk into the channel.
410        select! {
411            result = send_fut => {
412                // If the send future failed, the writer has failed or panicked.
413                if result.is_err() {
414                    return Err(self.handle_failed_task().await);
415                }
416            },
417            result = &mut self.future => {
418                // Under normal operation, the writer future should never complete until
419                // finish() is called. Therefore, we can assume the writer has failed.
420                // The writer future has failed, we need to propagate the error.
421                match result {
422                    Ok(_) => vortex_bail!("Internal error: writer future completed early"),
423                    Err(e) => return Err(e),
424                }
425            }
426        }
427
428        Ok(())
429    }
430
431    /// Push an entire [`ArrayStream`] into the writer, consuming it.
432    ///
433    /// A task is spawned to consume the stream and push it into the writer, with the current
434    /// thread being used to write buffers to the output.
435    pub async fn push_stream(&mut self, mut stream: SendableArrayStream) -> VortexResult<()> {
436        let arrays = self.arrays.clone().vortex_expect("missing arrays sender");
437        let stream_fut = async move {
438            while let Some(chunk) = stream.next().await {
439                arrays.send(chunk).await?;
440            }
441            Ok::<_, kanal::SendError>(())
442        }
443        .fuse();
444        pin_mut!(stream_fut);
445
446        // We poll the writer future to continue writing bytes to the output, while waiting for
447        // enough room to push the stream into the channel.
448        select! {
449            result = stream_fut => {
450                if let Err(_send_err) = result {
451                    // If the send future failed, the writer has failed or panicked.
452                    return Err(self.handle_failed_task().await);
453                }
454            }
455
456            result = &mut self.future => {
457                // Under normal operation, the writer future should never complete until
458                // finish() is called. Therefore, we can assume the writer has failed.
459                // The writer future has failed, we need to propagate the error.
460                match result {
461                    Ok(_) => vortex_bail!("Internal error: writer future completed early"),
462                    Err(e) => return Err(e),
463                }
464            }
465        }
466
467        Ok(())
468    }
469
470    /// Returns the number of bytes written to the file so far.
471    pub fn bytes_written(&self) -> u64 {
472        self.bytes_written.load(Ordering::Relaxed)
473    }
474
475    /// Returns the number of bytes currently buffered by the layout writers.
476    pub fn buffered_bytes(&self) -> u64 {
477        self.buffered_bytes.buffered_bytes()
478    }
479
480    /// Finish writing the Vortex file, flushing any remaining buffers and returning the
481    /// new file's footer.
482    pub async fn finish(mut self) -> VortexResult<WriteSummary> {
483        // Drop the input channel to signal EOF.
484        drop(self.arrays.take());
485
486        // Await the future task.
487        self.future.await
488    }
489
490    /// Assuming the writer task has failed, await it to get the error.
491    async fn handle_failed_task(&mut self) -> VortexError {
492        match (&mut self.future).await {
493            Ok(_) => vortex_err!(
494                "Internal error: writer task completed successfully but write future finished early"
495            ),
496            Err(e) => e,
497        }
498    }
499}
500
501/// Blocking adapter for [`VortexWriteOptions`].
502pub struct BlockingWrite<'rt, B: BlockingRuntime> {
503    options: VortexWriteOptions,
504    runtime: &'rt B,
505}
506
507impl<'rt, B: BlockingRuntime> BlockingWrite<'rt, B> {
508    /// Write a Vortex file into the given `Write` sink.
509    ///
510    /// The iterator is converted to an [`ArrayStream`] and driven to completion on
511    /// the configured blocking runtime.
512    pub fn write<W: Write + Unpin>(
513        self,
514        write: W,
515        iter: impl ArrayIterator + Send + 'static,
516    ) -> VortexResult<WriteSummary> {
517        self.runtime.block_on(async move {
518            self.options
519                .write(BlockingWriteAdapter(write), iter.into_array_stream())
520                .await
521        })
522    }
523
524    /// Create a blocking push-based writer for chunks with dtype `dtype`.
525    pub fn writer<'w, W: Write + Unpin + 'w>(
526        self,
527        write: W,
528        dtype: DType,
529    ) -> BlockingWriter<'rt, 'w, B> {
530        BlockingWriter {
531            writer: self.options.writer(BlockingWriteAdapter(write), dtype),
532            runtime: self.runtime,
533        }
534    }
535}
536
537/// A blocking adapter around a [`Writer`], allowing incremental writing of arrays to a Vortex file.
538pub struct BlockingWriter<'rt, 'w, B: BlockingRuntime> {
539    runtime: &'rt B,
540    writer: Writer<'w>,
541}
542
543impl<B: BlockingRuntime> BlockingWriter<'_, '_, B> {
544    /// Push one array chunk into the file.
545    pub fn push(&mut self, chunk: ArrayRef) -> VortexResult<()> {
546        self.runtime.block_on(self.writer.push(chunk))
547    }
548
549    /// Returns the number of bytes written to the sink so far.
550    pub fn bytes_written(&self) -> u64 {
551        self.writer.bytes_written()
552    }
553
554    /// Returns the number of bytes currently buffered by layout strategies.
555    pub fn buffered_bytes(&self) -> u64 {
556        self.writer.buffered_bytes()
557    }
558
559    /// Finish writing and return the written file summary.
560    pub fn finish(self) -> VortexResult<WriteSummary> {
561        self.runtime.block_on(self.writer.finish())
562    }
563}
564
565// TODO(ngates): this blocking API may change, for now we just run blocking I/O inline.
566struct BlockingWriteAdapter<W>(W);
567
568impl<W: Write + Unpin> VortexWrite for BlockingWriteAdapter<W> {
569    async fn write_all<B: IoBuf>(&mut self, buffer: B) -> io::Result<B> {
570        self.0.write_all(buffer.as_slice())?;
571        Ok(buffer)
572    }
573
574    fn flush(&mut self) -> impl Future<Output = io::Result<()>> {
575        ready(self.0.flush())
576    }
577
578    fn shutdown(&mut self) -> impl Future<Output = io::Result<()>> {
579        ready(Ok(()))
580    }
581}
582
583/// Summary returned after a Vortex file is written.
584pub struct WriteSummary {
585    footer: Footer,
586    size: u64,
587    // TODO(ngates): add a checksum
588}
589
590impl WriteSummary {
591    /// The footer of the written Vortex file.
592    pub fn footer(&self) -> &Footer {
593        &self.footer
594    }
595
596    /// The total size of the written Vortex file in bytes.
597    pub fn size(&self) -> u64 {
598        self.size
599    }
600
601    /// The total number of rows in the written Vortex file.
602    pub fn row_count(&self) -> u64 {
603        self.footer.row_count()
604    }
605
606    /// Returns the compressed size in bytes of each top-level column in schema order.
607    ///
608    /// A column's size includes every physical segment attributed to its layout subtree,
609    /// including auxiliary segments such as zone maps and dictionaries; see
610    /// [`Footer::compressed_field_sizes`] for the exact attribution semantics and for sizes of
611    /// nested fields. Bytes not attributable to a specific column (e.g. top-level struct
612    /// validity) are not included in any column's size.
613    ///
614    /// For a non-struct file, the returned vector contains a single entry for the root column.
615    pub fn compressed_column_sizes(&self) -> VortexResult<Vec<u64>> {
616        let sizes = self.footer.compressed_field_sizes()?;
617        let Some(fields) = self.footer.dtype().as_struct_fields_opt() else {
618            return Ok(vec![sizes.total()]);
619        };
620        Ok(fields
621            .names()
622            .iter()
623            .map(|name| {
624                sizes
625                    .get(&FieldPath::from_name(name.clone()))
626                    .unwrap_or_default()
627            })
628            .collect())
629    }
630}
631
632#[cfg(test)]
633mod tests {
634    use rstest::rstest;
635    use vortex_array::ArrayContext;
636    use vortex_array::VTable;
637    use vortex_array::array_session;
638    use vortex_array::arrays::Bool;
639    use vortex_array::arrays::Primitive;
640    use vortex_buffer::ByteBuffer;
641    use vortex_edition::Edition;
642    use vortex_edition::EditionDeclaration;
643    use vortex_edition::EditionId;
644    use vortex_edition::EditionSession;
645    use vortex_edition::EditionSessionExt;
646
647    use super::*;
648
649    #[test]
650    fn array_context_only_permits_enabled_encodings() -> Result<(), vortex_edition::EditionError> {
651        const EDITION: EditionId = EditionId::new("test", 2026, 7, 0);
652        static DECLARATION: EditionDeclaration = EditionDeclaration {
653            edition: Edition {
654                id: EDITION,
655                min_vortex_version: None,
656            },
657            added: &[&"vortex.primitive"],
658        };
659
660        let session = array_session().with::<EditionSession>();
661        session.register_edition(&DECLARATION)?;
662        session.enable_edition(EDITION)?;
663
664        let enabled_encoding_ids = session.enabled_encoding_ids();
665        let ctx = ArrayContext::new(enabled_encoding_ids.clone())
666            .with_allowed_ids(enabled_encoding_ids.into_iter().collect());
667        assert_eq!(ctx.to_ids(), [Primitive.id()]);
668        assert!(ctx.intern(&Bool.id()).is_none());
669        Ok(())
670    }
671
672    fn write_options_with_keys(keys: &[String]) -> VortexWriteOptions {
673        array_session().write_options().with_metadata_segments(
674            keys.iter()
675                .map(|key| (key.clone(), ByteBuffer::copy_from(b"value"))),
676        )
677    }
678
679    #[rstest]
680    #[case::empty_key(vec![String::new()], "non-empty")]
681    #[case::oversized_key(vec!["k".repeat(MAX_METADATA_KEY_BYTES + 1)], "keys must be at most")]
682    // The cap is on bytes, not characters.
683    #[case::oversized_multibyte_key(
684        vec!["é".repeat(MAX_METADATA_KEY_BYTES / "é".len() + 1)],
685        "keys must be at most"
686        )]
687    #[case::too_many_segments(
688        (0..=MAX_METADATA_SEGMENTS).map(|idx| format!("key-{idx}")).collect(),
689        "at most 16 metadata segments"
690        )]
691    fn validate_metadata_rejects(#[case] keys: Vec<String>, #[case] expected: &str) {
692        let Err(error) = write_options_with_keys(&keys).validate_metadata() else {
693            panic!("metadata must be rejected for {keys:?}");
694        };
695        assert!(
696            error.to_string().contains(expected),
697            "error should mention {expected:?}, got: {error}"
698        );
699    }
700
701    #[test]
702    fn validate_metadata_accepts_the_limits() -> VortexResult<()> {
703        // Distinct keys, each exactly at the key-length cap.
704        let keys = (0..MAX_METADATA_SEGMENTS)
705            .map(|idx| format!("{idx:0>width$}", width = MAX_METADATA_KEY_BYTES))
706            .collect::<Vec<_>>();
707        write_options_with_keys(&keys).validate_metadata()
708    }
709
710    #[test]
711    fn validate_metadata_accepts_no_metadata() -> VortexResult<()> {
712        write_options_with_keys(&[]).validate_metadata()
713    }
714}