vortex_layout/layouts/flat/
writer.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use async_trait::async_trait;
5use futures::StreamExt;
6use vortex_array::dtype::DType;
7use vortex_array::expr::stats::Precision;
8use vortex_array::expr::stats::Stat;
9use vortex_array::expr::stats::StatsProvider;
10use vortex_array::scalar::Scalar;
11use vortex_array::scalar::ScalarTruncation;
12use vortex_array::scalar::lower_bound;
13use vortex_array::scalar::upper_bound;
14use vortex_array::serde::SerializeOptions;
15use vortex_array::stats::StatsSetRef;
16use vortex_buffer::BufferString;
17use vortex_buffer::ByteBuffer;
18use vortex_error::VortexExpect;
19use vortex_error::VortexResult;
20use vortex_error::vortex_bail;
21use vortex_session::VortexSession;
22use vortex_session::registry::ReadContext;
23
24use crate::LayoutRef;
25use crate::LayoutStrategy;
26use crate::LayoutWriterContext;
27use crate::children::OwnedLayoutChildren;
28use crate::layouts::chunked::ChunkedLayout;
29use crate::layouts::flat::FlatLayout;
30use crate::layouts::flat::flat_layout_inline_array_node;
31use crate::segments::SegmentSinkRef;
32use crate::sequence::SendableSequentialStream;
33use crate::sequence::SequencePointer;
34
35#[derive(Clone)]
36pub struct FlatLayoutStrategy {
37    /// Whether to include padding for memory-mapped reads.
38    pub include_padding: bool,
39    /// Maximum length of variable length statistics
40    pub max_variable_length_statistics_size: usize,
41}
42
43impl Default for FlatLayoutStrategy {
44    fn default() -> Self {
45        Self {
46            include_padding: true,
47            max_variable_length_statistics_size: 64,
48        }
49    }
50}
51
52impl FlatLayoutStrategy {
53    /// Set whether to include padding for memory-mapped reads.
54    pub fn with_include_padding(mut self, include_padding: bool) -> Self {
55        self.include_padding = include_padding;
56        self
57    }
58
59    /// Set the maximum length of variable length statistics.
60    pub fn with_max_variable_length_statistics_size(mut self, size: usize) -> Self {
61        self.max_variable_length_statistics_size = size;
62        self
63    }
64}
65
66fn truncate_scalar_stat<F: Fn(Scalar) -> Option<(Scalar, bool)>>(
67    statistics: StatsSetRef<'_>,
68    stat: Stat,
69    truncation: F,
70) {
71    if let Some(sv) = statistics.get(stat).into_inner() {
72        if let Some((truncated_value, truncated)) = truncation(sv) {
73            if truncated && let Some(v) = truncated_value.into_value() {
74                statistics.set(stat, Precision::Inexact(v));
75            }
76        } else {
77            statistics.clear(stat)
78        }
79    }
80}
81
82#[async_trait]
83impl LayoutStrategy for FlatLayoutStrategy {
84    async fn write_stream(
85        &self,
86        ctx: LayoutWriterContext,
87        segment_sink: SegmentSinkRef,
88        mut stream: SendableSequentialStream,
89        _eof: SequencePointer,
90        session: &VortexSession,
91    ) -> VortexResult<LayoutRef> {
92        let Some(chunk) = stream.next().await else {
93            // an empty input has no segment to write.
94            return Ok(ChunkedLayout::new(
95                0,
96                stream.dtype().clone(),
97                OwnedLayoutChildren::layout_children(vec![]),
98            )
99            .into_layout());
100        };
101        let (sequence_id, chunk) = chunk?;
102
103        let row_count = chunk.len() as u64;
104
105        match chunk.dtype() {
106            DType::Utf8(n) => {
107                truncate_scalar_stat(chunk.statistics(), Stat::Min, |v| {
108                    lower_bound(
109                        BufferString::from_scalar(v)
110                            .vortex_expect("utf8 scalar must be a BufferString"),
111                        self.max_variable_length_statistics_size,
112                        *n,
113                    )
114                });
115                truncate_scalar_stat(chunk.statistics(), Stat::Max, |v| {
116                    upper_bound(
117                        BufferString::from_scalar(v)
118                            .vortex_expect("utf8 scalar must be a BufferString"),
119                        self.max_variable_length_statistics_size,
120                        *n,
121                    )
122                });
123            }
124            DType::Binary(n) => {
125                truncate_scalar_stat(chunk.statistics(), Stat::Min, |v| {
126                    lower_bound(
127                        ByteBuffer::from_scalar(v)
128                            .vortex_expect("binary scalar must be a ByteBuffer"),
129                        self.max_variable_length_statistics_size,
130                        *n,
131                    )
132                });
133                truncate_scalar_stat(chunk.statistics(), Stat::Max, |v| {
134                    upper_bound(
135                        ByteBuffer::from_scalar(v)
136                            .vortex_expect("binary scalar must be a ByteBuffer"),
137                        self.max_variable_length_statistics_size,
138                        *n,
139                    )
140                });
141            }
142            _ => {}
143        }
144
145        let buffers = chunk.serialize(
146            ctx.array_ctx(),
147            session,
148            &SerializeOptions {
149                offset: 0,
150                include_padding: self.include_padding,
151            },
152        )?;
153        // there is at least the flatbuffer and the length
154        assert!(buffers.len() >= 2);
155        let array_node =
156            flat_layout_inline_array_node().then(|| buffers[buffers.len() - 2].clone());
157        let segment_id = segment_sink.write(sequence_id, buffers).await?;
158
159        let None = stream.next().await else {
160            vortex_bail!("flat layout received stream with more than a single chunk");
161        };
162        Ok(FlatLayout::new_with_metadata(
163            row_count,
164            stream.dtype().clone(),
165            segment_id,
166            ReadContext::new(ctx.array_ctx().to_ids()),
167            array_node,
168        )
169        .into_layout())
170    }
171}
172
173#[cfg(test)]
174mod tests {
175    use std::sync::Arc;
176
177    use vortex_array::ArrayContext;
178    use vortex_array::ArrayRef;
179    use vortex_array::IntoArray;
180    use vortex_array::MaskFuture;
181    use vortex_array::VortexSessionExecute;
182    use vortex_array::array_session;
183    use vortex_array::arrays::BoolArray;
184    use vortex_array::arrays::PrimitiveArray;
185    use vortex_array::arrays::StructArray;
186    use vortex_array::arrays::struct_::StructArrayExt;
187    use vortex_array::builders::ArrayBuilder;
188    use vortex_array::builders::VarBinViewBuilder;
189    use vortex_array::dtype::DType;
190    use vortex_array::dtype::FieldName;
191    use vortex_array::dtype::FieldNames;
192    use vortex_array::dtype::Nullability;
193    use vortex_array::expr::root;
194    use vortex_array::expr::stats::Precision;
195    use vortex_array::expr::stats::Stat;
196    use vortex_array::expr::stats::StatsProviderExt;
197    use vortex_array::validity::Validity;
198    use vortex_buffer::BitBufferMut;
199    use vortex_buffer::buffer;
200    use vortex_error::VortexExpect;
201    use vortex_io::runtime::single::block_on;
202    use vortex_io::session::RuntimeSessionExt;
203    use vortex_mask::AllOr;
204
205    use crate::LayoutStrategy;
206    use crate::layouts::flat::writer::FlatLayoutStrategy;
207    use crate::segments::TestSegments;
208    use crate::sequence::SequenceId;
209    use crate::sequence::SequentialArrayStreamExt;
210    use crate::test::SESSION;
211    use crate::test::new_session;
212
213    // Currently, flat layouts do not force compute stats during write, they only retain
214    // pre-computed stats.
215    #[should_panic]
216    #[test]
217    fn flat_stats() {
218        block_on(|handle| async {
219            let session = new_session().with_handle(handle);
220            let ctx = ArrayContext::empty();
221            let segments = Arc::new(TestSegments::default());
222            let (ptr, eof) = SequenceId::root().split();
223            let array = PrimitiveArray::new(buffer![1, 2, 3, 4, 5], Validity::AllValid);
224            let layout = FlatLayoutStrategy::default()
225                .write_stream(
226                    ctx.into(),
227                    Arc::<TestSegments>::clone(&segments),
228                    array.into_array().to_array_stream().sequenced(ptr),
229                    eof,
230                    &session,
231                )
232                .await
233                .unwrap();
234
235            let reader = layout
236                .new_reader("".into(), segments, &SESSION, &Default::default())
237                .unwrap();
238            let expr = root().bind(reader.dtype()).unwrap();
239            let result = reader
240                .projection_evaluation(
241                    &(0..layout.row_count()),
242                    &expr,
243                    MaskFuture::new_true(layout.row_count().try_into().unwrap()),
244                )
245                .unwrap()
246                .await
247                .unwrap();
248
249            assert_eq!(
250                result.statistics().get_as::<bool>(Stat::IsSorted),
251                Precision::Exact(true)
252            );
253        })
254    }
255
256    #[test]
257    fn truncates_variable_size_stats() {
258        block_on(|handle| async {
259            let session = new_session().with_handle(handle);
260            let ctx = ArrayContext::empty();
261            let segments = Arc::new(TestSegments::default());
262            let (ptr, eof) = SequenceId::root().split();
263            let mut builder = VarBinViewBuilder::with_capacity_in(
264                DType::Utf8(Nullability::NonNullable),
265                2,
266                vortex_buffer::BufferAllocatorRef::statically_allocated(),
267            );
268            builder.append_value("Long value to test that the statistics are actually truncated, it needs a bit of extra padding though");
269            builder.append_value("Another string that's meant to be smaller than the previous value, though still need extra padding");
270            let array = builder.finish();
271            let mut stats_ctx = session.create_execution_ctx();
272            array.statistics().set_iter(
273                array
274                    .statistics()
275                    .compute_all(&Stat::all().collect::<Vec<_>>(), &mut stats_ctx)
276                    .vortex_expect("stats computation should succeed for test array")
277                    .into_iter(),
278            );
279
280            let layout = FlatLayoutStrategy::default()
281                .write_stream(
282                    ctx.into(),
283                    Arc::<TestSegments>::clone(&segments),
284                    array.into_array().to_array_stream().sequenced(ptr),
285                    eof,
286                    &session,
287                )
288                .await
289                .unwrap();
290
291            let reader = layout
292                .new_reader("".into(), segments, &SESSION, &Default::default())
293                .unwrap();
294            let expr = root().bind(reader.dtype()).unwrap();
295            let result = reader
296                .projection_evaluation(
297                    &(0..layout.row_count()),
298                    &expr,
299                    MaskFuture::new_true(layout.row_count().try_into().unwrap()),
300                )
301                .unwrap()
302                .await
303                .unwrap();
304
305            assert_eq!(
306                result.statistics().get_as::<String>(Stat::Min),
307                // The typo is correct, we need this to be truncated.
308                Precision::Inexact(
309                    // spellchecker:ignore-next-line
310                    "Another string that's meant to be smaller than the previous valu".to_string()
311                )
312            );
313            assert_eq!(
314                result.statistics().get_as::<String>(Stat::Max),
315                Precision::Inexact(
316                    "Long value to test that the statistics are actually truncated, j".to_string()
317                )
318            );
319        })
320    }
321
322    #[test]
323    fn struct_array_round_trip() {
324        block_on(|handle| async {
325            let mut ctx_exec = array_session().create_execution_ctx();
326            let session = new_session().with_handle(handle);
327            let mut validity_builder = BitBufferMut::with_capacity(2);
328            validity_builder.append(true);
329            validity_builder.append(false);
330            let validity_boolean_buffer = validity_builder.freeze();
331            let validity = Validity::Array(
332                BoolArray::new(validity_boolean_buffer.clone(), Validity::NonNullable).into_array(),
333            );
334            let array = StructArray::try_new(
335                FieldNames::from([FieldName::from("a"), FieldName::from("b")]),
336                vec![
337                    buffer![1_u64, 2].into_array(),
338                    buffer![3_u64, 4].into_array(),
339                ],
340                2,
341                validity,
342            )
343            .unwrap();
344
345            let ctx = ArrayContext::empty();
346
347            // Write the array into a byte buffer.
348            let (layout, segments) = {
349                let segments = Arc::new(TestSegments::default());
350                let (ptr, eof) = SequenceId::root().split();
351                let layout = FlatLayoutStrategy::default()
352                    .write_stream(
353                        ctx.into(),
354                        Arc::<TestSegments>::clone(&segments),
355                        array.into_array().to_array_stream().sequenced(ptr),
356                        eof,
357                        &session,
358                    )
359                    .await
360                    .unwrap();
361
362                (layout, segments)
363            };
364
365            // We should be able to read the array we just wrote.
366            let reader = layout
367                .new_reader("".into(), segments, &SESSION, &Default::default())
368                .unwrap();
369            let expr = root().bind(reader.dtype()).unwrap();
370            let result: ArrayRef = reader
371                .projection_evaluation(
372                    &(0..layout.row_count()),
373                    &expr,
374                    MaskFuture::new_true(layout.row_count().try_into().unwrap()),
375                )
376                .unwrap()
377                .await
378                .unwrap();
379
380            assert_eq!(
381                result
382                    .validity()
383                    .unwrap()
384                    .execute_mask(result.len(), &mut ctx_exec)
385                    .unwrap()
386                    .bit_buffer(),
387                AllOr::Some(&validity_boolean_buffer)
388            );
389            let result_struct = result
390                .clone()
391                .execute::<StructArray>(&mut ctx_exec)
392                .unwrap();
393            let field_a = result_struct
394                .unmasked_field_by_name("a")
395                .unwrap()
396                .clone()
397                .execute::<PrimitiveArray>(&mut ctx_exec)
398                .unwrap();
399            assert_eq!(field_a.as_slice::<u64>(), &[1, 2]);
400            let result_struct_b = result.execute::<StructArray>(&mut ctx_exec).unwrap();
401            let field_b = result_struct_b
402                .unmasked_field_by_name("b")
403                .unwrap()
404                .clone()
405                .execute::<PrimitiveArray>(&mut ctx_exec)
406                .unwrap();
407            assert_eq!(field_b.as_slice::<u64>(), &[3, 4]);
408        })
409    }
410}
vortex_layout/layouts/flat/writer.rs

vortex_layout/layouts/flat/
writer.rs
