common_datasource/
parquet_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
15use arrow::datatypes::{DataType, SchemaRef};
16use arrow::record_batch::RecordBatch;
17use bytes::Bytes;
18use futures::future::BoxFuture;
19use object_store::{ObjectStore, Writer};
20use parquet::arrow::ArrowWriter;
21use parquet::arrow::async_writer::AsyncFileWriter;
22use parquet::basic::{Compression, Encoding, ZstdLevel};
23use parquet::errors::ParquetError;
24use parquet::file::properties::WriterProperties;
25use parquet::schema::types::ColumnPath;
26use snafu::{IntoError, ResultExt, ensure};
27use tokio_util::sync::CancellationToken;
28
29use crate::DEFAULT_WRITE_BUFFER_SIZE;
30use crate::error::{self, Result};
31use crate::packed_writer::PackedTableWriter;
32
33/// Destination creation policy; conditional failures never authorize path deletion.
34#[derive(Clone, Copy, Debug, PartialEq, Eq)]
35pub enum ParquetCreationPolicy {
36    Overwrite,
37    IfNotExists,
38}
39
40/// Limits for one Parquet file. Flush thresholds are not hard memory caps.
41#[derive(Clone, Copy, Debug)]
42pub struct ParquetWriterLimits {
43    /// Maximum rows in a row group.
44    pub row_group_rows: usize,
45    /// Flush when encoder memory or encoded size reaches this threshold.
46    pub flush_threshold_bytes: usize,
47    /// Maximum row groups retained in the file footer.
48    pub max_row_groups: usize,
49}
50
51/// Encodes batches on a blocking runtime and writes one object-store file.
52/// Callers must await each operation before aborting; dropping an in-flight
53/// operation can leave a filesystem worker running after cleanup.
54pub struct ParquetFileWriter {
55    encoder: Option<ArrowWriter<Vec<u8>>>,
56    sink: ParquetSink,
57    rows: u64,
58    store: ObjectStore,
59    path: String,
60    limits: Option<ParquetWriterLimits>,
61    creation: ParquetCreationPolicy,
62    close_started: bool,
63}
64
65enum ParquetSink {
66    Object(Writer),
67    Packed(Box<PackedTableWriter>),
68}
69
70impl ParquetFileWriter {
71    /// Open a file using COPY's encoding settings. Destination ownership and
72    /// overwrite policy belong to the caller. None preserves Parquet's defaults.
73    pub async fn open(
74        schema: SchemaRef,
75        store: ObjectStore,
76        path: &str,
77        concurrency: usize,
78        limits: Option<ParquetWriterLimits>,
79    ) -> Result<Self> {
80        Self::open_with_creation(
81            schema,
82            store,
83            path,
84            concurrency,
85            limits,
86            ParquetCreationPolicy::Overwrite,
87        )
88        .await
89    }
90
91    /// Open with a conditional policy only when the caller verified backend support.
92    pub async fn open_with_creation(
93        schema: SchemaRef,
94        store: ObjectStore,
95        path: &str,
96        concurrency: usize,
97        limits: Option<ParquetWriterLimits>,
98        creation: ParquetCreationPolicy,
99    ) -> Result<Self> {
100        let encoder = build_encoder(schema, limits, path)?;
101        let sink = store
102            .writer_with(path)
103            .concurrent(concurrency)
104            .chunk(DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize)
105            .if_not_exists(creation == ParquetCreationPolicy::IfNotExists)
106            .await
107            .context(error::WriteObjectSnafu { path })?;
108        Ok(Self {
109            encoder: Some(encoder),
110            sink: ParquetSink::Object(sink),
111            rows: 0,
112            store,
113            path: path.to_owned(),
114            limits,
115            creation,
116            close_started: false,
117        })
118    }
119
120    /// Encode into a request-owned pack or standalone destination.
121    pub fn open_packed(
122        schema: SchemaRef,
123        store: ObjectStore,
124        path: &str,
125        limits: Option<ParquetWriterLimits>,
126        packed: PackedTableWriter,
127    ) -> Result<Self> {
128        Ok(Self {
129            encoder: Some(build_encoder(schema, limits, path)?),
130            sink: ParquetSink::Packed(Box::new(packed)),
131            rows: 0,
132            store,
133            path: path.into(),
134            limits,
135            creation: ParquetCreationPolicy::IfNotExists,
136            close_started: false,
137        })
138    }
139
140    /// Write a batch, enforcing file limits across batch and row-group boundaries.
141    pub async fn write(
142        &mut self,
143        batch: RecordBatch,
144        cancellation: Option<&CancellationToken>,
145    ) -> Result<()> {
146        let mut offset = 0;
147        while offset < batch.num_rows() {
148            check_cancelled(cancellation)?;
149            let mut encoder = self.encoder.take().ok_or_else(|| {
150                error::WriteParquetSnafu { path: &self.path }
151                    .into_error(ParquetError::General("Parquet writer is closed".into()))
152            })?;
153            let len = self.limits.map_or(batch.num_rows() - offset, |limits| {
154                (limits.row_group_rows - encoder.in_progress_rows()).min(batch.num_rows() - offset)
155            });
156            let slice = batch.slice(offset, len);
157            let limits = self.limits;
158            let path = self.path.clone();
159            let (encoder, bytes) = common_runtime::spawn_blocking_global(move || {
160                if let Some(limits) = limits {
161                    ensure!(
162                        encoder.flushed_row_groups().len() < limits.max_row_groups,
163                        error::ParquetWriterResourceSnafu {
164                            reason: "Parquet row-group metadata budget exceeded"
165                        }
166                    );
167                }
168                encoder
169                    .write(&slice)
170                    .context(error::WriteParquetSnafu { path: &path })?;
171                if limits.is_some_and(|limits| {
172                    encoder.memory_size() >= limits.flush_threshold_bytes
173                        || encoder.in_progress_size() >= limits.flush_threshold_bytes
174                }) {
175                    encoder
176                        .flush()
177                        .context(error::WriteParquetSnafu { path: &path })?;
178                }
179                // Draining preserves ArrowWriter's cumulative file offsets.
180                let bytes = std::mem::take(encoder.inner_mut());
181                Ok::<_, error::Error>((encoder, bytes))
182            })
183            .await
184            .context(error::JoinHandleSnafu)??;
185            self.encoder = Some(encoder);
186            check_cancelled(cancellation)?;
187            self.write_bytes(bytes).await?;
188            check_cancelled(cancellation)?;
189            self.rows += len as u64;
190            offset += len;
191        }
192        Ok(())
193    }
194
195    async fn write_bytes(&mut self, bytes: Vec<u8>) -> Result<()> {
196        let bytes = Bytes::from(bytes);
197        let sink = match &mut self.sink {
198            ParquetSink::Packed(packed) => return packed.write(bytes).await,
199            ParquetSink::Object(sink) => sink,
200        };
201        let chunk = DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize;
202        // Slices retain the complete encoded allocation until its last submission.
203        for offset in (0..bytes.len()).step_by(chunk) {
204            sink.write(bytes.slice(offset..(offset + chunk).min(bytes.len())))
205                .await
206                .context(error::WriteObjectSnafu { path: &self.path })?;
207        }
208        Ok(())
209    }
210
211    /// Write the footer and close the file. Retains the handle for cleanup on error.
212    pub async fn finish(&mut self, cancellation: Option<&CancellationToken>) -> Result<()> {
213        check_cancelled(cancellation)?;
214        let mut encoder = self.encoder.take().ok_or_else(|| {
215            error::WriteParquetSnafu { path: &self.path }
216                .into_error(ParquetError::General("Parquet writer is closed".into()))
217        })?;
218        let path = self.path.clone();
219        let bytes = common_runtime::spawn_blocking_global(move || {
220            encoder
221                .finish()
222                .context(error::WriteParquetSnafu { path: &path })?;
223            Ok::<_, error::Error>(std::mem::take(encoder.inner_mut()))
224        })
225        .await
226        .context(error::JoinHandleSnafu)??;
227        self.write_bytes(bytes).await?;
228        check_cancelled(cancellation)?;
229        let sink = match &mut self.sink {
230            ParquetSink::Packed(packed) => return packed.finish(self.rows, cancellation).await,
231            ParquetSink::Object(sink) => sink,
232        };
233        self.close_started = true;
234        sink.close()
235            .await
236            .context(error::WriteObjectSnafu { path: &self.path })?;
237        check_cancelled(cancellation)?;
238        Ok(())
239    }
240
241    /// Abort after all in-flight operations complete. Preserve ambiguous commits;
242    /// conditional callers delegate cleanup exclusively to the backend.
243    pub async fn abort(mut self) -> Result<()> {
244        let result = match &mut self.sink {
245            ParquetSink::Packed(packed) => return packed.abort().await,
246            ParquetSink::Object(sink) => sink.abort().await,
247        };
248        if self.creation == ParquetCreationPolicy::Overwrite
249            && result.as_ref().is_err_and(|error| {
250                error.kind() == object_store::ErrorKind::Unsupported
251                    && (!self.close_started
252                        || object_store::secure_fs::is_unsynced_overwrite_abort(error))
253            })
254        {
255            let store = self.store.clone();
256            let path = self.path.clone();
257            // Secure filesystem writers require dropping the handle before deletion.
258            drop(self);
259            store
260                .delete(&path)
261                .await
262                .context(error::WriteObjectSnafu { path })?;
263        } else {
264            result.context(error::WriteObjectSnafu { path: &self.path })?;
265        }
266        Ok(())
267    }
268}
269
270fn build_encoder(
271    schema: SchemaRef,
272    limits: Option<ParquetWriterLimits>,
273    path: &str,
274) -> Result<ArrowWriter<Vec<u8>>> {
275    let mut props = WriterProperties::builder()
276        .set_compression(Compression::ZSTD(ZstdLevel::default()))
277        .set_statistics_truncate_length(None)
278        .set_column_index_truncate_length(None);
279    if let Some(limits) = limits {
280        ensure!(
281            limits.row_group_rows > 0
282                && limits.flush_threshold_bytes > 0
283                && limits.max_row_groups > 0,
284            error::InvalidParquetWriterLimitsSnafu
285        );
286        props = props
287            .set_max_row_group_row_count(Some(limits.row_group_rows))
288            .set_max_row_group_bytes(None);
289    }
290    for field in schema.fields() {
291        if matches!(field.data_type(), DataType::Timestamp(_, _)) {
292            let column = ColumnPath::new(vec![field.name().clone()]);
293            props = props
294                .set_column_dictionary_enabled(column.clone(), false)
295                .set_column_encoding(column, Encoding::DELTA_BINARY_PACKED);
296        }
297    }
298    let encoder = ArrowWriter::try_new(Vec::new(), schema, Some(props.build()))
299        .context(error::WriteParquetSnafu { path })?;
300    Ok(encoder)
301}
302
303fn check_cancelled(cancellation: Option<&CancellationToken>) -> Result<()> {
304    ensure!(
305        cancellation.is_none_or(|token| !token.is_cancelled()),
306        error::ParquetWriteCancelledSnafu
307    );
308    Ok(())
309}
310
311/// Bridges opendal [Writer] with parquet [AsyncFileWriter].
312pub struct AsyncWriter {
313    inner: Writer,
314}
315
316impl AsyncWriter {
317    /// Create a [`AsyncWriter`] by given [`Writer`].
318    pub fn new(writer: Writer) -> Self {
319        Self { inner: writer }
320    }
321}
322
323impl AsyncFileWriter for AsyncWriter {
324    fn write(&mut self, bs: Bytes) -> BoxFuture<'_, parquet::errors::Result<()>> {
325        Box::pin(async move {
326            self.inner
327                .write(bs)
328                .await
329                .map_err(|err| ParquetError::External(Box::new(err)))
330        })
331    }
332
333    fn complete(&mut self) -> BoxFuture<'_, parquet::errors::Result<()>> {
334        Box::pin(async move {
335            self.inner
336                .close()
337                .await
338                .map(|_| ())
339                .map_err(|err| ParquetError::External(Box::new(err)))
340        })
341    }
342}
343
344#[cfg(test)]
345mod tests {
346    use std::sync::Arc;
347
348    use arrow::array::{Int64Array, TimestampMillisecondArray};
349    use common_error::ext::ErrorExt;
350    use common_error::status_code::StatusCode;
351    use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
352    use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
353
354    use super::*;
355    use crate::file_format::parquet::stream_to_parquet;
356
357    fn batch() -> RecordBatch {
358        RecordBatch::try_from_iter([
359            (
360                "value",
361                Arc::new(Int64Array::from(vec![Some(1), None, Some(3), Some(4)]))
362                    as arrow::array::ArrayRef,
363            ),
364            (
365                "ts",
366                Arc::new(TimestampMillisecondArray::from(vec![1, 2, 3, 4])),
367            ),
368        ])
369        .unwrap()
370    }
371
372    async fn read(store: &ObjectStore, path: &str) -> ParquetRecordBatchReaderBuilder<Bytes> {
373        ParquetRecordBatchReaderBuilder::try_new(store.read(path).await.unwrap().to_bytes())
374            .unwrap()
375    }
376
377    #[tokio::test]
378    async fn abort_preserves_collisions_and_ambiguous_commits() {
379        use object_store::layers::mock::{MockLayerBuilder, MockWriterFactory, oio};
380
381        struct AmbiguousCommit(oio::Writer);
382        impl oio::Write for AmbiguousCommit {
383            async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> {
384                self.0.write(bytes).await
385            }
386            async fn close(
387                &mut self,
388            ) -> object_store::Result<object_store::layers::mock::Metadata> {
389                self.0.close().await?;
390                Err(object_store::Error::new(
391                    object_store::ErrorKind::Unexpected,
392                    "lost close reply",
393                ))
394            }
395            async fn abort(&mut self) -> object_store::Result<()> {
396                Err(object_store::Error::new(
397                    object_store::ErrorKind::Unsupported,
398                    "cannot abort",
399                ))
400            }
401        }
402
403        for (creation, existing) in [
404            (ParquetCreationPolicy::Overwrite, false),
405            (ParquetCreationPolicy::IfNotExists, false),
406            (ParquetCreationPolicy::IfNotExists, true),
407        ] {
408            let directory = common_test_util::temp_dir::create_temp_dir("conditional_parquet");
409            let store = object_store::secure_fs::SecureFsRoot::open(directory.path())
410                .unwrap()
411                .build_operator();
412            let path = "conditional.parquet";
413            if existing {
414                store.write(path, "original").await.unwrap();
415            }
416            let factory: MockWriterFactory =
417                Arc::new(|_, _, writer| Box::new(AmbiguousCommit(writer)));
418            let store = store.layer(
419                MockLayerBuilder::default()
420                    .writer_factory(factory)
421                    .build()
422                    .unwrap(),
423            );
424            let mut writer = ParquetFileWriter::open_with_creation(
425                batch().schema(),
426                store.clone(),
427                path,
428                1,
429                None,
430                creation,
431            )
432            .await
433            .unwrap();
434            writer.write(batch(), None).await.unwrap();
435            assert!(writer.finish(None).await.is_err());
436            assert!(writer.abort().await.is_err());
437            if existing {
438                assert_eq!(
439                    store.read(path).await.unwrap().to_bytes(),
440                    Bytes::from_static(b"original")
441                );
442            } else {
443                assert_eq!(
444                    read(&store, path)
445                        .await
446                        .metadata()
447                        .file_metadata()
448                        .num_rows(),
449                    4
450                );
451            }
452        }
453    }
454
455    #[tokio::test]
456    async fn overwrite_abort_deletes_after_unsynced_close() {
457        use object_store::layers::mock::{MockLayerBuilder, MockWriterFactory, oio};
458
459        struct FailedClose(oio::Writer);
460        impl oio::Write for FailedClose {
461            async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> {
462                self.0.write(bytes).await
463            }
464            async fn close(
465                &mut self,
466            ) -> object_store::Result<object_store::layers::mock::Metadata> {
467                Err(object_store::Error::new(
468                    object_store::ErrorKind::Unexpected,
469                    "close failed before sync",
470                ))
471            }
472            async fn abort(&mut self) -> object_store::Result<()> {
473                self.0.abort().await
474            }
475        }
476
477        let directory = common_test_util::temp_dir::create_temp_dir("unsynced_parquet");
478        let store = object_store::secure_fs::SecureFsRoot::open(directory.path())
479            .unwrap()
480            .build_operator();
481        let path = "partial.parquet";
482        let factory: MockWriterFactory = Arc::new(|_, _, writer| Box::new(FailedClose(writer)));
483        let store = store.layer(
484            MockLayerBuilder::default()
485                .writer_factory(factory)
486                .build()
487                .unwrap(),
488        );
489        let mut writer = ParquetFileWriter::open(batch().schema(), store.clone(), path, 1, None)
490            .await
491            .unwrap();
492        writer.write(batch(), None).await.unwrap();
493        assert!(writer.finish(None).await.is_err());
494        assert!(store.exists(path).await.unwrap());
495        writer.abort().await.unwrap();
496        assert!(!store.exists(path).await.unwrap());
497    }
498
499    #[tokio::test]
500    async fn large_footer_flush_uses_bounded_submissions_in_one_parquet_stream() {
501        use std::sync::Mutex;
502
503        use object_store::layers::mock::{
504            MockLayerBuilder, MockWriterFactory, ObservedWriter, WriterObservations,
505        };
506        let directory = common_test_util::temp_dir::create_temp_dir("bounded_parquet");
507        let store = object_store::secure_fs::SecureFsRoot::open(directory.path())
508            .unwrap()
509            .build_operator();
510        let observations = Arc::new(Mutex::new(WriterObservations::default()));
511        let factory: MockWriterFactory = Arc::new({
512            let observations = observations.clone();
513            move |_, _, writer| Box::new(ObservedWriter::new(writer, observations.clone()))
514        });
515        let store = store.layer(
516            MockLayerBuilder::default()
517                .writer_factory(factory)
518                .build()
519                .unwrap(),
520        );
521        let mut state = 17u64;
522        let values = (0..600_000)
523            .map(|_| {
524                state ^= state << 13;
525                state ^= state >> 7;
526                state ^= state << 17;
527                state as i64
528            })
529            .collect::<Vec<_>>();
530        let array = Arc::new(Int64Array::from(values)) as arrow::array::ArrayRef;
531        let batch = RecordBatch::try_from_iter([("a", array.clone()), ("b", array)]).unwrap();
532        let mut writer =
533            ParquetFileWriter::open(batch.schema(), store.clone(), "large.parquet", 1, None)
534                .await
535                .unwrap();
536        // No storage-layer chunking: observe the application's actual submissions.
537        writer.sink = ParquetSink::Object(store.writer("large.parquet").await.unwrap());
538        writer.write(batch.clone(), None).await.unwrap();
539        writer.finish(None).await.unwrap();
540        let sizes = observations.lock().unwrap().write_sizes.clone();
541        let limit = DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize;
542        assert!(sizes.iter().sum::<usize>() > limit);
543        assert!(sizes.iter().all(|size| *size <= limit), "{sizes:?}");
544        let actual = read(&store, "large.parquet")
545            .await
546            .build()
547            .unwrap()
548            .collect::<std::result::Result<Vec<_>, _>>()
549            .unwrap();
550        assert_eq!(
551            arrow::compute::concat_batches(&batch.schema(), &actual).unwrap(),
552            batch
553        );
554    }
555
556    #[tokio::test]
557    async fn row_and_byte_limits_split_batches() {
558        let store = ObjectStore::new(object_store::services::Memory::default()).unwrap();
559        let batch = batch();
560        for (row_group_rows, flush_threshold_bytes, split) in [(2, usize::MAX, 3), (100, 1, 2)] {
561            let mut writer = ParquetFileWriter::open(
562                batch.schema(),
563                store.clone(),
564                "groups.parquet",
565                1,
566                Some(ParquetWriterLimits {
567                    row_group_rows,
568                    flush_threshold_bytes,
569                    max_row_groups: 2,
570                }),
571            )
572            .await
573            .unwrap();
574            writer.write(batch.slice(0, split), None).await.unwrap();
575            writer
576                .write(batch.slice(split, batch.num_rows() - split), None)
577                .await
578                .unwrap();
579            writer.finish(None).await.unwrap();
580            let reader = read(&store, "groups.parquet").await;
581            assert_eq!(reader.metadata().num_row_groups(), 2);
582            for group in reader.metadata().row_groups() {
583                assert_eq!(group.num_rows(), 2);
584                let encodings = group.column(1).encodings().collect::<Vec<_>>();
585                assert!(encodings.contains(&Encoding::DELTA_BINARY_PACKED));
586                assert!(!encodings.contains(&Encoding::RLE_DICTIONARY));
587                assert_eq!(
588                    group.column(0).compression(),
589                    Compression::ZSTD(ZstdLevel::default())
590                );
591            }
592            let actual = reader
593                .build()
594                .unwrap()
595                .collect::<std::result::Result<Vec<_>, _>>()
596                .unwrap();
597            assert_eq!(
598                arrow::compute::concat_batches(&batch.schema(), &actual).unwrap(),
599                batch
600            );
601        }
602    }
603
604    #[tokio::test]
605    async fn oversized_batch_stops_at_footer_limit_and_can_abort() {
606        let store = ObjectStore::new(object_store::services::Memory::default()).unwrap();
607        let batch = batch();
608        let mut writer = ParquetFileWriter::open(
609            batch.schema(),
610            store.clone(),
611            "limited.parquet",
612            1,
613            Some(ParquetWriterLimits {
614                row_group_rows: 2,
615                flush_threshold_bytes: usize::MAX,
616                max_row_groups: 1,
617            }),
618        )
619        .await
620        .unwrap();
621        let err = writer.write(batch, None).await.unwrap_err();
622        assert!(matches!(err, error::Error::ParquetWriterResource { .. }));
623        assert_eq!(err.status_code(), StatusCode::Suspended);
624        writer.abort().await.unwrap();
625        assert!(!store.exists("limited.parquet").await.unwrap());
626    }
627
628    #[tokio::test]
629    async fn copy_stream_preserves_values_and_empty_schema() {
630        let store = ObjectStore::new(object_store::services::Memory::default()).unwrap();
631        let batch = batch();
632        for (path, batches) in [
633            ("copy.parquet", vec![batch.clone()]),
634            ("empty.parquet", vec![]),
635        ] {
636            let stream = RecordBatchStreamAdapter::new(
637                batch.schema(),
638                futures::stream::iter(batches.clone().into_iter().map(Ok)),
639            );
640            assert_eq!(
641                stream_to_parquet(Box::pin(stream), store.clone(), path, 2)
642                    .await
643                    .unwrap(),
644                batches.iter().map(RecordBatch::num_rows).sum::<usize>()
645            );
646            let reader = read(&store, path).await;
647            assert_eq!(reader.schema().fields(), batch.schema().fields());
648            let actual = reader
649                .build()
650                .unwrap()
651                .collect::<std::result::Result<Vec<_>, _>>()
652                .unwrap();
653            assert_eq!(actual, batches);
654        }
655    }
656
657    #[tokio::test]
658    async fn failed_copy_preserves_untouched_securefs_destination() {
659        let directory = common_test_util::temp_dir::create_temp_dir("copy_existing");
660        let path = directory.path().join("existing.parquet");
661        std::fs::write(&path, b"original bytes").unwrap();
662        let access = crate::object_store::LocalFileAccess::sandboxed(directory.path()).unwrap();
663        let store = crate::object_store::build_backend_for_write(
664            &format!("{}/", directory.path().display()),
665            &Default::default(),
666            &access,
667        )
668        .await
669        .unwrap();
670        let stream = RecordBatchStreamAdapter::new(
671            batch().schema(),
672            futures::stream::iter(vec![Err(datafusion::error::DataFusionError::Execution(
673                "injected input failure".into(),
674            ))]),
675        );
676        let err = stream_to_parquet(Box::pin(stream), store, "existing.parquet", 1)
677            .await
678            .unwrap_err();
679        assert!(matches!(err, error::Error::ReadRecordBatch { .. }));
680        assert_eq!(std::fs::read(path).unwrap(), b"original bytes");
681    }
682
683    struct PausedFooterWriter {
684        inner: object_store::layers::mock::oio::Writer,
685        paused: bool,
686        started: Arc<tokio::sync::Notify>,
687        release: Arc<tokio::sync::Notify>,
688        closed: Arc<std::sync::atomic::AtomicBool>,
689    }
690
691    impl object_store::layers::mock::oio::Write for PausedFooterWriter {
692        async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> {
693            if !self.paused {
694                self.paused = true;
695                self.started.notify_one();
696                self.release.notified().await;
697            }
698            self.inner.write(bytes).await
699        }
700        async fn close(&mut self) -> object_store::Result<object_store::layers::mock::Metadata> {
701            self.closed.store(true, std::sync::atomic::Ordering::SeqCst);
702            self.inner.close().await
703        }
704        async fn abort(&mut self) -> object_store::Result<()> {
705            self.inner.abort().await
706        }
707    }
708
709    #[tokio::test]
710    async fn cancellation_during_footer_write_waits_then_aborts_before_close() {
711        use object_store::layers::mock::{MockLayerBuilder, MockWriterFactory};
712        let started = Arc::new(tokio::sync::Notify::new());
713        let release = Arc::new(tokio::sync::Notify::new());
714        let closed = Arc::new(std::sync::atomic::AtomicBool::new(false));
715        let factory: MockWriterFactory = Arc::new({
716            let (started, release, closed) = (started.clone(), release.clone(), closed.clone());
717            move |_, _, inner| {
718                Box::new(PausedFooterWriter {
719                    inner,
720                    paused: false,
721                    started: started.clone(),
722                    release: release.clone(),
723                    closed: closed.clone(),
724                })
725            }
726        });
727        let store = ObjectStore::new(object_store::services::Memory::default())
728            .unwrap()
729            .layer(
730                MockLayerBuilder::default()
731                    .writer_factory(factory)
732                    .build()
733                    .unwrap(),
734            );
735        let mut writer =
736            ParquetFileWriter::open(batch().schema(), store.clone(), "cancel.parquet", 1, None)
737                .await
738                .unwrap();
739        // Force footer bytes through the sink before its close operation.
740        writer.sink =
741            ParquetSink::Object(store.writer_with("cancel.parquet").chunk(1).await.unwrap());
742        let cancellation = CancellationToken::new();
743        let result = {
744            let finish = writer.finish(Some(&cancellation));
745            tokio::pin!(finish);
746            tokio::select! {
747                result = &mut finish => panic!("finished before footer write: {result:?}"),
748                _ = started.notified() => {},
749            }
750            cancellation.cancel();
751            assert!(futures::poll!(&mut finish).is_pending());
752            release.notify_one();
753            finish.await
754        };
755        assert!(matches!(
756            result,
757            Err(error::Error::ParquetWriteCancelled {})
758        ));
759        assert!(!closed.load(std::sync::atomic::Ordering::SeqCst));
760        writer.abort().await.unwrap();
761        assert!(!store.exists("cancel.parquet").await.unwrap());
762    }
763}
common_datasource/parquet_writer.rs

common_datasource/
parquet_writer.rs
