From 2891167015b4b3595246ec47ae438d60a9a0922a Mon Sep 17 00:00:00 2001 From: vrustx <279631638@qq.com> Date: Sun, 19 Jul 2026 01:32:56 +0800 Subject: [PATCH] fix(export): remove CSV and TSV buffer allocations --- crates/dbx-core/src/csv_export.rs | 227 ++++++++++++++++++++++++---- crates/dbx-core/src/table_export.rs | 21 ++- 2 files changed, 211 insertions(+), 37 deletions(-) diff --git a/crates/dbx-core/src/csv_export.rs b/crates/dbx-core/src/csv_export.rs index 6677c64b61..7568142065 100644 --- a/crates/dbx-core/src/csv_export.rs +++ b/crates/dbx-core/src/csv_export.rs @@ -1,5 +1,6 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; +use std::fmt; use std::fs::File; use std::io::{BufWriter, Write}; @@ -27,63 +28,177 @@ pub struct TableCsvExportOptions { pub timeout_secs: Option, } -pub(crate) fn escape_csv(value: &str) -> String { - format!("\"{}\"", value.replace('"', "\"\"")) +/// CSV 转义直写目标 buffer：包引号 + 内部 `"` 翻倍。值不含 `"` 时整段拷贝， +/// 不做 replace 分配（逐批流式导出对每个单元格调用，是导出热路径）。 +fn push_csv_escaped_content(out: &mut String, value: &str) { + let mut rest = value; + while let Some(pos) = rest.find('"') { + out.push_str(&rest[..=pos]); + out.push('"'); + rest = &rest[pos + 1..]; + } + out.push_str(rest); } -pub(crate) fn value_to_csv_text(value: &Value) -> String { - match value { - Value::Null => String::new(), - Value::Bool(v) => v.to_string(), - Value::Number(v) => v.to_string(), - Value::String(v) => v.clone(), - other => other.to_string(), +pub(crate) fn push_csv_escaped(out: &mut String, value: &str) { + out.push('"'); + push_csv_escaped_content(out, value); + out.push('"'); +} + +struct CsvEscapedWriter<'a>(&'a mut String); + +impl fmt::Write for CsvEscapedWriter<'_> { + fn write_str(&mut self, value: &str) -> fmt::Result { + push_csv_escaped_content(self.0, value); + Ok(()) } } +/// 将表导出 CSV 值直接写入已有 buffer；包括 NULL 在内的值均保留分页导出的带引号旧语义。 +pub(crate) fn push_csv_text_value(out: &mut String, value: &Value) { + out.push('"'); + match value { + Value::Null => {} + Value::String(value) => push_csv_escaped_content(out, value), + Value::Bool(value) => out.push_str(if *value { "true" } else { "false" }), + Value::Number(value) => { + fmt::write(out, format_args!("{value}")).expect("writing a number into a String cannot fail") + } + // 数组和对象可能包含引号，通过转义 writer 格式化，避免分配中间 JSON 字符串 + other => fmt::write(&mut CsvEscapedWriter(out), format_args!("{other}")) + .expect("writing JSON into a String cannot fail"), + } + out.push('"'); +} + +fn push_csv_value(out: &mut String, value: &Value) { + if value.is_null() { + return; + } + push_csv_text_value(out, value); +} + +/// TSV 转义直写：仅含特殊字符时包引号（语义与原 escape_tsv 一致）。 +fn push_tsv_escaped(out: &mut String, value: &str) { + if value.contains('\t') || value.contains('\n') || value.contains('\r') || value.contains('"') { + push_csv_escaped(out, value); + } else { + out.push_str(value); + } +} + +fn push_tsv_value(out: &mut String, value: &Value) { + match value { + Value::Null => {} + Value::String(v) => push_tsv_escaped(out, v), + Value::Bool(v) => out.push_str(if *v { "true" } else { "false" }), + Value::Number(value) => { + fmt::write(out, format_args!("{value}")).expect("writing a number into a String cannot fail") + } + other => push_tsv_escaped(out, &other.to_string()), + } +} + +/// 预分配粗估：按全部行的实际单元格数求和（不假设等宽），饱和运算防溢出， +/// 并设上限——估算只是性能提示，绝不能因病态输入放大成巨额分配 +const ROWS_CAPACITY_ESTIMATE_MAX: usize = 16 * 1024 * 1024; + +pub(crate) fn estimated_rows_capacity(rows: &[Vec]) -> usize { + let cells: usize = rows.iter().map(Vec::len).fold(0usize, usize::saturating_add); + cells.saturating_mul(12).min(ROWS_CAPACITY_ESTIMATE_MAX) +} + +pub(crate) fn escape_csv(value: &str) -> String { + let mut out = String::with_capacity(value.len() + 2); + push_csv_escaped(&mut out, value); + out +} + fn format_csv_value(value: &Value) -> String { - match value { - Value::Null => String::new(), - other => escape_csv(&value_to_csv_text(other)), - } + let mut out = String::new(); + push_csv_value(&mut out, value); + out } +#[cfg_attr(not(test), allow(dead_code))] pub(crate) fn escape_tsv(value: &str) -> String { - if value.contains('\t') || value.contains('\n') || value.contains('\r') || value.contains('"') { - format!("\"{}\"", value.replace('"', "\"\"")) - } else { - value.to_string() + let mut out = String::with_capacity(value.len()); + push_tsv_escaped(&mut out, value); + out +} + +fn push_tsv_rows(out: &mut String, rows: &[Vec]) { + for (row_index, row) in rows.iter().enumerate() { + if row_index > 0 { + out.push('\n'); + } + for (cell_index, cell) in row.iter().enumerate() { + if cell_index > 0 { + out.push('\t'); + } + push_tsv_value(out, cell); + } } } pub(crate) fn format_tsv_rows(rows: &[Vec]) -> String { - rows.iter() - .map(|row| row.iter().map(|cell| escape_tsv(&value_to_csv_text(cell))).collect::>().join("\t")) - .collect::>() - .join("\n") + let mut out = String::with_capacity(estimated_rows_capacity(rows)); + push_tsv_rows(&mut out, rows); + out } pub(crate) fn format_tsv(columns: &[String], rows: &[Vec]) -> String { - let header = columns.iter().map(|col| escape_tsv(col)).collect::>().join("\t"); - let body = format_tsv_rows(rows); - format!("{header}\n{body}") + let mut out = String::with_capacity( + estimated_rows_capacity(rows).saturating_add(columns.len().saturating_mul(12)).min(ROWS_CAPACITY_ESTIMATE_MAX), + ); + for (index, column) in columns.iter().enumerate() { + if index > 0 { + out.push('\t'); + } + push_tsv_escaped(&mut out, column); + } + out.push('\n'); + push_tsv_rows(&mut out, rows); + out } /// Format query-result rows as CSV text without a header row. Database NULLs /// use the same empty-cell representation as table-data exports. Used by the /// streaming query-result export for batches after the first. +fn push_query_result_csv_rows(out: &mut String, rows: &[Vec]) { + for (row_index, row) in rows.iter().enumerate() { + if row_index > 0 { + out.push('\n'); + } + for (cell_index, cell) in row.iter().enumerate() { + if cell_index > 0 { + out.push(','); + } + push_csv_value(out, cell); + } + } +} + pub fn format_query_result_csv_rows(rows: &[Vec]) -> String { - rows.iter().map(|row| row.iter().map(format_csv_value).collect::>().join(",")).collect::>().join("\n") + let mut out = String::with_capacity(estimated_rows_capacity(rows)); + push_query_result_csv_rows(&mut out, rows); + out } fn format_csv_with_value_formatter(columns: &[String], rows: &[Vec]) -> String { - let header = columns.iter().map(|col| escape_csv(col)).collect::>().join(","); - let body = rows - .iter() - .map(|row| row.iter().map(format_csv_value).collect::>().join(",")) - .collect::>() - .join("\n"); - format!("{header}\n{body}") + let mut out = String::with_capacity( + estimated_rows_capacity(rows).saturating_add(columns.len().saturating_mul(12)).min(ROWS_CAPACITY_ESTIMATE_MAX), + ); + for (index, column) in columns.iter().enumerate() { + if index > 0 { + out.push(','); + } + push_csv_escaped(&mut out, column); + } + out.push('\n'); + push_query_result_csv_rows(&mut out, rows); + out } pub fn format_csv(columns: &[String], rows: &[Vec]) -> String { @@ -236,5 +351,53 @@ mod tests { assert_eq!(out, "id\tnote\n1\t\n2\t\"line1\n\"\"line2\"\"\""); } + #[test] + fn capacity_estimate_sums_actual_cells_across_ragged_rows() { + // 不等宽行按实际单元格数求和，不得按首行宽度放大 + let wide_first = vec![vec![serde_json::Value::Null; 1000], vec![], vec![serde_json::Value::Null]]; + assert_eq!(super::estimated_rows_capacity(&wide_first), 1001 * 12); + assert!(super::estimated_rows_capacity(&[]) == 0); + } + + #[test] + fn escape_tsv_matches_reference_semantics() { + // TSV 仅在含 \t/\n/\r/引号时包引号；逗号不触发 + for input in ["", "plain", "with,comma", "tab\there", "line\nbreak", "cr\rhere", "quo\"te", "\t\"mix\""] { + let expected = + if input.contains('\t') || input.contains('\n') || input.contains('\r') || input.contains('"') { + format!("\"{}\"", input.replace('"', "\"\"")) + } else { + input.to_string() + }; + assert_eq!(super::escape_tsv(input), expected, "input: {input:?}"); + } + } + + #[test] + fn push_csv_escaped_matches_replace_reference() { + // 直写实现必须与原 replace 版本逐字节等价（含引号在首/尾/连续的边界） + for input in ["", "plain", "\"", "\"\"", "a\"b", "\"start", "end\"", "mid\"\"dle", "逗,号\n换行"] { + let expected = format!("\"{}\"", input.replace('"', "\"\"")); + assert_eq!(super::escape_csv(input), expected, "input: {input:?}"); + } + } + + #[test] + fn push_csv_text_value_preserves_paginated_export_semantics() { + let cases = [ + (serde_json::Value::Null, "\"\""), + (serde_json::json!(true), "\"true\""), + (serde_json::json!(42.5), "\"42.5\""), + (serde_json::json!("a\"b"), "\"a\"\"b\""), + (serde_json::json!({"key": "value"}), "\"{\"\"key\"\":\"\"value\"\"}\""), + ]; + + for (value, expected) in cases { + let mut out = String::from("prefix,"); + super::push_csv_text_value(&mut out, &value); + assert_eq!(out, format!("prefix,{expected}")); + } + } + use serde_json::Value; } diff --git a/crates/dbx-core/src/table_export.rs b/crates/dbx-core/src/table_export.rs index 5bb3e94fff..a426c1dac9 100644 --- a/crates/dbx-core/src/table_export.rs +++ b/crates/dbx-core/src/table_export.rs @@ -8,7 +8,7 @@ use tokio_util::sync::CancellationToken; use crate::connection::MysqlMode; use crate::connection::{config_for_pool_key, task_client_session_id, AppState, PoolKind}; -use crate::csv_export::{escape_csv, format_csv, format_tsv, format_tsv_rows, value_to_csv_text}; +use crate::csv_export::{format_csv, format_tsv, format_tsv_rows, push_csv_text_value}; pub use crate::database_export::ExportStatus; use crate::database_export::{ build_export_insert_statements, is_export_cancelled, is_internal_export_column, BuildExportInsertStatementsOptions, @@ -75,10 +75,21 @@ pub struct TableExportProgress { /// Format rows as CSV text without a header row. /// Used for streaming subsequent pagination batches. fn format_csv_rows(rows: &[Vec]) -> String { - rows.iter() - .map(|row| row.iter().map(|cell| escape_csv(&value_to_csv_text(cell))).collect::>().join(",")) - .collect::>() - .join("\n") + // 注意：该无表头批次路径的 Null 输出为 ""（带引号空串），与查询结果导出 + // 及首批 format_csv 的裸空单元格语义不同，直写化必须保留该差异 + let mut out = String::with_capacity(crate::csv_export::estimated_rows_capacity(rows)); + for (row_index, row) in rows.iter().enumerate() { + if row_index > 0 { + out.push('\n'); + } + for (cell_index, cell) in row.iter().enumerate() { + if cell_index > 0 { + out.push(','); + } + push_csv_text_value(&mut out, cell); + } + } + out } fn export_column_types(request: &TableExportRequest) -> Vec {