common_datasource/
packed_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//! Streaming destinations for one schema chunk of independent Parquet streams.
16
17use std::sync::Arc;
18
19use bytes::Bytes;
20use object_store::{ObjectStore, Writer};
21use snafu::{ResultExt, ensure};
22use tokio::sync::Mutex;
23use tokio_util::sync::CancellationToken;
24
25use crate::error::{self, Result};
26use crate::packed_snapshot::{ObjectKind, PACK_INDEX_FILE, PackIndex, PackObject, PackTable};
27
28pub const WRITE_BYTES: usize = 8 * 1024 * 1024;
29
30/// One request owns this sink. Finishers retain their encoder admission while
31/// waiting on the mutex, bounding both completed buffers and pending appends.
32pub type PackedWriterRef = Arc<Mutex<PackedWriter>>;
33
34#[derive(Clone, Copy)]
35struct UploadLimits {
36    part: usize,
37    object: u64,
38}
39
40impl UploadLimits {
41    fn for_store(store: &ObjectStore) -> Result<Self> {
42        let info = store.info();
43        let caps = info.capability();
44        let part = caps
45            .write_multi_max_size
46            .unwrap_or(WRITE_BYTES)
47            .min(WRITE_BYTES);
48        ensure!(
49            part > 0 && caps.write_multi_min_size.unwrap_or(0) <= part,
50            error::ParquetWriterResourceSnafu {
51                reason: "backend cannot accept bounded 8 MiB parts"
52            }
53        );
54        // OpenDAL exposes part size but not part count. These are the limits of
55        // its multipart implementations, including GCS's XML multipart writer.
56        let parts = match info.scheme() {
57            "s3" | "oss" | "gcs" => 10_000,
58            "azblob" => 50_000,
59            _ => u64::MAX,
60        };
61        let object = (caps.write_total_max_size.unwrap_or(usize::MAX) as u64)
62            .min(parts.saturating_mul(part as u64));
63        Ok(Self { part, object })
64    }
65}
66
67struct Upload {
68    writer: Writer,
69    store: ObjectStore,
70    path: String,
71    limits: UploadLimits,
72    length: u64,
73    conditional: bool,
74    close_started: bool,
75}
76
77impl Upload {
78    async fn open(store: &ObjectStore, path: String, limits: UploadLimits) -> Result<Self> {
79        let conditional = store.info().capability().write_with_if_not_exists;
80        if !conditional {
81            ensure!(
82                !store
83                    .exists(&path)
84                    .await
85                    .context(error::ReadObjectSnafu { path: &path })?,
86                error::InvalidPackedSnapshotSnafu {
87                    reason: format!("output already exists: {path}")
88                }
89            );
90        }
91        let writer = store
92            .writer_with(&path)
93            .chunk(limits.part)
94            .concurrent(1)
95            .if_not_exists(conditional)
96            .await
97            .context(error::WriteObjectSnafu { path: &path })?;
98        Ok(Self {
99            writer,
100            store: store.clone(),
101            path,
102            limits,
103            length: 0,
104            conditional,
105            close_started: false,
106        })
107    }
108
109    async fn write(&mut self, bytes: Bytes) -> Result<()> {
110        ensure!(
111            (bytes.len() as u64) <= self.limits.object.saturating_sub(self.length),
112            error::ParquetWriterResourceSnafu {
113                reason: "packed export object or multipart part-count limit exceeded"
114            }
115        );
116        for offset in (0..bytes.len()).step_by(self.limits.part) {
117            self.writer
118                .write(bytes.slice(offset..(offset + self.limits.part).min(bytes.len())))
119                .await
120                .context(error::WriteObjectSnafu { path: &self.path })?;
121        }
122        self.length += bytes.len() as u64;
123        Ok(())
124    }
125
126    async fn write_json_array<T: serde::Serialize>(
127        &mut self,
128        values: &[T],
129        token: &CancellationToken,
130    ) -> Result<()> {
131        self.write(Bytes::from_static(b"[")).await?;
132        for (i, value) in values.iter().enumerate() {
133            check_cancelled(Some(token))?;
134            if i > 0 {
135                self.write(Bytes::from_static(b",")).await?;
136            }
137            let bytes = serde_json::to_vec(value).map_err(|e| {
138                error::InvalidPackedSnapshotSnafu {
139                    reason: e.to_string(),
140                }
141                .build()
142            })?;
143            self.write(bytes.into()).await?;
144        }
145        self.write(Bytes::from_static(b"]")).await
146    }
147
148    async fn close(&mut self) -> Result<()> {
149        self.close_started = true;
150        self.writer
151            .close()
152            .await
153            .context(error::WriteObjectSnafu { path: &self.path })?;
154        Ok(())
155    }
156
157    async fn abort(mut self) -> Result<()> {
158        let result = self.writer.abort().await;
159        if !self.conditional
160            && result.as_ref().is_err_and(|e| {
161                e.kind() == object_store::ErrorKind::Unsupported
162                    && (!self.close_started
163                        || object_store::secure_fs::is_unsynced_overwrite_abort(e))
164            })
165        {
166            let store = self.store.clone();
167            let path = self.path.clone();
168            drop(self);
169            store
170                .delete(&path)
171                .await
172                .context(error::WriteObjectSnafu { path })?;
173        } else {
174            result.context(error::WriteObjectSnafu { path: &self.path })?;
175        }
176        Ok(())
177    }
178}
179
180/// Pack/index metadata is accounted separately from retained Arrow payloads.
181pub struct PackedWriter {
182    store: ObjectStore,
183    limits: UploadLimits,
184    pack: Option<Upload>,
185    index: PackIndex,
186    next_pack: usize,
187}
188
189impl PackedWriter {
190    pub fn new(store: ObjectStore) -> Result<PackedWriterRef> {
191        let limits = UploadLimits::for_store(&store)?;
192        Ok(Arc::new(Mutex::new(Self {
193            store,
194            limits,
195            pack: None,
196            index: PackIndex {
197                version: 1,
198                objects: vec![],
199                tables: vec![],
200            },
201            next_pack: 0,
202        })))
203    }
204
205    async fn close_pack(&mut self) -> Result<()> {
206        if let Some(pack) = &mut self.pack {
207            pack.close().await?;
208            self.index.objects.push(PackObject {
209                path: pack.path.clone(),
210                kind: ObjectKind::Pack,
211                length: pack.length,
212            });
213            self.pack = None;
214        }
215        Ok(())
216    }
217
218    async fn append(
219        &mut self,
220        name: String,
221        bytes: Bytes,
222        rows: u64,
223        token: Option<&CancellationToken>,
224    ) -> Result<()> {
225        check_cancelled(token)?;
226        if self
227            .pack
228            .as_ref()
229            .is_some_and(|p| bytes.len() as u64 > self.limits.object.saturating_sub(p.length))
230        {
231            self.close_pack().await?;
232        }
233        check_cancelled(token)?;
234        if self.pack.is_none() {
235            let path = format!("pack-{:06}.bin", self.next_pack);
236            self.next_pack += 1;
237            self.pack = Some(Upload::open(&self.store, path, self.limits).await?);
238        }
239        let pack = self.pack.as_mut().ok_or_else(|| {
240            error::InvalidPackedSnapshotSnafu {
241                reason: "missing pack writer",
242            }
243            .build()
244        })?;
245        let table = PackTable {
246            table_name: name,
247            object: pack.path.clone(),
248            offset: pack.length,
249            length: bytes.len() as u64,
250            row_count: rows,
251        };
252        pack.write(bytes).await?;
253        check_cancelled(token)?;
254        self.index.tables.push(table);
255        Ok(())
256    }
257
258    /// Publish the index only after every data object has closed successfully.
259    pub async fn finish(&mut self, token: &CancellationToken) -> Result<Vec<String>> {
260        check_cancelled(Some(token))?;
261        self.close_pack().await?;
262        check_cancelled(Some(token))?;
263        self.index
264            .tables
265            .sort_unstable_by(|a, b| a.table_name.cmp(&b.table_name));
266        self.index.validate()?;
267        let mut upload = Upload::open(&self.store, PACK_INDEX_FILE.into(), self.limits).await?;
268        let result = async {
269            upload
270                .write(Bytes::from_static(b"{\"version\":1,\"objects\":"))
271                .await?;
272            upload.write_json_array(&self.index.objects, token).await?;
273            upload.write(Bytes::from_static(b",\"tables\":")).await?;
274            upload.write_json_array(&self.index.tables, token).await?;
275            upload.write(Bytes::from_static(b"}")).await?;
276            check_cancelled(Some(token))?;
277            upload.close().await?;
278            check_cancelled(Some(token))
279        }
280        .await;
281        if let Err(error) = result {
282            if let Err(secondary) = upload.abort().await {
283                common_telemetry::warn!(secondary; "Failed to abort pack index");
284            }
285            return Err(error);
286        }
287        Ok(self
288            .index
289            .objects
290            .iter()
291            .map(|o| o.path.clone())
292            .chain([PACK_INDEX_FILE.into()])
293            .collect())
294    }
295
296    /// Call only after all admitted encoders have drained their I/O.
297    pub async fn abort(&mut self) -> Result<()> {
298        if let Some(pack) = self.pack.take() {
299            pack.abort().await?;
300        }
301        Ok(())
302    }
303}
304
305/// Buffers a small stream, or spills its prefix and continues the same stream.
306pub struct PackedTableWriter {
307    shared: PackedWriterRef,
308    name: String,
309    path: String,
310    buffer: Vec<u8>,
311    ordinary: bool,
312    standalone: Option<Upload>,
313}
314
315impl PackedTableWriter {
316    /// The caller retains encoder admission until finish/abort returns.
317    pub fn new(shared: PackedWriterRef, name: String, id: u32, ordinary: bool) -> Self {
318        Self {
319            shared,
320            name,
321            path: format!("table-{id}.parquet"),
322            buffer: if ordinary {
323                Vec::new()
324            } else {
325                Vec::with_capacity(WRITE_BYTES)
326            },
327            ordinary,
328            standalone: None,
329        }
330    }
331
332    pub(crate) async fn write(&mut self, bytes: Bytes) -> Result<()> {
333        if self.standalone.is_none()
334            && (self.ordinary || self.buffer.len().saturating_add(bytes.len()) > WRITE_BYTES)
335        {
336            let sink = self.shared.lock().await;
337            self.standalone =
338                Some(Upload::open(&sink.store, self.path.clone(), sink.limits).await?);
339        }
340        if let Some(upload) = &mut self.standalone {
341            if !self.buffer.is_empty() {
342                upload
343                    .write(std::mem::take(&mut self.buffer).into())
344                    .await?;
345            }
346            upload.write(bytes).await
347        } else {
348            self.buffer.extend_from_slice(&bytes);
349            Ok(())
350        }
351    }
352
353    pub(crate) async fn finish(
354        &mut self,
355        rows: u64,
356        token: Option<&CancellationToken>,
357    ) -> Result<()> {
358        check_cancelled(token)?;
359        if let Some(upload) = &mut self.standalone {
360            upload.close().await?;
361            check_cancelled(token)?;
362            let mut shared = self.shared.lock().await;
363            shared.index.objects.push(PackObject {
364                path: self.path.clone(),
365                kind: ObjectKind::Parquet,
366                length: upload.length,
367            });
368            shared.index.tables.push(PackTable {
369                table_name: self.name.clone(),
370                object: self.path.clone(),
371                offset: 0,
372                length: upload.length,
373                row_count: rows,
374            });
375        } else {
376            self.shared
377                .lock()
378                .await
379                .append(
380                    self.name.clone(),
381                    // OpenDAL may retain each slice until the pack part fills.
382                    // Release the encoder's 8 MiB capacity before queuing it.
383                    Bytes::from(std::mem::take(&mut self.buffer).into_boxed_slice()),
384                    rows,
385                    token,
386                )
387                .await?;
388        }
389        Ok(())
390    }
391
392    pub(crate) async fn abort(&mut self) -> Result<()> {
393        if let Some(upload) = self.standalone.take() {
394            upload.abort().await?;
395        }
396        Ok(())
397    }
398}
399
400fn check_cancelled(token: Option<&CancellationToken>) -> Result<()> {
401    ensure!(
402        token.is_none_or(|t| !t.is_cancelled()),
403        error::ParquetWriteCancelledSnafu
404    );
405    Ok(())
406}
407
408#[cfg(test)]
409mod tests {
410    use arrow::array::{ArrayRef, Int64Array, StringArray};
411    use arrow::record_batch::RecordBatch;
412    use futures::TryStreamExt;
413    use object_store::layers::mock::{Metadata, MockLayerBuilder, MockWriterFactory, oio};
414    use parquet::arrow::ParquetRecordBatchStreamBuilder;
415    use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
416
417    use super::*;
418    use crate::file_format::parquet::packed_reader::{PackReadWindows, PackedParquetReader};
419    use crate::parquet_writer::{ParquetFileWriter, ParquetWriterLimits};
420
421    #[tokio::test]
422    async fn independent_streams_empty_and_spilled_roundtrip() {
423        let directory = common_test_util::temp_dir::create_temp_dir("packed-writer");
424        let store = object_store::secure_fs::SecureFsRoot::open(directory.path())
425            .unwrap()
426            .build_operator();
427        let sizes = Arc::new(std::sync::Mutex::new(Vec::new()));
428        let factory: MockWriterFactory = Arc::new({
429            let sizes = sizes.clone();
430            move |_, _, writer| Box::new(Observe(writer, sizes.clone()))
431        });
432        let store = store.layer(
433            MockLayerBuilder::default()
434                .writer_factory(factory)
435                .build()
436                .unwrap(),
437        );
438        let shared = PackedWriter::new(store.clone()).unwrap();
439        let small = RecordBatch::try_from_iter([(
440            "tag",
441            Arc::new(StringArray::from(vec![Some(""), None, Some("value")])) as ArrayRef,
442        )])
443        .unwrap();
444        let mut state = 17u64;
445        let values = (0..600_000)
446            .map(|_| {
447                state ^= state << 13;
448                state ^= state >> 7;
449                state ^= state << 17;
450                state as i64
451            })
452            .collect::<Vec<_>>();
453        let random = Arc::new(Int64Array::from(values)) as ArrayRef;
454        let large = RecordBatch::try_from_iter([("a", random.clone()), ("b", random)]).unwrap();
455        let batches = [small.clone(), RecordBatch::new_empty(large.schema()), large];
456        for (id, batch) in batches.iter().enumerate() {
457            let sink =
458                PackedTableWriter::new(shared.clone(), format!("table{id}"), id as u32, false);
459            let mut writer = ParquetFileWriter::open_packed(
460                batch.schema(),
461                store.clone(),
462                "unused",
463                Some(ParquetWriterLimits {
464                    row_group_rows: 8192,
465                    flush_threshold_bytes: WRITE_BYTES,
466                    max_row_groups: 4096,
467                }),
468                sink,
469            )
470            .unwrap();
471            writer.write(batch.clone(), None).await.unwrap();
472            writer.finish(None).await.unwrap();
473        }
474        let inventory = shared
475            .lock()
476            .await
477            .finish(&CancellationToken::new())
478            .await
479            .unwrap();
480        let index: PackIndex =
481            serde_json::from_slice(&store.read(PACK_INDEX_FILE).await.unwrap().to_bytes()).unwrap();
482        index
483            .validate_membership(["table0", "table1", "table2"])
484            .unwrap();
485        assert_eq!(inventory.len(), 3);
486        assert!(sizes.lock().unwrap().iter().all(|n| *n <= WRITE_BYTES));
487        assert_eq!(
488            index
489                .objects
490                .iter()
491                .filter(|o| o.kind == ObjectKind::Pack)
492                .count(),
493            1
494        );
495        assert!(
496            index
497                .objects
498                .iter()
499                .any(|o| o.kind == ObjectKind::Parquet && o.length > WRITE_BYTES as u64)
500        );
501        let windows = PackReadWindows::new(store.clone());
502        for (entry, expected) in index.tables.iter().zip(batches) {
503            let object = index
504                .objects
505                .iter()
506                .find(|o| o.path == entry.object)
507                .unwrap();
508            let actual = if object.kind == ObjectKind::Pack {
509                let reader = PackedParquetReader::new(
510                    windows.clone(),
511                    object.path.clone(),
512                    object.length,
513                    entry.offset,
514                    entry.length,
515                )
516                .unwrap();
517                ParquetRecordBatchStreamBuilder::new(reader)
518                    .await
519                    .unwrap()
520                    .build()
521                    .unwrap()
522                    .try_collect::<Vec<_>>()
523                    .await
524                    .unwrap()
525            } else {
526                ParquetRecordBatchReaderBuilder::try_new(
527                    store.read(&object.path).await.unwrap().to_bytes(),
528                )
529                .unwrap()
530                .build()
531                .unwrap()
532                .collect::<std::result::Result<Vec<_>, _>>()
533                .unwrap()
534            };
535            assert_eq!(entry.row_count, expected.num_rows() as u64);
536            assert_eq!(
537                arrow::compute::concat_batches(&expected.schema(), &actual).unwrap(),
538                expected
539            );
540        }
541    }
542
543    struct Observe(oio::Writer, Arc<std::sync::Mutex<Vec<usize>>>);
544    impl oio::Write for Observe {
545        async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> {
546            self.1.lock().unwrap().push(bytes.len());
547            self.0.write(bytes).await
548        }
549        async fn close(&mut self) -> object_store::Result<Metadata> {
550            self.0.close().await
551        }
552        async fn abort(&mut self) -> object_store::Result<()> {
553            self.0.abort().await
554        }
555    }
556
557    #[tokio::test]
558    async fn rolls_at_table_boundary_and_aborts_standalone_at_part_limit() {
559        let directory = common_test_util::temp_dir::create_temp_dir("packed-limit");
560        let sizes = Arc::new(std::sync::Mutex::new(Vec::new()));
561        let factory: MockWriterFactory = Arc::new({
562            let sizes = sizes.clone();
563            move |_, _, writer| Box::new(Observe(writer, sizes.clone()))
564        });
565        let store = object_store::secure_fs::SecureFsRoot::open(directory.path())
566            .unwrap()
567            .build_operator()
568            .layer(
569                MockLayerBuilder::default()
570                    .writer_factory(factory)
571                    .build()
572                    .unwrap(),
573            );
574        let shared = PackedWriter::new(store.clone()).unwrap();
575        // Two data parts per object; no oversized fixture is needed.
576        shared.lock().await.limits = UploadLimits {
577            part: 512,
578            object: 1024,
579        };
580        for id in 0..3 {
581            let mut table = PackedTableWriter::new(shared.clone(), format!("t{id}"), id, false);
582            table.write(vec![id as u8; 600].into()).await.unwrap();
583            table.finish(0, None).await.unwrap();
584        }
585        let mut table = PackedTableWriter::new(shared.clone(), "large".into(), 4, true);
586        table.write(vec![0; 1024].into()).await.unwrap();
587        assert!(matches!(
588            table.write(Bytes::from_static(b"x")).await,
589            Err(error::Error::ParquetWriterResource { .. })
590        ));
591        table.abort().await.unwrap();
592        assert!(!store.exists("table-4.parquet").await.unwrap());
593        shared
594            .lock()
595            .await
596            .finish(&CancellationToken::new())
597            .await
598            .unwrap();
599        let index: PackIndex =
600            serde_json::from_slice(&store.read(PACK_INDEX_FILE).await.unwrap().to_bytes()).unwrap();
601        assert_eq!(index.objects.len(), 3);
602        assert!(
603            index
604                .tables
605                .iter()
606                .all(|t| t.offset == 0 && t.length == 600)
607        );
608        assert!(sizes.lock().unwrap().iter().all(|n| *n <= 512));
609    }
610
611    struct FailClose(oio::Writer);
612    impl oio::Write for FailClose {
613        async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> {
614            self.0.write(bytes).await
615        }
616        async fn close(&mut self) -> object_store::Result<Metadata> {
617            Err(object_store::Error::new(
618                object_store::ErrorKind::Unexpected,
619                "injected close failure",
620            ))
621        }
622        async fn abort(&mut self) -> object_store::Result<()> {
623            self.0.abort().await
624        }
625    }
626
627    #[tokio::test]
628    async fn close_failure_cancellation_and_collision_never_publish_index() {
629        for failure in ["pack-000000.bin", PACK_INDEX_FILE, "cancel", "collision"] {
630            let directory = common_test_util::temp_dir::create_temp_dir("packed-failure");
631            let store = object_store::secure_fs::SecureFsRoot::open(directory.path())
632                .unwrap()
633                .build_operator();
634            if failure == "collision" {
635                store.write("pack-000000.bin", "winner").await.unwrap();
636            }
637            let factory: MockWriterFactory = Arc::new(move |path, _, writer| {
638                if path.trim_start_matches('/') == failure {
639                    Box::new(FailClose(writer))
640                } else {
641                    writer
642                }
643            });
644            let store = store.layer(
645                MockLayerBuilder::default()
646                    .writer_factory(factory)
647                    .build()
648                    .unwrap(),
649            );
650            let shared = PackedWriter::new(store.clone()).unwrap();
651            let mut table = PackedTableWriter::new(shared.clone(), "t".into(), 0, false);
652            table.write(vec![0; 12].into()).await.unwrap();
653            table.finish(0, None).await.unwrap();
654            let token = CancellationToken::new();
655            if failure == "cancel" {
656                token.cancel();
657            }
658            assert!(shared.lock().await.finish(&token).await.is_err());
659            let _ = shared.lock().await.abort().await;
660            assert!(!store.exists(PACK_INDEX_FILE).await.unwrap());
661            if failure == "collision" {
662                assert_eq!(
663                    store.read("pack-000000.bin").await.unwrap().to_bytes(),
664                    Bytes::from_static(b"winner")
665                );
666            }
667        }
668    }
669}
common_datasource/packed_writer.rs

common_datasource/
packed_writer.rs
