1use std::sync::Arc;
16
17use async_trait::async_trait;
18use vortex_array::dtype::Field;
19use vortex_array::dtype::FieldName;
20use vortex_array::dtype::FieldPath;
21use vortex_error::VortexResult;
22use vortex_session::VortexSession;
23use vortex_utils::aliases::hash_map::HashMap;
24use vortex_utils::aliases::hash_set::HashSet;
25
26use crate::LayoutRef;
27use crate::LayoutStrategy;
28use crate::LayoutWriterContext;
29use crate::layouts::list::writer::ListLayoutStrategy;
30use crate::layouts::struct_::StructStrategy;
31use crate::segments::SegmentSinkRef;
32use crate::sequence::SendableSequentialStream;
33use crate::sequence::SequencePointer;
34
35type ListLayoutFactory = Arc<dyn Fn(ListLayoutStrategy) -> Arc<dyn LayoutStrategy> + Send + Sync>;
36
37pub struct TableStrategy {
52 leaf_writers: HashMap<FieldPath, Arc<dyn LayoutStrategy>>,
55 validity: Arc<dyn LayoutStrategy>,
57 leaf: Arc<dyn LayoutStrategy>,
59 list_layout_factory: Option<ListLayoutFactory>,
64}
65
66impl TableStrategy {
67 pub fn new(validity: Arc<dyn LayoutStrategy>, fallback: Arc<dyn LayoutStrategy>) -> Self {
86 Self {
87 leaf_writers: Default::default(),
88 validity,
89 leaf: fallback,
90 list_layout_factory: None,
91 }
92 }
93
94 pub fn with_field_writer(
123 mut self,
124 field_path: impl Into<FieldPath>,
125 writer: Arc<dyn LayoutStrategy>,
126 ) -> Self {
127 self.leaf_writers
128 .insert(self.validate_path(field_path.into()), writer);
129 self
130 }
131
132 pub fn with_field_writers(
136 mut self,
137 writers: impl IntoIterator<Item = (FieldPath, Arc<dyn LayoutStrategy>)>,
138 ) -> Self {
139 for (field_path, strategy) in writers {
140 self.leaf_writers
141 .insert(self.validate_path(field_path), strategy);
142 }
143 self
144 }
145
146 pub fn with_default_strategy(mut self, default: Arc<dyn LayoutStrategy>) -> Self {
148 self.leaf = default;
149 self
150 }
151
152 pub fn with_validity_strategy(mut self, validity: Arc<dyn LayoutStrategy>) -> Self {
154 self.validity = validity;
155 self
156 }
157
158 pub fn with_list_layout(self) -> Self {
163 self.with_list_layout_factory(|strategy| Arc::new(strategy))
164 }
165
166 pub fn with_list_layout_factory(
172 mut self,
173 factory: impl Fn(ListLayoutStrategy) -> Arc<dyn LayoutStrategy> + Send + Sync + 'static,
174 ) -> Self {
175 self.list_layout_factory = Some(Arc::new(factory));
176 self
177 }
178}
179
180impl TableStrategy {
181 fn struct_strategy(&self) -> StructStrategy {
186 let mut field_writers: HashMap<FieldName, Arc<dyn LayoutStrategy>> = HashMap::default();
187
188 let mut named_first: HashSet<FieldName> = HashSet::default();
191 for path in self.leaf_writers.keys() {
192 if let Some(Field::Name(name)) = path.parts().first() {
193 named_first.insert(name.clone());
194 }
195 }
196
197 for name in named_first {
198 let writer = match self.leaf_writers.get(&FieldPath::from_name(name.clone())) {
201 Some(exact) => Arc::clone(exact),
202 None => {
203 Arc::new(self.descend(&Field::Name(name.clone()))) as Arc<dyn LayoutStrategy>
204 }
205 };
206 field_writers.insert(name, writer);
207 }
208
209 StructStrategy::new(Arc::clone(&self.validity), Arc::new(self.descend_clean()))
210 .with_field_writers(field_writers)
211 }
212
213 fn list_strategy(&self) -> Option<Arc<dyn LayoutStrategy>> {
219 let factory = self.list_layout_factory.as_ref()?;
220 let list_layout = ListLayoutStrategy::default()
221 .with_elements(Arc::new(self.descend_clean()))
222 .with_offsets(Arc::clone(&self.leaf))
223 .with_validity(Arc::clone(&self.validity))
224 .with_fallback(Arc::clone(&self.leaf));
225 Some(factory(list_layout))
226 }
227
228 fn descend(&self, field: &Field) -> Self {
231 let mut new_writers = HashMap::with_capacity(self.leaf_writers.len());
232
233 for (field_path, strategy) in &self.leaf_writers {
234 if field_path.parts().first() == Some(field)
235 && let Some(subpath) = field_path.clone().step_into()
236 && !subpath.is_root()
237 {
238 new_writers.insert(subpath, Arc::clone(strategy));
239 }
240 }
241
242 Self {
243 leaf_writers: new_writers,
244 validity: Arc::clone(&self.validity),
245 leaf: Arc::clone(&self.leaf),
246 list_layout_factory: self.list_layout_factory.clone(),
247 }
248 }
249
250 fn descend_clean(&self) -> Self {
253 Self {
254 leaf_writers: HashMap::default(),
255 validity: Arc::clone(&self.validity),
256 leaf: Arc::clone(&self.leaf),
257 list_layout_factory: self.list_layout_factory.clone(),
258 }
259 }
260
261 fn validate_path(&self, path: FieldPath) -> FieldPath {
262 assert!(
263 !path.is_root(),
264 "Do not set override as a root strategy, instead set the default strategy"
265 );
266
267 for field_path in self.leaf_writers.keys() {
270 assert!(
271 !path.overlap(field_path),
272 "Override for field_path {path} conflicts with existing override for {field_path}"
273 );
274 }
275
276 path
277 }
278}
279
280#[async_trait]
282impl LayoutStrategy for TableStrategy {
283 async fn write_stream(
284 &self,
285 ctx: LayoutWriterContext,
286 segment_sink: SegmentSinkRef,
287 stream: SendableSequentialStream,
288 eof: SequencePointer,
289 session: &VortexSession,
290 ) -> VortexResult<LayoutRef> {
291 let dtype = stream.dtype().clone();
292
293 if dtype.is_struct() {
294 return self
295 .struct_strategy()
296 .write_stream(ctx, segment_sink, stream, eof, session)
297 .await;
298 }
299
300 if dtype.is_list()
301 && let Some(list_strategy) = self.list_strategy()
302 {
303 return list_strategy
304 .write_stream(ctx, segment_sink, stream, eof, session)
305 .await;
306 }
307
308 self.leaf
310 .write_stream(ctx, segment_sink, stream, eof, session)
311 .await
312 }
313}
314
315#[cfg(test)]
316mod tests {
317 use std::num::NonZeroUsize;
318 use std::sync::Arc;
319 use std::task::Poll;
320
321 use vortex_array::ArrayContext;
322 use vortex_array::ArrayRef;
323 use vortex_array::IntoArray;
324 use vortex_array::VortexSessionExecute;
325 use vortex_array::array_session;
326 use vortex_array::arrays::BoolArray;
327 use vortex_array::arrays::ChunkedArray;
328 use vortex_array::arrays::ListArray;
329 use vortex_array::arrays::PrimitiveArray;
330 use vortex_array::arrays::StructArray;
331 use vortex_array::dtype::DType;
332 use vortex_array::dtype::FieldPath;
333 use vortex_array::dtype::Nullability;
334 use vortex_array::dtype::PType;
335 use vortex_array::dtype::StructFields;
336 use vortex_array::field_path;
337 use vortex_array::validity::Validity;
338 use vortex_buffer::buffer;
339 use vortex_error::VortexExpect;
340 use vortex_error::VortexResult;
341 use vortex_io::session::RuntimeSessionExt;
342
343 use crate::LayoutRef;
344 use crate::LayoutStrategy;
345 use crate::layouts::chunked::Chunked;
346 use crate::layouts::chunked::writer::ChunkedLayoutStrategy;
347 use crate::layouts::flat::Flat;
348 use crate::layouts::flat::writer::FlatLayoutStrategy;
349 use crate::layouts::list::List;
350 use crate::layouts::repartition::RepartitionStrategy;
351 use crate::layouts::repartition::RepartitionWriterOptions;
352 use crate::layouts::table::TableStrategy;
353 use crate::layouts::zoned::Zoned;
354 use crate::layouts::zoned::writer::ZonedLayoutOptions;
355 use crate::layouts::zoned::writer::ZonedStrategy;
356 use crate::segments::TestSegments;
357 use crate::sequence::SequenceId;
358 use crate::sequence::SequentialArrayStreamExt;
359 use crate::sequence::SequentialStreamAdapter;
360 use crate::sequence::SequentialStreamExt;
361 use crate::test::SESSION;
362 use crate::test::new_session;
363
364 async fn write<S: LayoutStrategy>(strategy: &S, array: ArrayRef) -> VortexResult<LayoutRef> {
365 let segments = Arc::new(TestSegments::default());
366 let (ptr, eof) = SequenceId::root().split();
367 let stream = array.to_array_stream().sequenced(ptr);
368 let session = new_session().with_tokio();
369 strategy
370 .write_stream(
371 ArrayContext::empty().into(),
372 segments,
373 stream,
374 eof,
375 &session,
376 )
377 .await
378 }
379
380 fn flat_table() -> TableStrategy {
383 let flat: Arc<dyn LayoutStrategy> = Arc::new(FlatLayoutStrategy::default());
384 TableStrategy::new(Arc::clone(&flat), flat)
385 }
386
387 #[tokio::test]
389 async fn dispatches_struct() -> VortexResult<()> {
390 let struct_array = StructArray::from_fields(
391 [
392 ("a", buffer![1i32, 2, 3].into_array()),
393 ("b", buffer![10i32, 20, 30].into_array()),
394 ]
395 .as_slice(),
396 )?
397 .into_array();
398
399 let layout = write(&flat_table(), struct_array).await?;
400 insta::assert_snapshot!(layout.display_tree(), @r"
401 vortex.struct, dtype: {a=i32, b=i32}, children: 2
402 ├── a: vortex.flat, dtype: i32, segment: 0
403 └── b: vortex.flat, dtype: i32, segment: 1
404 ");
405 Ok(())
406 }
407
408 #[tokio::test]
411 async fn dispatches_nested_list() -> VortexResult<()> {
412 let mut ctx = array_session().create_execution_ctx();
413 let inner = ListArray::try_new(
414 buffer![1i32, 2, 3, 4, 5, 6].into_array(),
415 buffer![0u32, 2, 5, 5, 6].into_array(),
416 Validity::NonNullable,
417 &mut ctx,
418 )?
419 .into_array();
420 let outer = ListArray::try_new(
421 inner,
422 buffer![0u32, 2, 4].into_array(),
423 Validity::NonNullable,
424 &mut ctx,
425 )?
426 .into_array();
427
428 let layout = write(&flat_table().with_list_layout(), outer).await?;
429 insta::assert_snapshot!(layout.display_tree(), @r"
430 vortex.list, dtype: list(list(i32)), children: 2
431 ├── elements: vortex.list, dtype: list(i32), children: 2
432 │ ├── elements: vortex.flat, dtype: i32, segment: 1
433 │ └── offsets: vortex.flat, dtype: u64, segment: 2
434 └── offsets: vortex.flat, dtype: u64, segment: 0
435 ");
436 Ok(())
437 }
438
439 #[tokio::test]
442 async fn dispatches_struct_list_struct() -> VortexResult<()> {
443 let mut ctx = array_session().create_execution_ctx();
444 let inner_struct = StructArray::from_fields(
445 [
446 ("a", buffer![1i32, 2, 3, 4, 5].into_array()),
447 ("b", buffer![10i32, 20, 30, 40, 50].into_array()),
448 ]
449 .as_slice(),
450 )?
451 .into_array();
452 let items = ListArray::try_new(
453 inner_struct,
454 buffer![0u32, 2, 5, 5].into_array(),
455 Validity::Array(BoolArray::from_iter([true, false, true]).into_array()),
456 &mut ctx,
457 )?
458 .into_array();
459 let st = StructArray::from_fields([("items", items)].as_slice())?.into_array();
460
461 let layout = write(&flat_table().with_list_layout(), st).await?;
462 insta::assert_snapshot!(layout.display_tree(), @r"
463 vortex.struct, dtype: {items=list({a=i32, b=i32})?}, children: 1
464 └── items: vortex.list, dtype: list({a=i32, b=i32})?, children: 3
465 ├── elements: vortex.struct, dtype: {a=i32, b=i32}, children: 2
466 │ ├── a: vortex.flat, dtype: i32, segment: 2
467 │ └── b: vortex.flat, dtype: i32, segment: 3
468 ├── offsets: vortex.flat, dtype: u64, segment: 0
469 └── validity: vortex.flat, dtype: bool, segment: 1
470 ");
471 Ok(())
472 }
473
474 #[tokio::test]
478 async fn dispatches_chunked_list() -> VortexResult<()> {
479 let mut ctx = array_session().create_execution_ctx();
480 let chunk0 = ListArray::try_new(
481 buffer![1i32, 2, 3].into_array(),
482 buffer![0u32, 2, 3].into_array(),
483 Validity::NonNullable,
484 &mut ctx,
485 )?
486 .into_array();
487 let chunk1 = ListArray::try_new(
488 buffer![4i32, 5, 6, 7].into_array(),
489 buffer![0u32, 1, 4].into_array(),
490 Validity::NonNullable,
491 &mut ctx,
492 )?
493 .into_array();
494 let dtype = chunk0.dtype().clone();
495 let chunked = ChunkedArray::try_new(vec![chunk0, chunk1], dtype)?.into_array();
496
497 let flat: Arc<dyn LayoutStrategy> = Arc::new(FlatLayoutStrategy::default());
498 let dispatcher = TableStrategy::new(
499 Arc::clone(&flat),
500 Arc::new(ChunkedLayoutStrategy::new(FlatLayoutStrategy::default())),
501 )
502 .with_list_layout();
503 let layout = write(&dispatcher, chunked).await?;
504 insta::assert_snapshot!(layout.display_tree(), @r"
505 vortex.list, dtype: list(i32), children: 2
506 ├── elements: vortex.chunked, dtype: i32, children: 2
507 │ ├── [0]: vortex.flat, dtype: i32, segment: 0
508 │ └── [1]: vortex.flat, dtype: i32, segment: 1
509 └── offsets: vortex.chunked, dtype: u64, children: 2
510 ├── [0]: vortex.flat, dtype: u64, segment: 2
511 └── [1]: vortex.flat, dtype: u64, segment: 3
512 ");
513 Ok(())
514 }
515
516 #[tokio::test]
518 async fn wraps_list_strategy_before_decomposition() -> VortexResult<()> {
519 let mut ctx = array_session().create_execution_ctx();
520 let list = ListArray::try_new(
521 PrimitiveArray::from_iter(0..9_i32).into_array(),
522 PrimitiveArray::from_iter(0..=9_u32).into_array(),
523 Validity::NonNullable,
524 &mut ctx,
525 )?
526 .into_array();
527
528 let row_block_size = NonZeroUsize::new(4).vortex_expect("4 is non-zero");
529 let flat: Arc<dyn LayoutStrategy> = Arc::new(FlatLayoutStrategy::default());
530 let stats = Arc::clone(&flat);
531 let chunked: Arc<dyn LayoutStrategy> =
532 Arc::new(ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()));
533 let dispatcher = TableStrategy::new(Arc::clone(&flat), chunked).with_list_layout_factory(
534 move |list_layout| {
535 let zoned = ZonedStrategy::new(
536 list_layout,
537 Arc::clone(&stats),
538 ZonedLayoutOptions {
539 block_size: row_block_size,
540 ..Default::default()
541 },
542 );
543 Arc::new(RepartitionStrategy::new(
544 zoned,
545 RepartitionWriterOptions {
546 block_size_minimum: 0,
547 block_len_multiple: row_block_size.get(),
548 block_size_target: None,
549 canonicalize: false,
550 },
551 )) as Arc<dyn LayoutStrategy>
552 },
553 );
554
555 let layout = write(&dispatcher, list).await?;
556 let zoned = layout.as_::<Zoned>();
557 assert_eq!(zoned.zone_len(), 4);
558 assert_eq!(zoned.nzones(), 3);
559
560 let data = layout
561 .slot(0)?
562 .vortex_expect("ZonedLayout always has a data child");
563 assert!(data.is::<List>());
564 assert_eq!(data.row_count(), 9);
565 Ok(())
566 }
567
568 #[tokio::test]
570 async fn non_struct_input_uses_leaf() -> VortexResult<()> {
571 let primitive = PrimitiveArray::from_iter([1i32, 2, 3]).into_array();
572 let layout = write(&flat_table(), primitive).await?;
573 insta::assert_snapshot!(layout.display_tree(), @"vortex.flat, dtype: i32, segment: 0");
574 Ok(())
575 }
576
577 #[tokio::test]
579 async fn chunked_struct() -> VortexResult<()> {
580 let validity: Arc<dyn LayoutStrategy> = Arc::new(FlatLayoutStrategy::default());
581 let chunked_flat: Arc<dyn LayoutStrategy> =
582 Arc::new(ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()));
583 let dispatcher = TableStrategy::new(validity, chunked_flat);
584
585 let c0 = StructArray::from_fields(
586 [
587 ("a", buffer![1i32, 2].into_array()),
588 ("b", buffer![10i32, 20].into_array()),
589 ]
590 .as_slice(),
591 )?
592 .into_array();
593 let c1 = StructArray::from_fields(
594 [
595 ("a", buffer![3i32].into_array()),
596 ("b", buffer![30i32].into_array()),
597 ]
598 .as_slice(),
599 )?
600 .into_array();
601 let dtype = c0.dtype().clone();
602 let chunked = ChunkedArray::try_new(vec![c0, c1], dtype)?.into_array();
603
604 let layout = write(&dispatcher, chunked).await?;
605 insta::assert_snapshot!(layout.display_tree(), @r"
606 vortex.struct, dtype: {a=i32, b=i32}, children: 2
607 ├── a: vortex.chunked, dtype: i32, children: 2
608 │ ├── [0]: vortex.flat, dtype: i32, segment: 0
609 │ └── [1]: vortex.flat, dtype: i32, segment: 1
610 └── b: vortex.chunked, dtype: i32, children: 2
611 ├── [0]: vortex.flat, dtype: i32, segment: 2
612 └── [1]: vortex.flat, dtype: i32, segment: 3
613 ");
614 Ok(())
615 }
616
617 #[tokio::test]
619 async fn field_override_is_used() -> VortexResult<()> {
620 let struct_array = StructArray::from_fields(
621 [
622 ("a", buffer![1i32, 2, 3].into_array()),
623 ("b", buffer![10i32, 20, 30].into_array()),
624 ]
625 .as_slice(),
626 )?
627 .into_array();
628
629 let strategy = flat_table().with_field_writer(
630 field_path!(a),
631 Arc::new(RepartitionStrategy::new(
632 ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()),
633 RepartitionWriterOptions {
634 block_size_minimum: 0,
635 block_len_multiple: 2,
636 block_size_target: None,
637 canonicalize: false,
638 },
639 )),
640 );
641 let layout = write(&strategy, struct_array).await?;
642 assert_eq!(layout.nchildren(), 2);
643 let overridden = layout.slot(1)?.vortex_expect("field a");
644 assert!(overridden.is::<Chunked>());
645 assert_eq!(overridden.nchildren(), 2);
646 assert_eq!(
647 overridden.slot(0)?.vortex_expect("first chunk").row_count(),
648 2
649 );
650 assert_eq!(
651 overridden
652 .slot(1)?
653 .vortex_expect("second chunk")
654 .row_count(),
655 1
656 );
657 assert!(layout.slot(2)?.vortex_expect("field b").is::<Flat>());
658 Ok(())
659 }
660
661 #[test]
662 #[should_panic(
663 expected = "Override for field_path $a.$b conflicts with existing override for $a.$b.$c"
664 )]
665 fn test_overlapping_paths_fail() {
666 let flat = Arc::new(FlatLayoutStrategy::default());
667
668 let path = TableStrategy::new(
670 Arc::<FlatLayoutStrategy>::clone(&flat),
671 Arc::<FlatLayoutStrategy>::clone(&flat),
672 )
673 .with_field_writer(field_path!(a.b.c), Arc::<FlatLayoutStrategy>::clone(&flat));
674
675 let _path = path.with_field_writer(field_path!(a.b), flat);
677 }
678
679 #[test]
680 #[should_panic(
681 expected = "Do not set override as a root strategy, instead set the default strategy"
682 )]
683 fn test_root_override() {
684 let flat = Arc::new(FlatLayoutStrategy::default());
685 let _strategy = TableStrategy::new(
686 Arc::<FlatLayoutStrategy>::clone(&flat),
687 Arc::<FlatLayoutStrategy>::clone(&flat),
688 )
689 .with_field_writer(FieldPath::root(), flat);
690 }
691
692 #[tokio::test]
693 #[should_panic(expected = "panic while transposing table stream")]
694 async fn table_fanout_panic_propagates() {
695 let ctx = ArrayContext::empty();
696 let segments = Arc::new(TestSegments::default());
697 let (_, eof) = SequenceId::root().split();
698 let dtype = DType::Struct(
699 StructFields::from_iter([(
700 "a",
701 DType::Primitive(PType::I32, Nullability::NonNullable),
702 )]),
703 Nullability::NonNullable,
704 );
705 let stream =
706 futures::stream::poll_fn(|_| -> Poll<Option<VortexResult<(SequenceId, ArrayRef)>>> {
707 panic!("panic while transposing table stream");
708 });
709 let strategy = TableStrategy::new(
710 Arc::new(FlatLayoutStrategy::default()),
711 Arc::new(FlatLayoutStrategy::default()),
712 );
713
714 strategy
715 .write_stream(
716 ctx.into(),
717 segments,
718 SequentialStreamAdapter::new(dtype, stream).sendable(),
719 eof,
720 &SESSION,
721 )
722 .await
723 .unwrap();
724 }
725}