1use 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
64pub 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
81pub trait WriteOptionsSessionExt: SessionExt {
83 fn write_options(&self) -> VortexWriteOptions {
85 VortexWriteOptions::new(self.session())
86 }
87}
88impl<S: SessionExt> WriteOptionsSessionExt for S {}
89
90impl VortexWriteOptions {
91 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 pub fn with_strategy(mut self, strategy: Arc<dyn LayoutStrategy>) -> Self {
113 self.strategy = strategy;
114 self
115 }
116
117 pub fn buffered_bytes_tracker(&self) -> BufferedBytesTracker {
124 self.buffered_bytes.clone()
125 }
126
127 pub fn exclude_dtype(mut self) -> Self {
131 self.exclude_dtype = true;
132 self
133 }
134
135 pub fn with_file_statistics(mut self, file_statistics: Vec<Stat>) -> Self {
139 self.file_statistics = file_statistics;
140 self
141 }
142
143 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 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 pub fn validate_metadata(&self) -> VortexResult<()> {
177 validate_metadata_segments(&self.metadata)
178 }
179}
180
181impl VortexWriteOptions {
182 pub fn blocking<B: BlockingRuntime>(self, runtime: &B) -> BlockingWrite<'_, B> {
187 BlockingWrite {
188 options: self,
189 runtime,
190 }
191 }
192
193 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 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 write.write_all(ByteBuffer::copy_from(MAGIC_BYTES)).await?;
240 let mut position = MAGIC_BYTES.len() as u64;
241
242 let (send, recv) = kanal::bounded_async(1);
244
245 let segments = Arc::new(BufferedSegmentSink::new(send, position));
246
247 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 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 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 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 pub fn writer<'w, W: VortexWrite + Unpin + 'w>(self, write: W, dtype: DType) -> Writer<'w> {
325 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 let enabled_encoding_ids = session.enabled_encoding_ids();
352 ArrayContext::new(enabled_encoding_ids.iter().cloned().sorted().collect())
353 .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
389pub struct Writer<'w> {
391 arrays: Option<kanal::AsyncSender<VortexResult<ArrayRef>>>,
393 future: Fuse<LocalBoxFuture<'w, VortexResult<WriteSummary>>>,
395 bytes_written: Arc<AtomicU64>,
397 buffered_bytes: BufferedBytesTracker,
399}
400
401impl Writer<'_> {
402 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 select! {
411 result = send_fut => {
412 if result.is_err() {
414 return Err(self.handle_failed_task().await);
415 }
416 },
417 result = &mut self.future => {
418 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 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 select! {
449 result = stream_fut => {
450 if let Err(_send_err) = result {
451 return Err(self.handle_failed_task().await);
453 }
454 }
455
456 result = &mut self.future => {
457 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 pub fn bytes_written(&self) -> u64 {
472 self.bytes_written.load(Ordering::Relaxed)
473 }
474
475 pub fn buffered_bytes(&self) -> u64 {
477 self.buffered_bytes.buffered_bytes()
478 }
479
480 pub async fn finish(mut self) -> VortexResult<WriteSummary> {
483 drop(self.arrays.take());
485
486 self.future.await
488 }
489
490 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
501pub struct BlockingWrite<'rt, B: BlockingRuntime> {
503 options: VortexWriteOptions,
504 runtime: &'rt B,
505}
506
507impl<'rt, B: BlockingRuntime> BlockingWrite<'rt, B> {
508 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 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
537pub struct BlockingWriter<'rt, 'w, B: BlockingRuntime> {
539 runtime: &'rt B,
540 writer: Writer<'w>,
541}
542
543impl<B: BlockingRuntime> BlockingWriter<'_, '_, B> {
544 pub fn push(&mut self, chunk: ArrayRef) -> VortexResult<()> {
546 self.runtime.block_on(self.writer.push(chunk))
547 }
548
549 pub fn bytes_written(&self) -> u64 {
551 self.writer.bytes_written()
552 }
553
554 pub fn buffered_bytes(&self) -> u64 {
556 self.writer.buffered_bytes()
557 }
558
559 pub fn finish(self) -> VortexResult<WriteSummary> {
561 self.runtime.block_on(self.writer.finish())
562 }
563}
564
565struct 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
583pub struct WriteSummary {
585 footer: Footer,
586 size: u64,
587 }
589
590impl WriteSummary {
591 pub fn footer(&self) -> &Footer {
593 &self.footer
594 }
595
596 pub fn size(&self) -> u64 {
598 self.size
599 }
600
601 pub fn row_count(&self) -> u64 {
603 self.footer.row_count()
604 }
605
606 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 #[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 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}