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::{
414        Metadata, MockLayerBuilder, MockWriterFactory, ObservedWriter, WriterObservations, oio,
415    };
416    use parquet::arrow::ParquetRecordBatchStreamBuilder;
417    use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
418
419    use super::*;
420    use crate::file_format::parquet::packed_reader::{PackReadWindows, PackedParquetReader};
421    use crate::parquet_writer::{ParquetFileWriter, ParquetWriterLimits};
422
423    #[tokio::test]
424    async fn independent_streams_empty_and_spilled_roundtrip() {
425        let directory = common_test_util::temp_dir::create_temp_dir("packed-writer");
426        let store = object_store::secure_fs::SecureFsRoot::open(directory.path())
427            .unwrap()
428            .build_operator();
429        let observations = Arc::new(std::sync::Mutex::new(WriterObservations::default()));
430        let factory: MockWriterFactory = Arc::new({
431            let observations = observations.clone();
432            move |_, _, writer| Box::new(ObservedWriter::new(writer, observations.clone()))
433        });
434        let store = store.layer(
435            MockLayerBuilder::default()
436                .writer_factory(factory)
437                .build()
438                .unwrap(),
439        );
440        let shared = PackedWriter::new(store.clone()).unwrap();
441        let small = RecordBatch::try_from_iter([(
442            "tag",
443            Arc::new(StringArray::from(vec![Some(""), None, Some("value")])) as ArrayRef,
444        )])
445        .unwrap();
446        let mut state = 17u64;
447        let values = (0..600_000)
448            .map(|_| {
449                state ^= state << 13;
450                state ^= state >> 7;
451                state ^= state << 17;
452                state as i64
453            })
454            .collect::<Vec<_>>();
455        let random = Arc::new(Int64Array::from(values)) as ArrayRef;
456        let large = RecordBatch::try_from_iter([("a", random.clone()), ("b", random)]).unwrap();
457        let batches = [small.clone(), RecordBatch::new_empty(large.schema()), large];
458        for (id, batch) in batches.iter().enumerate() {
459            let sink =
460                PackedTableWriter::new(shared.clone(), format!("table{id}"), id as u32, false);
461            let mut writer = ParquetFileWriter::open_packed(
462                batch.schema(),
463                store.clone(),
464                "unused",
465                Some(ParquetWriterLimits {
466                    row_group_rows: 8192,
467                    flush_threshold_bytes: WRITE_BYTES,
468                    max_row_groups: 4096,
469                }),
470                sink,
471            )
472            .unwrap();
473            writer.write(batch.clone(), None).await.unwrap();
474            writer.finish(None).await.unwrap();
475        }
476        let inventory = shared
477            .lock()
478            .await
479            .finish(&CancellationToken::new())
480            .await
481            .unwrap();
482        let index: PackIndex =
483            serde_json::from_slice(&store.read(PACK_INDEX_FILE).await.unwrap().to_bytes()).unwrap();
484        index
485            .validate_membership(["table0", "table1", "table2"])
486            .unwrap();
487        assert_eq!(inventory.len(), 3);
488        assert!(
489            observations
490                .lock()
491                .unwrap()
492                .write_sizes
493                .iter()
494                .all(|n| *n <= WRITE_BYTES)
495        );
496        assert_eq!(
497            index
498                .objects
499                .iter()
500                .filter(|o| o.kind == ObjectKind::Pack)
501                .count(),
502            1
503        );
504        assert!(
505            index
506                .objects
507                .iter()
508                .any(|o| o.kind == ObjectKind::Parquet && o.length > WRITE_BYTES as u64)
509        );
510        let windows = PackReadWindows::new(store.clone());
511        for (entry, expected) in index.tables.iter().zip(batches) {
512            let object = index
513                .objects
514                .iter()
515                .find(|o| o.path == entry.object)
516                .unwrap();
517            let actual = if object.kind == ObjectKind::Pack {
518                let reader = PackedParquetReader::new(
519                    windows.clone(),
520                    object.path.clone(),
521                    object.length,
522                    entry.offset,
523                    entry.length,
524                )
525                .unwrap();
526                ParquetRecordBatchStreamBuilder::new(reader)
527                    .await
528                    .unwrap()
529                    .build()
530                    .unwrap()
531                    .try_collect::<Vec<_>>()
532                    .await
533                    .unwrap()
534            } else {
535                ParquetRecordBatchReaderBuilder::try_new(
536                    store.read(&object.path).await.unwrap().to_bytes(),
537                )
538                .unwrap()
539                .build()
540                .unwrap()
541                .collect::<std::result::Result<Vec<_>, _>>()
542                .unwrap()
543            };
544            assert_eq!(entry.row_count, expected.num_rows() as u64);
545            assert_eq!(
546                arrow::compute::concat_batches(&expected.schema(), &actual).unwrap(),
547                expected
548            );
549        }
550    }
551
552    #[tokio::test]
553    async fn rolls_at_table_boundary_and_aborts_standalone_at_part_limit() {
554        let directory = common_test_util::temp_dir::create_temp_dir("packed-limit");
555        let observations = Arc::new(std::sync::Mutex::new(WriterObservations::default()));
556        let factory: MockWriterFactory = Arc::new({
557            let observations = observations.clone();
558            move |_, _, writer| Box::new(ObservedWriter::new(writer, observations.clone()))
559        });
560        let store = object_store::secure_fs::SecureFsRoot::open(directory.path())
561            .unwrap()
562            .build_operator()
563            .layer(
564                MockLayerBuilder::default()
565                    .writer_factory(factory)
566                    .build()
567                    .unwrap(),
568            );
569        let shared = PackedWriter::new(store.clone()).unwrap();
570        // Two data parts per object; no oversized fixture is needed.
571        shared.lock().await.limits = UploadLimits {
572            part: 512,
573            object: 1024,
574        };
575        for id in 0..3 {
576            let mut table = PackedTableWriter::new(shared.clone(), format!("t{id}"), id, false);
577            table.write(vec![id as u8; 600].into()).await.unwrap();
578            table.finish(0, None).await.unwrap();
579        }
580        let mut table = PackedTableWriter::new(shared.clone(), "large".into(), 4, true);
581        table.write(vec![0; 1024].into()).await.unwrap();
582        assert!(matches!(
583            table.write(Bytes::from_static(b"x")).await,
584            Err(error::Error::ParquetWriterResource { .. })
585        ));
586        table.abort().await.unwrap();
587        assert!(!store.exists("table-4.parquet").await.unwrap());
588        shared
589            .lock()
590            .await
591            .finish(&CancellationToken::new())
592            .await
593            .unwrap();
594        let index: PackIndex =
595            serde_json::from_slice(&store.read(PACK_INDEX_FILE).await.unwrap().to_bytes()).unwrap();
596        assert_eq!(index.objects.len(), 3);
597        assert!(
598            index
599                .tables
600                .iter()
601                .all(|t| t.offset == 0 && t.length == 600)
602        );
603        assert!(
604            observations
605                .lock()
606                .unwrap()
607                .write_sizes
608                .iter()
609                .all(|n| *n <= 512)
610        );
611    }
612
613    struct FailClose(oio::Writer);
614    impl oio::Write for FailClose {
615        async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> {
616            self.0.write(bytes).await
617        }
618        async fn close(&mut self) -> object_store::Result<Metadata> {
619            Err(object_store::Error::new(
620                object_store::ErrorKind::Unexpected,
621                "injected close failure",
622            ))
623        }
624        async fn abort(&mut self) -> object_store::Result<()> {
625            self.0.abort().await
626        }
627    }
628
629    #[tokio::test]
630    async fn close_failure_cancellation_and_collision_never_publish_index() {
631        for failure in ["pack-000000.bin", PACK_INDEX_FILE, "cancel", "collision"] {
632            let directory = common_test_util::temp_dir::create_temp_dir("packed-failure");
633            let store = object_store::secure_fs::SecureFsRoot::open(directory.path())
634                .unwrap()
635                .build_operator();
636            if failure == "collision" {
637                store.write("pack-000000.bin", "winner").await.unwrap();
638            }
639            let factory: MockWriterFactory = Arc::new(move |path, _, writer| {
640                if path.trim_start_matches('/') == failure {
641                    Box::new(FailClose(writer))
642                } else {
643                    writer
644                }
645            });
646            let store = store.layer(
647                MockLayerBuilder::default()
648                    .writer_factory(factory)
649                    .build()
650                    .unwrap(),
651            );
652            let shared = PackedWriter::new(store.clone()).unwrap();
653            let mut table = PackedTableWriter::new(shared.clone(), "t".into(), 0, false);
654            table.write(vec![0; 12].into()).await.unwrap();
655            table.finish(0, None).await.unwrap();
656            let token = CancellationToken::new();
657            if failure == "cancel" {
658                token.cancel();
659            }
660            assert!(shared.lock().await.finish(&token).await.is_err());
661            let _ = shared.lock().await.abort().await;
662            assert!(!store.exists(PACK_INDEX_FILE).await.unwrap());
663            if failure == "collision" {
664                assert_eq!(
665                    store.read("pack-000000.bin").await.unwrap().to_bytes(),
666                    Bytes::from_static(b"winner")
667                );
668            }
669        }
670    }
671}
common_datasource/packed_writer.rs

common_datasource/
packed_writer.rs
