vortex_layout/layouts/indexed/
writer.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Write-time assembly: forward chunks to the data child while feeding index builders, then write
5//! each index's content after all data segments.
6
7use std::num::NonZeroU64;
8use std::sync::Arc;
9
10use async_trait::async_trait;
11use futures::StreamExt;
12use futures::stream;
13use parking_lot::Mutex;
14use roaring::RoaringBitmap;
15use tracing::trace;
16use vortex_array::ArrayRef;
17use vortex_array::ExecutionCtx;
18use vortex_array::VortexSessionExecute;
19use vortex_array::dtype::DType;
20use vortex_array::stream::ArrayStreamAdapter;
21use vortex_array::stream::ArrayStreamExt;
22use vortex_array::stream::SendableArrayStream;
23use vortex_error::VortexResult;
24use vortex_error::vortex_bail;
25use vortex_error::vortex_ensure;
26use vortex_session::VortexSession;
27
28use crate::BufferedBytesReservation;
29use crate::BufferedBytesTracker;
30use crate::LayoutRef;
31use crate::LayoutStrategy;
32use crate::LayoutWriterContext;
33use crate::layouts::indexed::IndexPartitioning;
34use crate::layouts::indexed::IndexSpec;
35use crate::layouts::indexed::IndexedLayout;
36use crate::layouts::indexed::index::IndexBuilder;
37use crate::layouts::indexed::index::IndexVTable;
38use crate::layouts::indexed::index::IndexVTableRef;
39use crate::segments::SegmentSinkRef;
40use crate::sequence::SendableSequentialStream;
41use crate::sequence::SequencePointer;
42use crate::sequence::SequentialArrayStreamExt;
43use crate::sequence::SequentialStreamAdapter;
44use crate::sequence::SequentialStreamExt;
45
46/// An index to attach to a column, as configured on the write side.
47#[derive(Clone, Debug)]
48pub struct IndexConfig {
49    vtable: IndexVTableRef,
50    options: Vec<u8>,
51    partition_len: Option<NonZeroU64>,
52}
53
54impl IndexConfig {
55    /// Configure `vtable` with kind-defined options.
56    pub fn new(vtable: IndexVTableRef, options: Vec<u8>) -> Self {
57        Self {
58            vtable,
59            options,
60            partition_len: None,
61        }
62    }
63
64    /// Configure `vtable` with its default options.
65    pub fn with_defaults(vtable: IndexVTableRef) -> Self {
66        Self::new(vtable, Vec::new())
67    }
68
69    /// Build and probe this index in independent partitions of `partition_len` data rows, rather
70    /// than as one index over the whole column.
71    ///
72    /// Partitions are independent of how the data child is chunked, so a scan over part of the
73    /// column reads only the overlapping partitions' index rows. `partition_len` must be a
74    /// multiple of the strategy's [`IndexedStrategy::with_data_block_len`], which must be set.
75    pub fn with_partition_len(mut self, partition_len: NonZeroU64) -> Self {
76        self.partition_len = Some(partition_len);
77        self
78    }
79}
80
81/// Wraps a data-child strategy with one or more index builders.
82///
83/// Sits in the same slot as `ZonedStrategy`, meaning above the repartition step, so it knows the
84/// data child's row block size and can hand it to block-granular index kinds, making their blocks
85/// line up with the data child's chunks by default.
86pub struct IndexedStrategy {
87    data: Arc<dyn LayoutStrategy>,
88    index: Arc<dyn LayoutStrategy>,
89    configs: Arc<[IndexConfig]>,
90    data_block_len: Option<u64>,
91}
92
93impl IndexedStrategy {
94    /// Create a strategy writing data through `data` and each index's content through `index`.
95    pub fn new<D: LayoutStrategy, I: LayoutStrategy>(
96        data: D,
97        index: I,
98        configs: Vec<IndexConfig>,
99    ) -> Self {
100        Self {
101            data: Arc::new(data),
102            index: Arc::new(index),
103            configs: configs.into(),
104            data_block_len: None,
105        }
106    }
107
108    /// Tell block-granular index kinds the data child's row block size, so pruned blocks align
109    /// with chunk and segment boundaries.
110    pub fn with_data_block_len(mut self, data_block_len: u64) -> Self {
111        self.data_block_len = Some(data_block_len);
112        self
113    }
114}
115
116/// An index kind with its configured options parsed, once per write.
117pub(crate) trait IndexWriter: Send + Sync {
118    fn index_dtype(&self, dtype: &DType) -> Option<DType>;
119
120    /// A builder for one partition, whose output must be of `index_dtype`.
121    fn builder(
122        &self,
123        dtype: &DType,
124        index_dtype: &DType,
125        data_block_len: Option<u64>,
126        session: &VortexSession,
127    ) -> VortexResult<Box<dyn PartitionBuilder>>;
128}
129
130pub(crate) struct TypedIndexWriter<V: IndexVTable> {
131    vtable: Arc<V>,
132    options: V::Options,
133}
134
135/// One partition's builder, with the kind's types erased. This is where a kind's content and
136/// options become bytes, through [`IndexVTable::encode`] and [`IndexVTable::serialize_options`].
137pub(crate) trait PartitionBuilder: Send {
138    fn push(
139        &mut self,
140        chunk: &ArrayRef,
141        row_offset: u64,
142        ctx: &mut ExecutionCtx,
143    ) -> VortexResult<()>;
144
145    fn finish(self: Box<Self>) -> VortexResult<PartitionOutput>;
146
147    fn buffered_bytes(&self) -> u64;
148}
149
150struct TypedPartitionBuilder<V: IndexVTable> {
151    vtable: Arc<V>,
152    builder: V::Builder,
153    index_dtype: DType,
154    session: VortexSession,
155}
156
157impl<V: IndexVTable> PartitionBuilder for TypedPartitionBuilder<V> {
158    fn push(
159        &mut self,
160        chunk: &ArrayRef,
161        row_offset: u64,
162        ctx: &mut ExecutionCtx,
163    ) -> VortexResult<()> {
164        self.builder.push(chunk, row_offset, ctx)
165    }
166
167    fn finish(self: Box<Self>) -> VortexResult<PartitionOutput> {
168        let Self {
169            vtable,
170            builder,
171            index_dtype,
172            session,
173        } = *self;
174        let Some((chunks, options)) = builder.finish()? else {
175            return Ok(None);
176        };
177
178        let mut ctx = session.create_execution_ctx();
179        let arrays = chunks
180            .iter()
181            .map(|chunk| {
182                let array = vtable.encode(chunk, &options, &mut ctx)?;
183                vortex_ensure!(
184                    array.dtype() == &index_dtype,
185                    "Index {} encoded a chunk of dtype {}, but the index kind declared {index_dtype}",
186                    vtable.id(),
187                    array.dtype()
188                );
189                Ok(array)
190            })
191            .collect::<VortexResult<Vec<_>>>()?;
192
193        let content =
194            ArrayStreamAdapter::new(index_dtype, stream::iter(arrays.into_iter().map(Ok)));
195        Ok(Some((
196            ArrayStreamExt::boxed(content),
197            vtable.serialize_options(&options),
198        )))
199    }
200
201    fn buffered_bytes(&self) -> u64 {
202        self.builder.buffered_bytes()
203    }
204}
205
206impl<V: IndexVTable> TypedIndexWriter<V> {
207    pub(crate) fn try_new(vtable: Arc<V>, options: &[u8]) -> VortexResult<Self> {
208        let options = vtable.deserialize_options(options)?;
209        Ok(Self { vtable, options })
210    }
211}
212
213impl<V: IndexVTable> IndexWriter for TypedIndexWriter<V> {
214    fn index_dtype(&self, dtype: &DType) -> Option<DType> {
215        self.vtable.index_dtype(dtype, &self.options)
216    }
217
218    fn builder(
219        &self,
220        dtype: &DType,
221        index_dtype: &DType,
222        data_block_len: Option<u64>,
223        session: &VortexSession,
224    ) -> VortexResult<Box<dyn PartitionBuilder>> {
225        Ok(Box::new(TypedPartitionBuilder {
226            vtable: Arc::clone(&self.vtable),
227            builder: self
228                .vtable
229                .builder(dtype, &self.options, data_block_len, session)?,
230            index_dtype: index_dtype.clone(),
231            session: session.clone(),
232        }))
233    }
234}
235
236/// Everything needed to start another partition's builder mid-stream.
237struct BuilderFactory {
238    dtype: DType,
239    data_block_len: Option<u64>,
240    session: VortexSession,
241}
242
243impl BuilderFactory {
244    fn builder(
245        &self,
246        writer: &dyn IndexWriter,
247        index_dtype: &DType,
248    ) -> VortexResult<Box<dyn PartitionBuilder>> {
249        writer.builder(&self.dtype, index_dtype, self.data_block_len, &self.session)
250    }
251}
252
253/// A finished partition's output: its content stream and final options, or `None` if it declined.
254type PartitionOutput = Option<(SendableArrayStream, Vec<u8>)>;
255
256/// One configured index, built partition by partition.
257struct IndexState {
258    vtable: IndexVTableRef,
259    /// The kind with its configured options parsed, starting every partition's builder.
260    writer: Arc<dyn IndexWriter>,
261    /// What the kind declared it builds, which every partition's output must match.
262    index_dtype: DType,
263    partition_len: Option<u64>,
264    builder: Box<dyn PartitionBuilder>,
265    /// Rows the current builder has seen, which is also the next row's partition-local offset.
266    partition_rows: u64,
267    finished: Vec<PartitionOutput>,
268    /// What each finished builder last reported buffering; its content stream still holds it.
269    finished_bytes: u64,
270}
271
272impl IndexState {
273    fn push(
274        &mut self,
275        chunk: &ArrayRef,
276        factory: &BuilderFactory,
277        ctx: &mut ExecutionCtx,
278    ) -> VortexResult<()> {
279        let Some(partition_len) = self.partition_len else {
280            return self
281                .builder
282                .push(chunk, self.partition_rows, ctx)
283                .map(|()| {
284                    self.partition_rows += chunk.len() as u64;
285                });
286        };
287
288        // Chunks come from upstream of any repartitioning, so one can straddle several
289        // partitions; each builder must see only its own partition's rows.
290        let mut start = 0;
291        while start < chunk.len() {
292            let room = usize::try_from(partition_len - self.partition_rows)?;
293            let end = chunk.len().min(start + room);
294            let piece = if start == 0 && end == chunk.len() {
295                chunk.clone()
296            } else {
297                chunk.slice(start..end)?
298            };
299            self.builder.push(&piece, self.partition_rows, ctx)?;
300            self.partition_rows += (end - start) as u64;
301            start = end;
302
303            if self.partition_rows == partition_len {
304                let next = factory.builder(self.writer.as_ref(), &self.index_dtype)?;
305                self.finish_partition(next)?;
306            }
307        }
308        Ok(())
309    }
310
311    fn finish_partition(&mut self, next: Box<dyn PartitionBuilder>) -> VortexResult<()> {
312        let done = std::mem::replace(&mut self.builder, next);
313        self.finished_bytes += done.buffered_bytes();
314        self.finished.push(done.finish()?);
315        self.partition_rows = 0;
316        Ok(())
317    }
318
319    fn buffered_bytes(&self) -> u64 {
320        self.builder.buffered_bytes() + self.finished_bytes
321    }
322
323    /// Finish the trailing partition and return every partition's output in row order.
324    ///
325    /// An unpartitioned index always yields exactly one partition, even over zero rows, leaving
326    /// the empty case to the builder. A partitioned index yields one per started partition, so
327    /// the builder opened after an exactly-full final partition is dropped unused.
328    fn finish(mut self) -> VortexResult<Vec<PartitionOutput>> {
329        if self.partition_len.is_none() || self.partition_rows > 0 {
330            self.finished.push(self.builder.finish()?);
331        }
332        Ok(self.finished)
333    }
334}
335
336/// Index state shared between the stream-mapping closure and the finishing code. Index building
337/// is globally stateful, so pushes stay sequential in stream order.
338struct BuilderState {
339    factory: BuilderFactory,
340    indexes: Vec<IndexState>,
341    /// Reservation reflecting the indexes' current combined buffered bytes.
342    ///
343    /// Builders report a running total rather than a per-chunk delta, so each push replaces this
344    /// reservation (drop the old, reserve the new total) instead of accumulating one reservation
345    /// per chunk the way `BufferedStrategy` does for known-size chunks.
346    buffered: Option<BufferedBytesReservation>,
347}
348
349impl BuilderState {
350    fn push(&mut self, chunk: &ArrayRef, tracker: &BufferedBytesTracker) -> VortexResult<()> {
351        let mut ctx = self.factory.session.create_execution_ctx();
352        for index in &mut self.indexes {
353            index.push(chunk, &self.factory, &mut ctx)?;
354        }
355
356        let total: u64 = self.indexes.iter().map(IndexState::buffered_bytes).sum();
357        self.buffered = Some(tracker.reserve(total));
358        Ok(())
359    }
360}
361
362/// Every built partition of one index, concatenated into a single stream for the index child.
363struct IndexContent {
364    stream: SendableArrayStream,
365    options: Vec<u8>,
366    dtype: DType,
367    /// Index rows written per partition, filled in as the index child consumes `stream`.
368    partition_rows: Arc<Mutex<Vec<u64>>>,
369    declined: RoaringBitmap,
370}
371
372impl IndexContent {
373    /// Concatenate the built partitions, each already of the declared `dtype`, or `None` if every
374    /// partition declined.
375    fn assemble(
376        id: impl std::fmt::Display,
377        dtype: DType,
378        partitions: Vec<PartitionOutput>,
379    ) -> VortexResult<Option<Self>> {
380        let mut declined = RoaringBitmap::new();
381        let mut built = Vec::with_capacity(partitions.len());
382        let mut options: Option<Vec<u8>> = None;
383        let partition_count = partitions.len();
384
385        for (partition, output) in partitions.into_iter().enumerate() {
386            let Some((content, partition_options)) = output else {
387                declined.insert(u32::try_from(partition)?);
388                continue;
389            };
390            match &options {
391                None => options = Some(partition_options),
392                // One spec records one options blob, so partitions may not disagree about it.
393                Some(first) => vortex_ensure!(
394                    *first == partition_options,
395                    "Index {id} partition {partition} finished with different options than earlier partitions"
396                ),
397            }
398            built.push((partition, content));
399        }
400
401        let Some(options) = options else {
402            return Ok(None);
403        };
404
405        let partition_rows = Arc::new(Mutex::new(vec![0u64; partition_count]));
406        let counts = Arc::clone(&partition_rows);
407        let chunks = stream::iter(built).flat_map(move |(partition, content)| {
408            let counts = Arc::clone(&counts);
409            content.map(move |chunk| {
410                if let Ok(chunk) = &chunk {
411                    counts.lock()[partition] += chunk.len() as u64;
412                }
413                chunk
414            })
415        });
416
417        Ok(Some(Self {
418            stream: ArrayStreamExt::boxed(ArrayStreamAdapter::new(dtype.clone(), chunks)),
419            options,
420            dtype,
421            partition_rows,
422            declined,
423        }))
424    }
425}
426
427/// The partitioning to record for an index, from the index rows each partition wrote.
428fn partitioning(
429    partition_len: u64,
430    partition_rows: &[u64],
431    declined: RoaringBitmap,
432) -> IndexPartitioning {
433    let ends = partition_rows
434        .iter()
435        .scan(0u64, |end, rows| {
436            *end += rows;
437            Some(*end)
438        })
439        .collect();
440    IndexPartitioning::new(partition_len, ends, declined)
441}
442
443#[async_trait]
444impl LayoutStrategy for IndexedStrategy {
445    async fn write_stream(
446        &self,
447        ctx: LayoutWriterContext,
448        segment_sink: SegmentSinkRef,
449        stream: SendableSequentialStream,
450        mut eof: SequencePointer,
451        session: &VortexSession,
452    ) -> VortexResult<LayoutRef> {
453        let dtype = stream.dtype().clone();
454
455        // Checked before the dtype filter below, so a misconfigured partitioning fails every
456        // write rather than only the ones whose column happens to be indexable.
457        for config in self.configs.iter() {
458            let Some(partition_len) = config.partition_len else {
459                continue;
460            };
461            let Some(block_len) = self.data_block_len.filter(|len| *len > 0) else {
462                vortex_bail!(
463                    "Index {} is partitioned every {partition_len} rows, which requires a non-zero \
464                     data block length (IndexedStrategy::with_data_block_len)",
465                    config.vtable.id()
466                );
467            };
468            vortex_ensure!(
469                partition_len.get() % block_len == 0,
470                "Index {} partition length {partition_len} must be a multiple of the data block \
471                 length {block_len}",
472                config.vtable.id()
473            );
474        }
475
476        let factory = BuilderFactory {
477            dtype: dtype.clone(),
478            data_block_len: self.data_block_len,
479            session: session.clone(),
480        };
481        let mut indexes = Vec::with_capacity(self.configs.len());
482        for config in self.configs.iter() {
483            let writer = config.vtable.open_writer(&config.options)?;
484            let Some(index_dtype) = writer.index_dtype(&dtype) else {
485                continue;
486            };
487            indexes.push(IndexState {
488                builder: factory.builder(writer.as_ref(), &index_dtype)?,
489                writer,
490                index_dtype,
491                vtable: config.vtable.clone(),
492                partition_len: config.partition_len.map(NonZeroU64::get),
493                partition_rows: 0,
494                finished: Vec::new(),
495                finished_bytes: 0,
496            });
497        }
498
499        // Nothing to index for this dtype: don't emit a wrapper at all, so readers see the plain
500        // data layout.
501        if indexes.is_empty() {
502            return self
503                .data
504                .write_stream(ctx, segment_sink, stream, eof, session)
505                .await;
506        }
507
508        let state = Arc::new(Mutex::new(BuilderState {
509            factory,
510            indexes,
511            buffered: None,
512        }));
513
514        let feed_state = Arc::clone(&state);
515        let feed_tracker = ctx.buffered_bytes_tracker().clone();
516        let stream = SequentialStreamAdapter::new(
517            dtype,
518            stream.map(move |item| {
519                let (sequence_id, chunk) = item?;
520                feed_state.lock().push(&chunk, &feed_tracker)?;
521                Ok((sequence_id, chunk))
522            }),
523        )
524        .sendable();
525
526        // Data segments come first, so a reader that ignores indexes keeps its locality and a
527        // streaming writer never has to seek back.
528        let data_eof = eof.split_off();
529        let data_layout = self
530            .data
531            .write_stream(
532                ctx.clone(),
533                Arc::clone(&segment_sink),
534                stream,
535                data_eof,
536                session,
537            )
538            .await?;
539
540        // The stream is drained, so every builder has seen every chunk. Bytes now move from being
541        // buffered in memory to being written out below, so the reservation is released here
542        // rather than left to drop at the end of the function.
543        let indexes = {
544            let mut state = state.lock();
545            state.buffered.take();
546            std::mem::take(&mut state.indexes)
547        };
548
549        let mut index_layouts = Vec::with_capacity(indexes.len());
550        let mut specs = Vec::with_capacity(indexes.len());
551        for index in indexes {
552            let vtable = index.vtable.clone();
553            let partition_len = index.partition_len;
554            let declared = index.index_dtype.clone();
555            // An index whose every partition declined leaves no trace: no child, no spec, and no
556            // sequence pointer, since the splits below are what allocate one.
557            let Some(IndexContent {
558                stream: content,
559                options,
560                dtype: index_dtype,
561                partition_rows,
562                declined,
563            }) = IndexContent::assemble(vtable.id(), declared, index.finish()?)?
564            else {
565                trace!(index = %vtable.id(), "index builder declined, writing no child");
566                continue;
567            };
568
569            // Each index child gets its own (stream pointer, eof) pair, all ordered after the data
570            // segments. A chunking index strategy keeps each partition in its own segments, which
571            // is what lets a probe of one partition skip the others' bytes.
572            let content_ptr = eof.split_off();
573            let child_eof = eof.split_off();
574            let layout = self
575                .index
576                .write_stream(
577                    ctx.clone(),
578                    Arc::clone(&segment_sink),
579                    content.sequenced(content_ptr),
580                    child_eof,
581                    session,
582                )
583                .await?;
584
585            // The child has consumed the stream, so every partition's row count is final.
586            let partitioning =
587                partition_len.map(|len| partitioning(len, &partition_rows.lock(), declined));
588            specs.push(IndexSpec::new(vtable, options, index_dtype, partitioning));
589            index_layouts.push(layout);
590        }
591
592        // Every builder declined, so there is nothing to wrap. The data layout is already written
593        // and stands on its own, so hand it back as if no index had been configured.
594        if index_layouts.is_empty() {
595            return Ok(data_layout);
596        }
597
598        Ok(IndexedLayout::try_new(data_layout, index_layouts, specs)?.into_layout())
599    }
600}
vortex_layout/layouts/indexed/writer.rs

vortex_layout/layouts/indexed/
writer.rs
