feat: export Metric snapshots with packed Parquet objects (#9382)

* feat: export Metric snapshots with packed Parquet objects Signed-off-by: jeremyhi <fengjiachun@gmail.com> * fix: reject empty packed export time ranges Signed-off-by: jeremyhi <fengjiachun@gmail.com> * test: cover cancellation during packed export I/O Signed-off-by: jeremyhi <fengjiachun@gmail.com> --------- Signed-off-by: jeremyhi <fengjiachun@gmail.com>
2026-10-26 21:56:17 +00:00 · 2026-09-28 15:22:02 +00:00
parent e2b7e5e0a8
commit 3c0e2a8d55
24 changed files with 1584 additions and 127 deletions
@@ -24,6 +24,7 @@ use common_datasource::file_format::csv::stream_to_csv;
 use common_datasource::file_format::json::stream_to_json;
 use common_datasource::file_format::parquet::stream_to_parquet;
 use common_datasource::object_store::build_backend_for_write_with_path;
+use common_datasource::packed_writer::PackedTableWriter;
 use common_datasource::parquet_writer::ParquetFileWriter;
 use common_query::Output;
 use common_recordbatch::adapter::DfRecordBatchStreamAdapter;
@@ -127,7 +128,7 @@ impl StatementExecutor {
        req: CopyTableRequest,
        query_ctx: QueryContextRef,
    ) -> Result<usize> {
-        self.copy_captured_table_to_managed(table, req, query_ctx, None)
+        self.copy_captured_table_to_managed(table, req, query_ctx, None, None)
            .await
    }

@@ -137,6 +138,7 @@ impl StatementExecutor {
        req: CopyTableRequest,
        query_ctx: QueryContextRef,
        managed: Option<(&ExportWriteBudget, &CancellationToken)>,
+        packed: Option<PackedTableWriter>,
    ) -> Result<usize> {
        let info = table.table_info();
        let table_ref = TableReference::full(&info.catalog_name, &info.schema_name, &info.name);
@@ -185,7 +187,7 @@ impl StatementExecutor {
        } = &req;

        debug!("Copy table: {table_id} to location: {location}");
-        self.copy_to_file_managed(&format, output, location, connection, managed)
+        self.copy_to_file_managed(&format, output, location, connection, managed, packed)
            .await
    }

@@ -196,7 +198,7 @@ impl StatementExecutor {
        location: &str,
        connection: &HashMap<String, String>,
    ) -> Result<usize> {
-        self.copy_to_file_managed(format, output, location, connection, None)
+        self.copy_to_file_managed(format, output, location, connection, None, None)
            .await
    }

@@ -207,6 +209,7 @@ impl StatementExecutor {
        location: &str,
        connection: &HashMap<String, String>,
        managed: Option<(&ExportWriteBudget, &CancellationToken)>,
+        packed: Option<PackedTableWriter>,
    ) -> Result<usize> {
        let output = if managed.is_none() {
            output
@@ -229,7 +232,15 @@ impl StatementExecutor {
            violated: format!("Expected filename, path: {location}"),
        })?;
        if let Some((budget, token)) = managed {
-            stream_to_managed_parquet(stream, backend.object_store, &filename, budget, token).await
+            stream_to_managed_parquet_with_packed(
+                stream,
+                backend.object_store,
+                &filename,
+                budget,
+                token,
+                packed,
+            )
+            .await
        } else {
            self.stream_to_file(stream, format, backend.object_store, &filename)
                .await
@@ -237,12 +248,24 @@ impl StatementExecutor {
    }
 }

+#[cfg(test)]
 pub(crate) async fn stream_to_managed_parquet(
+    stream: SendableRecordBatchStream,
+    store: ObjectStore,
+    path: &str,
+    budget: &ExportWriteBudget,
+    token: &CancellationToken,
+) -> Result<usize> {
+    stream_to_managed_parquet_with_packed(stream, store, path, budget, token, None).await
+}
+
+async fn stream_to_managed_parquet_with_packed(
    mut stream: SendableRecordBatchStream,
    store: ObjectStore,
    path: &str,
    budget: &ExportWriteBudget,
    token: &CancellationToken,
+    packed: Option<PackedTableWriter>,
 ) -> Result<usize> {
    use common_recordbatch::{RecordBatch, map_dictionary_to_values_schema};
    let original = stream.schema();
@@ -261,10 +284,21 @@ pub(crate) async fn stream_to_managed_parquet(
    } else {
        expanded_schema.clone()
    };
-    let mut writer =
-        ParquetFileWriter::open(output_schema.arrow_schema().clone(), store, path, 1, None)
-            .await
-            .context(error::WriteStreamToFileSnafu { path })?;
+    let packed_destination = packed.is_some();
+    let mut writer = if let Some(packed) = packed {
+        ParquetFileWriter::open_packed(
+            output_schema.arrow_schema().clone(),
+            store,
+            path,
+            Some(
+                crate::statement::export_logical_tables::LogicalTableExportLimits::default().writer,
+            ),
+            packed,
+        )
+    } else {
+        ParquetFileWriter::open(output_schema.arrow_schema().clone(), store, path, 1, None).await
+    }
+    .context(error::WriteStreamToFileSnafu { path })?;
    let mut started = false;
    let result = async {
        let mut rows = 0;
@@ -358,7 +392,9 @@ pub(crate) async fn stream_to_managed_parquet(
    .await;
    if result.is_err() {
        token.cancel();
-        if started && let Err(error) = writer.abort().await {
+        if (started || packed_destination)
+            && let Err(error) = writer.abort().await
+        {
            common_telemetry::warn!(error; "Failed to abort ordinary export file");
        }
    }
@@ -60,7 +60,7 @@ pub(crate) fn parse_parallelism_from_option_map(options: &HashMap<String, String
        .clamp(1, Semaphore::MAX_PERMITS)
 }

-/// Rejects import-only layouts before either database export path creates output.
+/// Rejects layouts unsupported by the per-table export path.
 pub(crate) fn validate_database_export_layout(options: &HashMap<String, String>) -> Result<()> {
    if let Some(layout) = options.get("metric_data_layout") {
        return error::InvalidCopyParameterSnafu {
@@ -20,6 +20,7 @@ use std::future::Future;

 use common_datasource::file_format::Format;
 use common_datasource::object_store::{FILE_SCHEMA, FS_SCHEMA, build_backend_for_write, parse_url};
+use common_datasource::packed_writer::{PackedTableWriter, PackedWriter};
 use common_meta::key::table_route::TableRouteValue;
 use futures::StreamExt;
 use futures::stream::FuturesUnordered;
@@ -70,7 +71,13 @@ impl StatementExecutor {
        req: CopyDatabaseRequest,
        tables: Vec<TableRef>,
    ) -> Result<PreparedDatabaseExport> {
-        validate_database_export_layout(&req.with)?;
+        if req
+            .with
+            .get("metric_data_layout")
+            .is_none_or(|v| v != "packed")
+        {
+            validate_database_export_layout(&req.with)?;
+        }
        validate_database_directory(&req.location)?;
        let format = Format::try_from(&req.with).context(error::ParseFileFormatSnafu)?;
        ensure!(
@@ -205,9 +212,27 @@ impl StatementExecutor {
        let req = &plan.request;
        let parallelism = parse_parallelism_from_option_map(&req.with);
        let budget = ExportWriteBudget::new(parallelism);
+        let packed = if req
+            .with
+            .get("metric_data_layout")
+            .is_some_and(|v| v == "packed")
+        {
+            let store =
+                build_backend_for_write(&req.location, &req.connection, &self.local_file_access)
+                    .await
+                    .context(error::BuildBackendSnafu)?;
+            Some(
+                PackedWriter::new(store).context(error::WriteStreamToFileSnafu {
+                    path: &req.location,
+                })?,
+            )
+        } else {
+            None
+        };
        let rows = run_database_export_jobs(plan.jobs, parallelism, cancellation, |job, token| {
            let ctx = ctx.clone();
            let budget = budget.clone();
+            let packed = packed.clone();
            async move {
                match job {
                    DatabaseExportJob::Metric(unit) => self
@@ -220,6 +245,7 @@ impl StatementExecutor {
                            &token,
                            ctx,
                            budget,
+                            packed,
                        )
                        .await
                        .map(|summary| summary.rows),
@@ -238,19 +264,57 @@ impl StatementExecutor {
                            limit: None,
                        };
                        let _permit = budget.writer(&token).await?;
+                        let destination = packed.map(|shared| {
+                            PackedTableWriter::new(shared, info.name.clone(), info.table_id(), true)
+                        });
                        self.copy_captured_table_to_managed(
                            table,
                            copy,
                            ctx,
                            Some((&budget, &token)),
+                            destination,
                        )
                        .await
                    }
                }
            }
        })
-        .await?;
-        Ok(DatabaseExportSummary { rows, output_files })
+        .await;
+        if let Some(packed) = packed {
+            let mut packed = packed.lock().await;
+            let result = match rows {
+                Ok(rows) => packed
+                    .finish(cancellation)
+                    .await
+                    .context(error::WriteStreamToFileSnafu {
+                        path: &req.location,
+                    })
+                    .and_then(|files| {
+                        Ok(DatabaseExportSummary {
+                            rows,
+                            output_files: files
+                                .into_iter()
+                                .map(|file| {
+                                    DatabaseExportFile::new(&req.location, &file, "")
+                                        .map(|f| f.location)
+                                })
+                                .collect::<Result<_>>()?,
+                        })
+                    }),
+                Err(error) => Err(error),
+            };
+            if result.is_err()
+                && let Err(error) = packed.abort().await
+            {
+                common_telemetry::warn!(error; "Failed to abort Metric pack");
+            }
+            result
+        } else {
+            Ok(DatabaseExportSummary {
+                rows: rows?,
+                output_files,
+            })
+        }
    }
 }

@@ -31,6 +31,7 @@ use arrow::datatypes::{DataType, SchemaRef};
 use arrow::downcast_dictionary_array;
 use arrow::record_batch::RecordBatch;
 use common_datasource::object_store::build_backend_for_write;
+use common_datasource::packed_writer::{PackedTableWriter, PackedWriterRef};
 use common_datasource::parquet_writer::{
    ParquetCreationPolicy, ParquetFileWriter, ParquetWriterLimits,
 };
@@ -108,6 +109,7 @@ pub struct LogicalTableExport {
 }

 pub(crate) struct LogicalTableProjection {
+    name: String,
    output: DatabaseExportFile,
    schema: SchemaRef,
    projection: Vec<usize>,
@@ -196,6 +198,7 @@ impl LogicalTableExport {
                    .insert(
                        info.table_id(),
                        LogicalTableProjection {
+                            name: name.clone(),
                            output: DatabaseExportFile::new(directory, name, ".parquet")?,
                            schema,
                            projection: indices,
@@ -322,6 +325,7 @@ impl StatementExecutor {
            &cancellation.child_token(),
            query_ctx,
            ExportWriteBudget::new(1),
+            None,
        )
        .await
    }
@@ -337,6 +341,7 @@ impl StatementExecutor {
        cancellation: &CancellationToken,
        query_ctx: QueryContextRef,
        budget: Arc<ExportWriteBudget>,
+        packed: Option<PackedWriterRef>,
    ) -> Result<LogicalTableExportSummary> {
        limits.validate()?;
        let (store, stream) = tokio::select! {
@@ -356,7 +361,7 @@ impl StatementExecutor {
                Ok((store, stream))
            } => result?,
        };
-        export_stream_managed(unit, stream, &store, limits, cancellation, budget).await
+        export_stream_managed(unit, stream, &store, limits, cancellation, budget, packed).await
    }
 }

@@ -375,6 +380,7 @@ async fn export_stream(
        limits,
        &cancellation.child_token(),
        ExportWriteBudget::new(1),
+        None,
    )
    .await
 }
@@ -386,8 +392,10 @@ async fn export_stream_managed(
    limits: LogicalTableExportLimits,
    cancellation: &CancellationToken,
    budget: Arc<ExportWriteBudget>,
+    packed: Option<PackedWriterRef>,
 ) -> Result<LogicalTableExportSummary> {
    let mut writers = TableWriters::new(budget.clone());
+    writers.packed = packed;
    let result = write_tables(
        unit,
        stream,
@@ -535,6 +543,26 @@ struct ActiveWriter {
 }

 impl ActiveWriter {
+    fn open_packed(
+        table: &LogicalTableProjection,
+        id: u32,
+        store: &ObjectStore,
+        limits: LogicalTableExportLimits,
+        shared: PackedWriterRef,
+    ) -> Result<Self> {
+        let path = format!("table-{id}.parquet");
+        let packed = PackedTableWriter::new(shared, table.name.clone(), id, false);
+        let writer = ParquetFileWriter::open_packed(
+            table.schema.clone(),
+            store.clone(),
+            &path,
+            Some(limits.writer),
+            packed,
+        )
+        .map_err(|e| map_writer_error(e, &path))?;
+        Ok(Self { path, writer })
+    }
+
    async fn open(
        table: &LogicalTableProjection,
        store: &ObjectStore,
@@ -146,6 +146,7 @@ async fn routes_across_batches_and_writes_empty_files() {
        export_limits(),
        &CancellationToken::new(),
        budget.clone(),
+        None,
    )
    .await
    .unwrap();
@@ -400,12 +401,17 @@ async fn native_histogram_parquet_roundtrip() {

 struct PausedFileWriter {
    inner: Option<object_store::layers::mock::oio::Writer>,
+    paused: bool,
    started: Arc<tokio::sync::Notify>,
    release: Arc<tokio::sync::Notify>,
 }

 impl object_store::layers::mock::oio::Write for PausedFileWriter {
    async fn write(&mut self, bytes: object_store::Buffer) -> object_store::Result<()> {
+        if self.paused {
+            return self.inner.as_mut().unwrap().write(bytes).await;
+        }
+        self.paused = true;
        let mut inner = self.inner.take().unwrap();
        let started = self.started.clone();
        let release = self.release.clone();
@@ -439,7 +445,7 @@ impl object_store::layers::mock::oio::Write for PausedFileWriter {

 #[tokio::test]
 async fn cancellation_drains_storage_and_preserves_committed_files() {
-    for large in [false, true] {
+    for (large, packed) in [(false, false), (true, false), (false, true)] {
        use object_store::layers::mock::{MockLayerBuilder, MockWriterFactory};
        let directory = common_test_util::temp_dir::create_temp_dir("metric_export_pending_open");
        let access =
@@ -459,6 +465,7 @@ async fn cancellation_drains_storage_and_preserves_committed_files() {
            move |_, _, inner| {
                Box::new(PausedFileWriter {
                    inner: Some(inner),
+                    paused: false,
                    started: started.clone(),
                    release: release.clone(),
                })
@@ -470,6 +477,16 @@ async fn cancellation_drains_storage_and_preserves_committed_files() {
                .build()
                .unwrap(),
        );
+        let store = store.layer(object_store::layers::CapabilityOverrideLayer::new(
+            move |mut capability| {
+                if packed {
+                    capability.write_multi_max_size = Some(256);
+                }
+                capability
+            },
+        ));
+        let packed = packed
+            .then(|| common_datasource::packed_writer::PackedWriter::new(store.clone()).unwrap());
        let cancellation = CancellationToken::new();
        let (unit, input) = if large {
            let field = Field::new("value", DataType::Utf8, true);
@@ -518,6 +535,7 @@ async fn cancellation_drains_storage_and_preserves_committed_files() {
            limits,
            &cancellation,
            budget.clone(),
+            packed.clone(),
        );
        tokio::pin!(export);
        tokio::select! {
@@ -528,6 +546,7 @@ async fn cancellation_drains_storage_and_preserves_committed_files() {
            assert!(budget.available().1 < 64 * 1024 * 1024);
        }
        let held = budget.available();
+        assert_eq!(held.0, 0);
        cancellation.cancel();
        assert!(futures::poll!(&mut export).is_pending());
        assert_eq!(budget.available().0, held.0);
@@ -542,8 +561,16 @@ async fn cancellation_drains_storage_and_preserves_committed_files() {
            result,
            Err(error::Error::LogicalTableExportCancelled { .. })
        ));
-        assert_eq!(store.exists("cpu.v1.parquet").await.unwrap(), !large);
-        if !large {
+        if let Some(packed) = &packed {
+            packed.lock().await.abort().await.unwrap();
+            assert!(!store.exists("pack-000000.bin").await.unwrap());
+            assert!(!store.exists("pack-index.json").await.unwrap());
+        }
+        assert_eq!(
+            store.exists("cpu.v1.parquet").await.unwrap(),
+            !large && packed.is_none()
+        );
+        if !large && packed.is_none() {
            let (_, batches) = read(&store, "cpu.v1.parquet").await;
            assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
        }
@@ -829,7 +856,8 @@ async fn groups_and_ordinary_files_share_writer_admission_and_drain() {
                        &store,
                        export_limits(),
                        &token,
-                        budget.clone()
+                        budget.clone(),
+                        None,
                    ),
                    export_stream_managed(
                        &b,
@@ -837,7 +865,8 @@ async fn groups_and_ordinary_files_share_writer_admission_and_drain() {
                        &store,
                        export_limits(),
                        &token,
-                        budget.clone()
+                        budget.clone(),
+                        None,
                    ),
                    ordinary,
                )
@@ -15,6 +15,7 @@
 use std::sync::Arc;

 use arrow::record_batch::RecordBatch;
+use common_datasource::packed_writer::PackedWriterRef;
 use futures::stream::FuturesUnordered;
 use futures::{FutureExt, StreamExt};
 use object_store::ObjectStore;
@@ -111,6 +112,7 @@ pub(crate) struct Payload {
 }

 pub(crate) struct TableWriters {
+    pub(crate) packed: Option<PackedWriterRef>,
    current: Option<(u32, mpsc::Sender<Payload>)>,
    tasks: FuturesUnordered<JoinHandle<Result<()>>>,
    budget: Arc<ExportWriteBudget>,
@@ -120,6 +122,7 @@ impl TableWriters {
    pub(crate) fn new(budget: Arc<ExportWriteBudget>) -> Self {
        Self {
            current: None,
+            packed: None,
            tasks: FuturesUnordered::new(),
            budget,
        }
@@ -150,7 +153,11 @@ impl TableWriters {
        self.close_input();
        let permit = self.budget.writer(token).await?;
        self.reap_for_admission(token).await?;
-        let writer = ActiveWriter::open(table, store, limits).await?;
+        let writer = if let Some(packed) = &self.packed {
+            ActiveWriter::open_packed(table, id, store, limits, packed.clone())?
+        } else {
+            ActiveWriter::open(table, store, limits).await?
+        };
        let (sender, receiver) = mpsc::channel(2);
        self.current = Some((id, sender));
        let token = token.clone();