1use std::sync::Arc;
15
16use async_trait::async_trait;
17use vortex_array::ArrayContext;
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::layouts::struct_::StructStrategy;
29use crate::segments::SegmentSinkRef;
30use crate::sequence::SendableSequentialStream;
31use crate::sequence::SequencePointer;
32
33pub struct TableStrategy {
43 leaf_writers: HashMap<FieldPath, Arc<dyn LayoutStrategy>>,
46 validity: Arc<dyn LayoutStrategy>,
48 leaf: Arc<dyn LayoutStrategy>,
50}
51
52impl TableStrategy {
53 pub fn new(validity: Arc<dyn LayoutStrategy>, fallback: Arc<dyn LayoutStrategy>) -> Self {
72 Self {
73 leaf_writers: Default::default(),
74 validity,
75 leaf: fallback,
76 }
77 }
78
79 pub fn with_field_writer(
108 mut self,
109 field_path: impl Into<FieldPath>,
110 writer: Arc<dyn LayoutStrategy>,
111 ) -> Self {
112 self.leaf_writers
113 .insert(self.validate_path(field_path.into()), writer);
114 self
115 }
116
117 pub fn with_field_writers(
121 mut self,
122 writers: impl IntoIterator<Item = (FieldPath, Arc<dyn LayoutStrategy>)>,
123 ) -> Self {
124 for (field_path, strategy) in writers {
125 self.leaf_writers
126 .insert(self.validate_path(field_path), strategy);
127 }
128 self
129 }
130
131 pub fn with_default_strategy(mut self, default: Arc<dyn LayoutStrategy>) -> Self {
133 self.leaf = default;
134 self
135 }
136
137 pub fn with_validity_strategy(mut self, validity: Arc<dyn LayoutStrategy>) -> Self {
139 self.validity = validity;
140 self
141 }
142}
143
144impl TableStrategy {
145 fn struct_strategy(&self) -> StructStrategy {
150 let mut field_writers: HashMap<FieldName, Arc<dyn LayoutStrategy>> = HashMap::default();
151
152 let mut named_first: HashSet<FieldName> = HashSet::default();
155 for path in self.leaf_writers.keys() {
156 if let Some(Field::Name(name)) = path.parts().first() {
157 named_first.insert(name.clone());
158 }
159 }
160
161 for name in named_first {
162 let writer = match self.leaf_writers.get(&FieldPath::from_name(name.clone())) {
165 Some(exact) => Arc::clone(exact),
166 None => {
167 Arc::new(self.descend(&Field::Name(name.clone()))) as Arc<dyn LayoutStrategy>
168 }
169 };
170 field_writers.insert(name, writer);
171 }
172
173 StructStrategy::new(Arc::clone(&self.validity), Arc::new(self.descend_clean()))
174 .with_field_writers(field_writers)
175 }
176
177 fn descend(&self, field: &Field) -> Self {
180 let mut new_writers = HashMap::with_capacity(self.leaf_writers.len());
181
182 for (field_path, strategy) in &self.leaf_writers {
183 if field_path.parts().first() == Some(field)
184 && let Some(subpath) = field_path.clone().step_into()
185 && !subpath.is_root()
186 {
187 new_writers.insert(subpath, Arc::clone(strategy));
188 }
189 }
190
191 Self {
192 leaf_writers: new_writers,
193 validity: Arc::clone(&self.validity),
194 leaf: Arc::clone(&self.leaf),
195 }
196 }
197
198 fn descend_clean(&self) -> Self {
201 Self {
202 leaf_writers: HashMap::default(),
203 validity: Arc::clone(&self.validity),
204 leaf: Arc::clone(&self.leaf),
205 }
206 }
207
208 fn validate_path(&self, path: FieldPath) -> FieldPath {
209 assert!(
210 !path.is_root(),
211 "Do not set override as a root strategy, instead set the default strategy"
212 );
213
214 for field_path in self.leaf_writers.keys() {
217 assert!(
218 !path.overlap(field_path),
219 "Override for field_path {path} conflicts with existing override for {field_path}"
220 );
221 }
222
223 path
224 }
225}
226
227#[async_trait]
229impl LayoutStrategy for TableStrategy {
230 async fn write_stream(
231 &self,
232 ctx: ArrayContext,
233 segment_sink: SegmentSinkRef,
234 stream: SendableSequentialStream,
235 eof: SequencePointer,
236 session: &VortexSession,
237 ) -> VortexResult<LayoutRef> {
238 let dtype = stream.dtype().clone();
239
240 if dtype.is_struct() {
241 return self
242 .struct_strategy()
243 .write_stream(ctx, segment_sink, stream, eof, session)
244 .await;
245 }
246
247 self.leaf
249 .write_stream(ctx, segment_sink, stream, eof, session)
250 .await
251 }
252}
253
254#[cfg(test)]
255mod tests {
256 use std::sync::Arc;
257 use std::task::Poll;
258
259 use vortex_array::ArrayContext;
260 use vortex_array::ArrayRef;
261 use vortex_array::IntoArray;
262 use vortex_array::arrays::ChunkedArray;
263 use vortex_array::arrays::PrimitiveArray;
264 use vortex_array::arrays::StructArray;
265 use vortex_array::dtype::DType;
266 use vortex_array::dtype::FieldPath;
267 use vortex_array::dtype::Nullability;
268 use vortex_array::dtype::PType;
269 use vortex_array::dtype::StructFields;
270 use vortex_array::field_path;
271 use vortex_buffer::buffer;
272 use vortex_error::VortexResult;
273 use vortex_io::runtime::single::block_on;
274 use vortex_io::session::RuntimeSessionExt;
275
276 use crate::LayoutRef;
277 use crate::LayoutStrategy;
278 use crate::layouts::chunked::writer::ChunkedLayoutStrategy;
279 use crate::layouts::flat::writer::FlatLayoutStrategy;
280 use crate::layouts::table::TableStrategy;
281 use crate::segments::TestSegments;
282 use crate::sequence::SequenceId;
283 use crate::sequence::SequentialArrayStreamExt;
284 use crate::sequence::SequentialStreamAdapter;
285 use crate::sequence::SequentialStreamExt;
286 use crate::test::SESSION;
287
288 async fn write<S: LayoutStrategy>(strategy: &S, array: ArrayRef) -> VortexResult<LayoutRef> {
289 let segments = Arc::new(TestSegments::default());
290 let (ptr, eof) = SequenceId::root().split();
291 let stream = array.to_array_stream().sequenced(ptr);
292 strategy
293 .write_stream(ArrayContext::empty(), segments, stream, eof, &SESSION)
294 .await
295 }
296
297 fn flat_table() -> TableStrategy {
300 let flat: Arc<dyn LayoutStrategy> = Arc::new(FlatLayoutStrategy::default());
301 TableStrategy::new(Arc::clone(&flat), flat)
302 }
303
304 #[tokio::test]
306 async fn dispatches_struct() -> VortexResult<()> {
307 let struct_array = StructArray::from_fields(
308 [
309 ("a", buffer![1i32, 2, 3].into_array()),
310 ("b", buffer![10i32, 20, 30].into_array()),
311 ]
312 .as_slice(),
313 )?
314 .into_array();
315
316 let layout = write(&flat_table(), struct_array).await?;
317 insta::assert_snapshot!(layout.display_tree(), @r"
318 vortex.struct, dtype: {a=i32, b=i32}, children: 2
319 ├── a: vortex.flat, dtype: i32, segment: 0
320 └── b: vortex.flat, dtype: i32, segment: 1
321 ");
322 Ok(())
323 }
324
325 #[tokio::test]
327 async fn non_struct_input_uses_leaf() -> VortexResult<()> {
328 let primitive = PrimitiveArray::from_iter([1i32, 2, 3]).into_array();
329 let layout = write(&flat_table(), primitive).await?;
330 insta::assert_snapshot!(layout.display_tree(), @"vortex.flat, dtype: i32, segment: 0");
331 Ok(())
332 }
333
334 #[tokio::test]
336 async fn chunked_struct() -> VortexResult<()> {
337 let validity: Arc<dyn LayoutStrategy> = Arc::new(FlatLayoutStrategy::default());
338 let chunked_flat: Arc<dyn LayoutStrategy> =
339 Arc::new(ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()));
340 let dispatcher = TableStrategy::new(validity, chunked_flat);
341
342 let c0 = StructArray::from_fields(
343 [
344 ("a", buffer![1i32, 2].into_array()),
345 ("b", buffer![10i32, 20].into_array()),
346 ]
347 .as_slice(),
348 )?
349 .into_array();
350 let c1 = StructArray::from_fields(
351 [
352 ("a", buffer![3i32].into_array()),
353 ("b", buffer![30i32].into_array()),
354 ]
355 .as_slice(),
356 )?
357 .into_array();
358 let dtype = c0.dtype().clone();
359 let chunked = ChunkedArray::try_new(vec![c0, c1], dtype)?.into_array();
360
361 let layout = write(&dispatcher, chunked).await?;
362 insta::assert_snapshot!(layout.display_tree(), @r"
363 vortex.struct, dtype: {a=i32, b=i32}, children: 2
364 ├── a: vortex.chunked, dtype: i32, children: 2
365 │ ├── [0]: vortex.flat, dtype: i32, segment: 0
366 │ └── [1]: vortex.flat, dtype: i32, segment: 1
367 └── b: vortex.chunked, dtype: i32, children: 2
368 ├── [0]: vortex.flat, dtype: i32, segment: 2
369 └── [1]: vortex.flat, dtype: i32, segment: 3
370 ");
371 Ok(())
372 }
373
374 #[tokio::test]
376 async fn field_override_is_used() -> VortexResult<()> {
377 let struct_array = StructArray::from_fields(
378 [
379 ("a", buffer![1i32, 2, 3].into_array()),
380 ("b", buffer![10i32, 20, 30].into_array()),
381 ]
382 .as_slice(),
383 )?
384 .into_array();
385
386 let strategy =
387 flat_table().with_field_writer(field_path!(a), Arc::new(FlatLayoutStrategy::default()));
388 let layout = write(&strategy, struct_array).await?;
389 insta::assert_snapshot!(layout.display_tree(), @r"
390 vortex.struct, dtype: {a=i32, b=i32}, children: 2
391 ├── a: vortex.flat, dtype: i32, segment: 0
392 └── b: vortex.flat, dtype: i32, segment: 1
393 ");
394 Ok(())
395 }
396
397 #[test]
398 #[should_panic(
399 expected = "Override for field_path $a.$b conflicts with existing override for $a.$b.$c"
400 )]
401 fn test_overlapping_paths_fail() {
402 let flat = Arc::new(FlatLayoutStrategy::default());
403
404 let path = TableStrategy::new(
406 Arc::<FlatLayoutStrategy>::clone(&flat),
407 Arc::<FlatLayoutStrategy>::clone(&flat),
408 )
409 .with_field_writer(field_path!(a.b.c), Arc::<FlatLayoutStrategy>::clone(&flat));
410
411 let _path = path.with_field_writer(field_path!(a.b), flat);
413 }
414
415 #[test]
416 #[should_panic(
417 expected = "Do not set override as a root strategy, instead set the default strategy"
418 )]
419 fn test_root_override() {
420 let flat = Arc::new(FlatLayoutStrategy::default());
421 let _strategy = TableStrategy::new(
422 Arc::<FlatLayoutStrategy>::clone(&flat),
423 Arc::<FlatLayoutStrategy>::clone(&flat),
424 )
425 .with_field_writer(FieldPath::root(), flat);
426 }
427
428 #[test]
429 #[should_panic(expected = "panic while transposing table stream")]
430 fn table_fanout_panic_propagates() {
431 let ctx = ArrayContext::empty();
432 let segments = Arc::new(TestSegments::default());
433 let (_, eof) = SequenceId::root().split();
434 let dtype = DType::Struct(
435 StructFields::from_iter([(
436 "a",
437 DType::Primitive(PType::I32, Nullability::NonNullable),
438 )]),
439 Nullability::NonNullable,
440 );
441 let stream =
442 futures::stream::poll_fn(|_| -> Poll<Option<VortexResult<(SequenceId, ArrayRef)>>> {
443 panic!("panic while transposing table stream");
444 });
445 let strategy = TableStrategy::new(
446 Arc::new(FlatLayoutStrategy::default()),
447 Arc::new(FlatLayoutStrategy::default()),
448 );
449
450 block_on(|handle| async move {
451 let session = SESSION.clone().with_handle(handle);
452 strategy
453 .write_stream(
454 ctx,
455 segments,
456 SequentialStreamAdapter::new(dtype, stream).sendable(),
457 eof,
458 &session,
459 )
460 .await
461 .unwrap();
462 });
463 }
464}