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    #[test]
216    fn flat_does_not_compute_missing_stats() {
217        block_on(|handle| async {
218            let session = new_session().with_handle(handle);
219            let ctx = ArrayContext::empty();
220            let segments = Arc::new(TestSegments::default());
221            let (ptr, eof) = SequenceId::root().split();
222            let array = PrimitiveArray::new(buffer![1, 2, 3, 4, 5], Validity::AllValid);
223            let layout = FlatLayoutStrategy::default()
224                .write_stream(
225                    ctx.into(),
226                    Arc::<TestSegments>::clone(&segments),
227                    array.into_array().to_array_stream().sequenced(ptr),
228                    eof,
229                    &session,
230                )
231                .await
232                .unwrap();
233
234            let reader = layout
235                .new_reader("".into(), segments, &SESSION, &Default::default())
236                .unwrap();
237            let expr = root().bind(reader.dtype()).unwrap();
238            let result = reader
239                .projection_evaluation(
240                    &(0..layout.row_count()),
241                    &expr,
242                    MaskFuture::new_true(layout.row_count().try_into().unwrap()),
243                )
244                .unwrap()
245                .await
246                .unwrap();
247
248            assert_eq!(
249                result.statistics().get_as::<bool>(Stat::IsSorted),
250                Precision::Absent
251            );
252        })
253    }
254
255    #[test]
256    fn truncates_variable_size_stats() {
257        block_on(|handle| async {
258            let session = new_session().with_handle(handle);
259            let ctx = ArrayContext::empty();
260            let segments = Arc::new(TestSegments::default());
261            let (ptr, eof) = SequenceId::root().split();
262            let mut builder = VarBinViewBuilder::with_capacity_in(
263                DType::Utf8(Nullability::NonNullable),
264                2,
265                vortex_buffer::BufferAllocatorRef::statically_allocated(),
266            );
267            builder.append_value("Long value to test that the statistics are actually truncated, it needs a bit of extra padding though");
268            builder.append_value("Another string that's meant to be smaller than the previous value, though still need extra padding");
269            let array = builder.finish();
270            let mut stats_ctx = session.create_execution_ctx();
271            array.statistics().set_iter(
272                array
273                    .statistics()
274                    .compute_all(&Stat::all().collect::<Vec<_>>(), &mut stats_ctx)
275                    .vortex_expect("stats computation should succeed for test array")
276                    .into_iter(),
277            );
278
279            let layout = FlatLayoutStrategy::default()
280                .write_stream(
281                    ctx.into(),
282                    Arc::<TestSegments>::clone(&segments),
283                    array.into_array().to_array_stream().sequenced(ptr),
284                    eof,
285                    &session,
286                )
287                .await
288                .unwrap();
289
290            let reader = layout
291                .new_reader("".into(), segments, &SESSION, &Default::default())
292                .unwrap();
293            let expr = root().bind(reader.dtype()).unwrap();
294            let result = reader
295                .projection_evaluation(
296                    &(0..layout.row_count()),
297                    &expr,
298                    MaskFuture::new_true(layout.row_count().try_into().unwrap()),
299                )
300                .unwrap()
301                .await
302                .unwrap();
303
304            assert_eq!(
305                result.statistics().get_as::<String>(Stat::Min),
306                // The typo is correct, we need this to be truncated.
307                Precision::Inexact(
308                    // spellchecker:ignore-next-line
309                    "Another string that's meant to be smaller than the previous valu".to_string()
310                )
311            );
312            assert_eq!(
313                result.statistics().get_as::<String>(Stat::Max),
314                Precision::Inexact(
315                    "Long value to test that the statistics are actually truncated, j".to_string()
316                )
317            );
318        })
319    }
320
321    #[test]
322    fn struct_array_round_trip() {
323        block_on(|handle| async {
324            let mut ctx_exec = array_session().create_execution_ctx();
325            let session = new_session().with_handle(handle);
326            let mut validity_builder = BitBufferMut::with_capacity(2);
327            validity_builder.append(true);
328            validity_builder.append(false);
329            let validity_boolean_buffer = validity_builder.freeze();
330            let validity = Validity::Array(
331                BoolArray::new(validity_boolean_buffer.clone(), Validity::NonNullable).into_array(),
332            );
333            let array = StructArray::try_new(
334                FieldNames::from([FieldName::from("a"), FieldName::from("b")]),
335                vec![
336                    buffer![1_u64, 2].into_array(),
337                    buffer![3_u64, 4].into_array(),
338                ],
339                2,
340                validity,
341            )
342            .unwrap();
343
344            let ctx = ArrayContext::empty();
345
346            // Write the array into a byte buffer.
347            let (layout, segments) = {
348                let segments = Arc::new(TestSegments::default());
349                let (ptr, eof) = SequenceId::root().split();
350                let layout = FlatLayoutStrategy::default()
351                    .write_stream(
352                        ctx.into(),
353                        Arc::<TestSegments>::clone(&segments),
354                        array.into_array().to_array_stream().sequenced(ptr),
355                        eof,
356                        &session,
357                    )
358                    .await
359                    .unwrap();
360
361                (layout, segments)
362            };
363
364            // We should be able to read the array we just wrote.
365            let reader = layout
366                .new_reader("".into(), segments, &SESSION, &Default::default())
367                .unwrap();
368            let expr = root().bind(reader.dtype()).unwrap();
369            let result: ArrayRef = reader
370                .projection_evaluation(
371                    &(0..layout.row_count()),
372                    &expr,
373                    MaskFuture::new_true(layout.row_count().try_into().unwrap()),
374                )
375                .unwrap()
376                .await
377                .unwrap();
378
379            assert_eq!(
380                result
381                    .validity()
382                    .unwrap()
383                    .execute_mask(result.len(), &mut ctx_exec)
384                    .unwrap()
385                    .bit_buffer(),
386                AllOr::Some(&validity_boolean_buffer)
387            );
388            let result_struct = result
389                .clone()
390                .execute::<StructArray>(&mut ctx_exec)
391                .unwrap();
392            let field_a = result_struct
393                .unmasked_field_by_name("a")
394                .unwrap()
395                .clone()
396                .execute::<PrimitiveArray>(&mut ctx_exec)
397                .unwrap();
398            assert_eq!(field_a.as_slice::<u64>(), &[1, 2]);
399            let result_struct_b = result.execute::<StructArray>(&mut ctx_exec).unwrap();
400            let field_b = result_struct_b
401                .unmasked_field_by_name("b")
402                .unwrap()
403                .clone()
404                .execute::<PrimitiveArray>(&mut ctx_exec)
405                .unwrap();
406            assert_eq!(field_b.as_slice::<u64>(), &[3, 4]);
407        })
408    }
409}
vortex_layout/layouts/flat/writer.rs

vortex_layout/layouts/flat/
writer.rs
