mito2/sst/parquet/
parallel_writer.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Arrow to Parquet writer that encodes the columns of a row group in parallel.
16
17use std::mem;
18use std::sync::Arc;
19use std::time::{Duration, Instant};
20
21use bytes::Bytes;
22use common_runtime::JoinHandle;
23use datatypes::arrow::datatypes::SchemaRef;
24use datatypes::arrow::record_batch::RecordBatch;
25use parquet::arrow::ArrowWriter;
26use parquet::arrow::arrow_writer::{
27    ArrowColumnChunk, ArrowColumnWriter, ArrowRowGroupWriterFactory, compute_leaves,
28};
29use parquet::arrow::async_writer::AsyncFileWriter;
30use parquet::errors::{ParquetError, Result};
31use parquet::file::metadata::ParquetMetaData;
32use parquet::file::properties::WriterProperties;
33use parquet::file::writer::SerializedFileWriter;
34
35/// Runtime whose blocking pool encodes the columns of a writer.
36#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
37pub enum EncodeRuntime {
38    /// The global runtime.
39    #[default]
40    Global,
41    /// The compaction runtime.
42    Compact,
43}
44
45impl EncodeRuntime {
46    fn spawn<F, R>(self, f: F) -> JoinHandle<R>
47    where
48        F: FnOnce() -> R + Send + 'static,
49        R: Send + 'static,
50    {
51        match self {
52            EncodeRuntime::Global => common_runtime::spawn_blocking_global(f),
53            EncodeRuntime::Compact => common_runtime::spawn_blocking_compact(f),
54        }
55    }
56}
57
58/// Column writers of one top-level field: its index and its leaf writers.
59type FieldWriters = (usize, Vec<(usize, ArrowColumnWriter)>);
60
61/// Output of one encode task: column chunks by leaf index, and the time spent
62/// on each field.
63type EncodedFields = (Vec<(usize, ArrowColumnChunk)>, Vec<(usize, Duration)>);
64
65/// Writes Parquet like [`parquet::arrow::AsyncArrowWriter`], but encodes the
66/// columns of each row group on up to `parallelism` blocking tasks.
67///
68/// A row group encodes in the background while the next one is buffered, so at
69/// most two row groups of input are held in memory. Every column writer sees the
70/// same input as with [`ArrowWriter`], so the output is byte-for-byte identical.
71/// The bytes of a row group reach the async writer when the next row group
72/// starts or on [`Self::flush`], one row group later than with `AsyncArrowWriter`.
73pub struct ParallelArrowWriter<W> {
74    /// Assembles row groups and the footer into an in-memory buffer.
75    file_writer: SerializedFileWriter<Vec<u8>>,
76    /// Creates the column writers of each row group.
77    factory: Arc<ArrowRowGroupWriterFactory>,
78    arrow_schema: SchemaRef,
79    /// Top-level field of each leaf column.
80    leaf_fields: Vec<usize>,
81    /// Encode time of each field in the last row group, used to balance tasks.
82    field_costs: Vec<Duration>,
83    async_writer: W,
84    max_row_group_rows: usize,
85    parallelism: usize,
86    runtime: EncodeRuntime,
87    /// Slices of input buffered for the next row group.
88    buffered: Vec<RecordBatch>,
89    buffered_rows: usize,
90    /// Tasks encoding the previous row group.
91    encoding: Vec<JoinHandle<Result<EncodedFields>>>,
92    num_row_groups: usize,
93}
94
95impl<W: AsyncFileWriter> ParallelArrowWriter<W> {
96    /// Returns whether the writer can reproduce [`ArrowWriter`]'s row groups for
97    /// `props`, which holds only when row groups are limited by row count alone.
98    pub fn supports(props: &WriterProperties) -> bool {
99        props.max_row_group_row_count().is_some() && props.max_row_group_bytes().is_none()
100    }
101
102    /// Creates a writer. `props` must be [supported](Self::supports).
103    pub fn try_new(
104        async_writer: W,
105        arrow_schema: SchemaRef,
106        props: WriterProperties,
107        parallelism: usize,
108        runtime: EncodeRuntime,
109    ) -> Result<Self> {
110        let max_row_group_rows = match props.max_row_group_row_count() {
111            Some(rows) if Self::supports(&props) => rows.max(1),
112            _ => {
113                return Err(ParquetError::General(
114                    "parallel writer requires a row-count-only row group limit".to_string(),
115                ));
116            }
117        };
118        let (file_writer, factory) =
119            ArrowWriter::try_new(Vec::new(), arrow_schema.clone(), Some(props))?
120                .into_serialized_writer()?;
121        let schema = file_writer.schema_descr();
122        let leaf_fields = (0..schema.num_columns())
123            .map(|leaf| schema.get_column_root_idx(leaf))
124            .collect();
125        let num_fields = arrow_schema.fields().len();
126        Ok(Self {
127            file_writer,
128            factory: Arc::new(factory),
129            arrow_schema,
130            leaf_fields,
131            field_costs: vec![Duration::ZERO; num_fields],
132            async_writer,
133            max_row_group_rows,
134            parallelism: parallelism.max(1),
135            runtime,
136            buffered: Vec::new(),
137            buffered_rows: 0,
138            encoding: Vec::new(),
139            num_row_groups: 0,
140        })
141    }
142
143    /// Buffers `batch`, starting to encode each row group once it is full.
144    pub async fn write(&mut self, batch: &RecordBatch) -> Result<()> {
145        let mut offset = 0;
146        while offset < batch.num_rows() {
147            let len = (self.max_row_group_rows - self.buffered_rows).min(batch.num_rows() - offset);
148            self.buffered.push(batch.slice(offset, len));
149            self.buffered_rows += len;
150            offset += len;
151            if self.buffered_rows == self.max_row_group_rows {
152                self.encode_buffered().await?;
153            }
154        }
155        Ok(())
156    }
157
158    /// Writes all buffered rows as a row group, like [`ArrowWriter::flush`].
159    pub async fn flush(&mut self) -> Result<()> {
160        if self.buffered_rows > 0 {
161            self.encode_buffered().await?;
162        }
163        self.finish_encoding().await
164    }
165
166    /// Writes the remaining rows and the footer, then completes the async writer.
167    pub async fn close(mut self) -> Result<ParquetMetaData> {
168        self.flush().await?;
169        let metadata = self.file_writer.finish()?;
170        self.write_buffered_bytes().await?;
171        self.async_writer.complete().await?;
172        Ok(metadata)
173    }
174
175    /// Discards buffered input and waits for encoders without writing a failed SST.
176    pub(crate) async fn abort(mut self) {
177        self.buffered.clear();
178        abort_encoding_tasks(&mut self.encoding).await;
179    }
180
181    /// Installs controlled tasks to test SST write-error cleanup independently of codec timing.
182    #[cfg(test)]
183    pub(crate) fn set_encoding_tasks_for_test(
184        &mut self,
185        tasks: Vec<JoinHandle<Result<EncodedFields>>>,
186    ) {
187        assert!(self.encoding.is_empty());
188        self.encoding = tasks;
189    }
190
191    /// Starts encoding the buffered rows once the previous row group is written.
192    async fn encode_buffered(&mut self) -> Result<()> {
193        self.finish_encoding().await?;
194
195        let batches: Arc<[RecordBatch]> = mem::take(&mut self.buffered).into();
196        self.buffered_rows = 0;
197        let writers = self.factory.create_column_writers(self.num_row_groups)?;
198        self.num_row_groups += 1;
199
200        // One array produces all leaves of a field, so a field is never split.
201        let mut fields: Vec<FieldWriters> = (0..self.arrow_schema.fields().len())
202            .map(|field| (field, Vec::new()))
203            .collect();
204        for (leaf, writer) in writers.into_iter().enumerate() {
205            fields[self.leaf_fields[leaf]].1.push((leaf, writer));
206        }
207        self.encoding = assign_fields(fields, &self.field_costs, self.parallelism)
208            .into_iter()
209            .map(|fields| {
210                let batches = batches.clone();
211                let schema = self.arrow_schema.clone();
212                self.runtime
213                    .spawn(move || encode_fields(&schema, &batches, fields))
214            })
215            .collect();
216        Ok(())
217    }
218
219    /// Waits for the row group being encoded and appends it to the file.
220    async fn finish_encoding(&mut self) -> Result<()> {
221        if self.encoding.is_empty() {
222            return Ok(());
223        }
224        let mut chunks = Vec::with_capacity(self.leaf_fields.len());
225        let mut tasks = mem::take(&mut self.encoding).into_iter();
226        while let Some(task) = tasks.next() {
227            let encoded = task
228                .await
229                .map_err(|e| ParquetError::External(Box::new(e)))
230                .and_then(|encoded| encoded);
231            let (task_chunks, costs) = match encoded {
232                Ok(encoded) => encoded,
233                Err(error) => {
234                    abort_encoding_tasks(tasks.as_mut_slice()).await;
235                    return Err(error);
236                }
237            };
238            chunks.extend(task_chunks);
239            for (field, cost) in costs {
240                self.field_costs[field] = cost;
241            }
242        }
243        chunks.sort_unstable_by_key(|(leaf, _)| *leaf);
244        let mut row_group = self.file_writer.next_row_group()?;
245        for (_, chunk) in chunks {
246            chunk.append_to_row_group(&mut row_group)?;
247        }
248        row_group.close()?;
249        self.write_buffered_bytes().await
250    }
251
252    /// Moves the bytes assembled so far to the async writer.
253    async fn write_buffered_bytes(&mut self) -> Result<()> {
254        let bytes = mem::take(self.file_writer.inner_mut());
255        if bytes.is_empty() {
256            return Ok(());
257        }
258        self.async_writer.write(Bytes::from(bytes)).await
259    }
260}
261
262/// Aborts queued blocking work and joins running work before its budget is released.
263async fn abort_encoding_tasks(tasks: &mut [JoinHandle<Result<EncodedFields>>]) {
264    for task in tasks.iter() {
265        task.abort();
266    }
267    for task in tasks {
268        // Cleanup failures must not replace the original encoding or source error.
269        let _ = task.await;
270    }
271}
272
273/// Splits `fields` into at most `parallelism` tasks, giving the most expensive
274/// field to the least loaded task first. Fields cost nothing before any row
275/// group is measured, which spreads them evenly.
276fn assign_fields(
277    mut fields: Vec<FieldWriters>,
278    costs: &[Duration],
279    parallelism: usize,
280) -> Vec<Vec<FieldWriters>> {
281    let num_tasks = parallelism.min(fields.len()).max(1);
282    fields.sort_by_key(|(field, _)| std::cmp::Reverse(costs[*field]));
283    let mut tasks: Vec<(Duration, Vec<FieldWriters>)> = (0..num_tasks)
284        .map(|_| (Duration::ZERO, Vec::new()))
285        .collect();
286    for field in fields {
287        let key = |(load, fields): &(Duration, Vec<FieldWriters>)| (*load, fields.len());
288        let mut lightest = 0;
289        for task in 1..tasks.len() {
290            if key(&tasks[task]) < key(&tasks[lightest]) {
291                lightest = task;
292            }
293        }
294        tasks[lightest].0 += costs[field.0];
295        tasks[lightest].1.push(field);
296    }
297    tasks
298        .into_iter()
299        .map(|(_, fields)| fields)
300        .filter(|fields| !fields.is_empty())
301        .collect()
302}
303
304/// Encodes `fields` of all `batches` and closes their column writers.
305fn encode_fields(
306    schema: &SchemaRef,
307    batches: &[RecordBatch],
308    fields: Vec<FieldWriters>,
309) -> Result<EncodedFields> {
310    let mut chunks = Vec::new();
311    let mut costs = Vec::with_capacity(fields.len());
312    for (field, mut writers) in fields {
313        let start = Instant::now();
314        for batch in batches {
315            let leaves = compute_leaves(schema.field(field), batch.column(field))?;
316            if leaves.len() != writers.len() {
317                return Err(ParquetError::General(format!(
318                    "field {} has {} leaf arrays but {} column writers",
319                    schema.field(field).name(),
320                    leaves.len(),
321                    writers.len()
322                )));
323            }
324            for (leaf, (_, writer)) in leaves.iter().zip(writers.iter_mut()) {
325                writer.write(leaf)?;
326            }
327        }
328        for (leaf, writer) in writers {
329            chunks.push((leaf, writer.close()?));
330        }
331        costs.push((field, start.elapsed()));
332    }
333    Ok((chunks, costs))
334}
335
336#[cfg(test)]
337mod tests {
338    use std::sync::Arc;
339    use std::sync::atomic::{AtomicBool, Ordering};
340
341    use common_runtime::runtime::{BuilderBuild, RuntimeTrait};
342    use datatypes::arrow::array::{
343        ArrayRef, DictionaryArray, Float64Array, Int32Array, Int64Array, StringArray, StructArray,
344    };
345    use datatypes::arrow::datatypes::{DataType, Field, Fields, Schema, UInt32Type};
346    use parquet::basic::{Compression, ZstdLevel};
347    use tokio::sync::oneshot;
348    use tokio::time::timeout;
349
350    use super::*;
351    use crate::compaction::memory_manager::new_compaction_memory_manager;
352
353    fn schema() -> SchemaRef {
354        let nested = Fields::from(vec![
355            Field::new("a", DataType::Int32, true),
356            Field::new("b", DataType::Utf8, true),
357        ]);
358        Arc::new(Schema::new(vec![
359            Field::new("int", DataType::Int64, false),
360            Field::new("float", DataType::Float64, true),
361            Field::new(
362                "tag",
363                DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)),
364                false,
365            ),
366            Field::new("nested", DataType::Struct(nested), true),
367        ]))
368    }
369
370    fn batch(start: usize, len: usize) -> RecordBatch {
371        let rows = start..start + len;
372        let ints = Int64Array::from_iter_values(rows.clone().map(|i| i as i64 * 7));
373        let floats = Float64Array::from_iter(
374            rows.clone()
375                .map(|i| (i % 5 != 0).then_some((i % 13) as f64 / 3.0)),
376        );
377        let tags: DictionaryArray<UInt32Type> = rows
378            .clone()
379            .map(|i| format!("tag-{}", i % 3))
380            .collect::<Vec<_>>()
381            .iter()
382            .map(String::as_str)
383            .collect();
384        let nested = StructArray::from(vec![
385            (
386                Arc::new(Field::new("a", DataType::Int32, true)),
387                Arc::new(Int32Array::from_iter(
388                    rows.clone().map(|i| (i % 4 != 0).then_some(i as i32)),
389                )) as ArrayRef,
390            ),
391            (
392                Arc::new(Field::new("b", DataType::Utf8, true)),
393                Arc::new(StringArray::from_iter(
394                    rows.map(|i| (i % 6 != 0).then(|| format!("v{i}"))),
395                )) as ArrayRef,
396            ),
397        ]);
398        RecordBatch::try_new(
399            schema(),
400            vec![
401                Arc::new(ints),
402                Arc::new(floats),
403                Arc::new(tags),
404                Arc::new(nested),
405            ],
406        )
407        .unwrap()
408    }
409
410    fn props(row_group_rows: usize) -> WriterProperties {
411        WriterProperties::builder()
412            .set_compression(Compression::ZSTD(ZstdLevel::default()))
413            .set_max_row_group_row_count(Some(row_group_rows))
414            .set_data_page_row_count_limit(16)
415            .build()
416    }
417
418    /// One step of a write sequence: a batch of `len` rows, or a flush.
419    #[derive(Clone, Copy)]
420    enum Step {
421        Write(usize),
422        Flush,
423    }
424
425    fn write_serial(steps: &[Step], props: WriterProperties) -> Vec<u8> {
426        let mut buf = Vec::new();
427        let mut writer = ArrowWriter::try_new(&mut buf, schema(), Some(props)).unwrap();
428        let mut start = 0;
429        for step in steps {
430            match *step {
431                Step::Write(len) => {
432                    writer.write(&batch(start, len)).unwrap();
433                    start += len;
434                }
435                Step::Flush => writer.flush().unwrap(),
436            }
437        }
438        writer.close().unwrap();
439        buf
440    }
441
442    async fn write_parallel(
443        steps: &[Step],
444        props: WriterProperties,
445        parallelism: usize,
446    ) -> Vec<u8> {
447        let mut buf = Vec::new();
448        let mut writer = ParallelArrowWriter::try_new(
449            &mut buf,
450            schema(),
451            props,
452            parallelism,
453            EncodeRuntime::Global,
454        )
455        .unwrap();
456        let mut start = 0;
457        for step in steps {
458            match *step {
459                Step::Write(len) => {
460                    writer.write(&batch(start, len)).await.unwrap();
461                    start += len;
462                }
463                Step::Flush => writer.flush().await.unwrap(),
464            }
465        }
466        writer.close().await.unwrap();
467        buf
468    }
469
470    #[tokio::test]
471    async fn test_parallel_writer_matches_arrow_writer() {
472        use Step::*;
473        let cases: Vec<(&str, Vec<Step>)> = vec![
474            ("empty", vec![]),
475            ("partial row group", vec![Write(10)]),
476            ("exact row groups", vec![Write(64), Write(64)]),
477            (
478                "batches across row groups",
479                vec![Write(7), Write(100), Write(3), Write(250), Write(1)],
480            ),
481            ("batch spanning many row groups", vec![Write(1000)]),
482            (
483                "explicit flushes",
484                vec![Write(10), Flush, Flush, Write(70), Flush, Write(5)],
485            ),
486        ];
487        for (name, steps) in cases {
488            let expected = write_serial(&steps, props(64));
489            for parallelism in [1, 2, 3, 16] {
490                let actual = write_parallel(&steps, props(64), parallelism).await;
491                assert!(
492                    expected == actual,
493                    "case {name}, parallelism {parallelism}: {} bytes != {} bytes",
494                    expected.len(),
495                    actual.len()
496                );
497            }
498        }
499    }
500
501    #[rstest::rstest]
502    #[case::encode_error(false)]
503    #[case::encode_panic(true)]
504    #[tokio::test]
505    async fn test_encode_failure_drains_tasks_before_releasing_budget(#[case] panics: bool) {
506        let runtime: common_runtime::Runtime = common_runtime::Builder::default()
507            .worker_threads(1)
508            .max_blocking_threads(1)
509            .build()
510            .unwrap();
511        let mut writer = ParallelArrowWriter::try_new(
512            Vec::new(),
513            schema(),
514            props(64),
515            3,
516            EncodeRuntime::Compact,
517        )
518        .unwrap();
519        let memory = new_compaction_memory_manager(1024 * 1024);
520        let budget = memory.try_acquire(memory.limit_bytes()).unwrap();
521        let batches: Arc<[RecordBatch]> = vec![batch(0, 4)].into();
522        let retained_batches = Arc::downgrade(&batches);
523
524        let failed = runtime.spawn_blocking(move || -> Result<EncodedFields> {
525            if panics {
526                panic!("encode task failed");
527            }
528            Err(ParquetError::General("encode task failed".to_string()))
529        });
530        let (started_tx, started_rx) = oneshot::channel();
531        let (release_tx, release_rx) = oneshot::channel();
532        let column = writer.factory.create_column_writers(0).unwrap().remove(0);
533        let running_batches = batches.clone();
534        let running = runtime.spawn_blocking(move || {
535            started_tx.send(()).unwrap();
536            release_rx.blocking_recv().unwrap();
537            let _encoded =
538                encode_fields(&schema(), &running_batches, vec![(0, vec![(0, column)])])?;
539            Err(ParquetError::General("later encoding failure".to_string()))
540        });
541        // With one blocking thread, this also proves the first task has ended.
542        timeout(Duration::from_secs(10), started_rx)
543            .await
544            .unwrap()
545            .unwrap();
546        let queued_started = Arc::new(AtomicBool::new(false));
547        let queued_started_task = queued_started.clone();
548        let queued = runtime.spawn_blocking(move || {
549            queued_started_task.store(true, Ordering::SeqCst);
550            drop(batches);
551            Ok((Vec::new(), Vec::new()))
552        });
553        writer.encoding = vec![failed, running, queued];
554
555        // Model the compaction caller's guard: it is released when flush returns.
556        let mut finish = Box::pin(async {
557            let _budget = budget;
558            writer.flush().await
559        });
560        assert!(futures::poll!(finish.as_mut()).is_pending());
561        assert!(memory.try_acquire(memory.limit_bytes()).is_none());
562        assert!(retained_batches.upgrade().is_some());
563
564        release_tx.send(()).unwrap();
565        let error = timeout(Duration::from_secs(10), finish)
566            .await
567            .unwrap()
568            .unwrap_err();
569        if panics {
570            let ParquetError::External(source) = error else {
571                panic!("expected the original task panic, got {error}");
572            };
573            assert!(
574                source
575                    .downcast_ref::<common_runtime::JoinError>()
576                    .unwrap()
577                    .is_panic()
578            );
579        } else {
580            assert!(
581                matches!(error, ParquetError::General(message) if message == "encode task failed")
582            );
583        }
584        assert!(!queued_started.load(Ordering::SeqCst));
585        assert!(retained_batches.upgrade().is_none());
586        assert!(memory.try_acquire(memory.limit_bytes()).is_some());
587        assert!(writer.file_writer.flushed_row_groups().is_empty());
588        assert!(writer.async_writer.is_empty());
589    }
590
591    #[test]
592    fn test_parallel_writer_requires_row_count_limit() {
593        assert!(ParallelArrowWriter::<Vec<u8>>::supports(&props(64)));
594        let byte_limited = WriterProperties::builder()
595            .set_max_row_group_bytes(Some(1024))
596            .build();
597        assert!(!ParallelArrowWriter::<Vec<u8>>::supports(&byte_limited));
598        let unlimited = WriterProperties::builder()
599            .set_max_row_group_row_count(None)
600            .build();
601        assert!(!ParallelArrowWriter::<Vec<u8>>::supports(&unlimited));
602        assert!(
603            ParallelArrowWriter::try_new(Vec::new(), schema(), unlimited, 2, EncodeRuntime::Global)
604                .is_err()
605        );
606    }
607
608    #[test]
609    fn test_assign_fields_balances_measured_costs() {
610        let (file_writer, factory) = ArrowWriter::try_new(Vec::new(), schema(), Some(props(64)))
611            .unwrap()
612            .into_serialized_writer()
613            .unwrap();
614        let schema_descr = file_writer.schema_descr();
615        let mut fields: Vec<FieldWriters> = (0..4).map(|field| (field, Vec::new())).collect();
616        for (leaf, writer) in factory
617            .create_column_writers(0)
618            .unwrap()
619            .into_iter()
620            .enumerate()
621        {
622            fields[schema_descr.get_column_root_idx(leaf)]
623                .1
624                .push((leaf, writer));
625        }
626        // The nested field owns both struct leaves.
627        assert_eq!(2, fields[3].1.len());
628
629        let costs = [10, 9, 1, 1].map(Duration::from_millis);
630        let tasks = assign_fields(fields, &costs, 2);
631        let mut loads: Vec<_> = tasks
632            .iter()
633            .map(|task| {
634                task.iter()
635                    .map(|(field, _)| costs[*field])
636                    .sum::<Duration>()
637            })
638            .collect();
639        loads.sort();
640        assert_eq!(
641            vec![Duration::from_millis(10), Duration::from_millis(11)],
642            loads
643        );
644        let mut assigned: Vec<_> = tasks.iter().flatten().map(|(field, _)| *field).collect();
645        assigned.sort();
646        assert_eq!(vec![0, 1, 2, 3], assigned);
647    }
648}
mito2/sst/parquet/parallel_writer.rs

mito2/sst/parquet/
parallel_writer.rs
