From 1cd201d261cc8ae68291608aeb588515406ff58a Mon Sep 17 00:00:00 2001 From: vrustx <279631638@qq.com> Date: Fri, 17 Jul 2026 13:48:19 +0800 Subject: [PATCH] perf(postgres): classify result types once per column --- crates/dbx-core/src/db/postgres.rs | 300 ++++++++++++++++++++--------- crates/dbx-core/src/query.rs | 17 +- 2 files changed, 225 insertions(+), 92 deletions(-) diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index 68761449f5..39828c4278 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -304,123 +304,195 @@ fn format_pg_timestamptz(value: DateTime) -> String { value.to_rfc3339() } -pub(crate) fn pg_value_to_json(row: &Row, idx: usize, type_name: &str) -> serde_json::Value { +/// 时间类型解码失败后的回退目标，与原 if 链中时间分支之后的匹配顺序一一对应。 +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub(crate) enum PgTemporalFallback { + /// 数组类型名（下划线开头）落到通用数组解码。 + GenericArray, + /// `VECTOR(...)` 形式的类型名落到 pgvector 解码。 + Vector, + /// 其余落到通用试探链。 + Probe, +} + +/// 每列一次的类型分类结果，避免在逐单元格路径上重复 `to_uppercase` 与字符串比较链。 +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub(crate) enum PgColType { + Bytea, + Json, + Bool, + Temporal { fallback: PgTemporalFallback }, + Numeric, + Uuid, + Inet { cidr: bool }, + MacAddr, + BitString, + TsVector, + SystemU32, + InetArray { cidr: bool }, + MacAddrArray, + BitStringArray, + GenericArray, + Vector, + Geometry, + Other, +} + +pub(crate) fn classify_pg_type(type_name: &str) -> PgColType { let upper = type_name.to_uppercase(); if upper == "BYTEA" { - return row - .try_get::<_, Vec>(idx) - .map(|bytes| super::binary_value_to_json(&bytes)) - .unwrap_or(serde_json::Value::Null); + return PgColType::Bytea; } - if upper == "JSON" || upper == "JSONB" { - if let Ok(v) = row.try_get::<_, serde_json::Value>(idx) { - return serde_json::Value::String(v.to_string()); - } - if let Ok(v) = row.try_get::<_, String>(idx) { - return serde_json::Value::String(v); - } - return serde_json::Value::Null; + return PgColType::Json; } - if upper == "BOOL" { - return row.try_get::<_, bool>(idx).map(serde_json::Value::Bool).unwrap_or(serde_json::Value::Null); + return PgColType::Bool; } - if upper.contains("TIMESTAMP") || upper == "DATE" || upper == "TIME" || upper == "TIMETZ" || upper.contains("INTERVAL") { - if let Some(v) = pg_temporal_to_json_value(row, idx) { - return v; - } + let fallback = if upper.starts_with('_') { + PgTemporalFallback::GenericArray + } else if upper.starts_with("VECTOR(") { + PgTemporalFallback::Vector + } else { + PgTemporalFallback::Probe + }; + return PgColType::Temporal { fallback }; } - if upper == "NUMERIC" || upper == "DECIMAL" || upper == "MONEY" { - return row + return PgColType::Numeric; + } + if upper == "UUID" { + return PgColType::Uuid; + } + if matches!(upper.as_str(), "INET" | "CIDR") { + return PgColType::Inet { cidr: upper == "CIDR" }; + } + if matches!(upper.as_str(), "MACADDR" | "MACADDR8") { + return PgColType::MacAddr; + } + if matches!(upper.as_str(), "BIT" | "VARBIT") { + return PgColType::BitString; + } + if upper == "TSVECTOR" { + return PgColType::TsVector; + } + if matches!(upper.as_str(), "OID" | "XID" | "CID") { + return PgColType::SystemU32; + } + if matches!(upper.as_str(), "_INET" | "_CIDR") { + return PgColType::InetArray { cidr: upper == "_CIDR" }; + } + if matches!(upper.as_str(), "_MACADDR" | "_MACADDR8") { + return PgColType::MacAddrArray; + } + if matches!(upper.as_str(), "_BIT" | "_VARBIT") { + return PgColType::BitStringArray; + } + if upper.starts_with('_') { + return PgColType::GenericArray; + } + if upper == "VECTOR" || upper.starts_with("VECTOR(") { + return PgColType::Vector; + } + if upper == "GEOMETRY" || upper == "GEOGRAPHY" { + return PgColType::Geometry; + } + PgColType::Other +} + +pub(crate) fn classify_pg_column_types(column_types: &[String]) -> Vec { + column_types.iter().map(|type_name| classify_pg_type(type_name)).collect() +} + +pub(crate) fn pg_value_to_json_classified(row: &Row, idx: usize, col_type: PgColType) -> serde_json::Value { + match col_type { + PgColType::Bytea => row + .try_get::<_, Vec>(idx) + .map(|bytes| super::binary_value_to_json(&bytes)) + .unwrap_or(serde_json::Value::Null), + PgColType::Json => { + if let Ok(v) = row.try_get::<_, serde_json::Value>(idx) { + return serde_json::Value::String(v.to_string()); + } + if let Ok(v) = row.try_get::<_, String>(idx) { + return serde_json::Value::String(v); + } + serde_json::Value::Null + } + PgColType::Bool => row.try_get::<_, bool>(idx).map(serde_json::Value::Bool).unwrap_or(serde_json::Value::Null), + PgColType::Temporal { fallback } => { + if let Some(v) = pg_temporal_to_json_value(row, idx) { + return v; + } + match fallback { + PgTemporalFallback::GenericArray => pg_array_to_json_value(row, idx).unwrap_or(serde_json::Value::Null), + PgTemporalFallback::Vector => pg_vector_value_to_json(row, idx), + PgTemporalFallback::Probe => pg_fallback_value_to_json(row, idx), + } + } + PgColType::Numeric => row .try_get::<_, Decimal>(idx) .map(|v: Decimal| serde_json::Value::String(v.to_string())) - .unwrap_or(serde_json::Value::Null); - } - - if upper == "UUID" { - return row + .unwrap_or(serde_json::Value::Null), + PgColType::Uuid => row .try_get::<_, uuid::Uuid>(idx) .map(|v| serde_json::Value::String(v.to_string())) - .unwrap_or(serde_json::Value::Null); - } - - if matches!(upper.as_str(), "INET" | "CIDR") { - return pg_network_address_to_json_value(row, idx, upper == "CIDR").unwrap_or(serde_json::Value::Null); - } - - if matches!(upper.as_str(), "MACADDR" | "MACADDR8") { - return pg_macaddr_to_json_value(row, idx).unwrap_or(serde_json::Value::Null); - } - - if matches!(upper.as_str(), "BIT" | "VARBIT") { - return pg_bit_string_to_json_value(row, idx).unwrap_or(serde_json::Value::Null); - } - - if upper == "TSVECTOR" { - return row + .unwrap_or(serde_json::Value::Null), + PgColType::Inet { cidr } => pg_network_address_to_json_value(row, idx, cidr).unwrap_or(serde_json::Value::Null), + PgColType::MacAddr => pg_macaddr_to_json_value(row, idx).unwrap_or(serde_json::Value::Null), + PgColType::BitString => pg_bit_string_to_json_value(row, idx).unwrap_or(serde_json::Value::Null), + PgColType::TsVector => row .try_get::<_, PgRawBytes>(idx) .ok() .and_then(|raw| decode_tsvector_bytes(&raw.0)) .map(serde_json::Value::String) - .unwrap_or(serde_json::Value::Null); - } - - if matches!(upper.as_str(), "OID" | "XID" | "CID") { - return pg_system_u32_to_json(row, idx).unwrap_or(serde_json::Value::Null); - } - - if matches!(upper.as_str(), "_INET" | "_CIDR") { - return pg_network_address_array_to_json_value(row, idx, upper == "_CIDR").unwrap_or(serde_json::Value::Null); - } - - if matches!(upper.as_str(), "_MACADDR" | "_MACADDR8") { - return pg_macaddr_array_to_json_value(row, idx).unwrap_or(serde_json::Value::Null); - } - - if matches!(upper.as_str(), "_BIT" | "_VARBIT") { - return pg_bit_string_array_to_json_value(row, idx).unwrap_or(serde_json::Value::Null); - } - - if upper.starts_with('_') { - return pg_array_to_json_value(row, idx).unwrap_or(serde_json::Value::Null); - } - - if upper == "VECTOR" || upper.starts_with("VECTOR(") { - if let Ok(PgRawBytes(raw)) = row.try_get::<_, PgRawBytes>(idx) { - if let Some(floats) = decode_pgvector_bytes(&raw) { - return serde_json::Value::Array( - floats - .into_iter() - .map(|v| { - serde_json::Number::from_f64((v as f64 * 1_000_000.0).round() / 1_000_000.0) - .map(serde_json::Value::Number) - .unwrap_or(serde_json::Value::Null) - }) - .collect(), - ); + .unwrap_or(serde_json::Value::Null), + PgColType::SystemU32 => pg_system_u32_to_json(row, idx).unwrap_or(serde_json::Value::Null), + PgColType::InetArray { cidr } => { + pg_network_address_array_to_json_value(row, idx, cidr).unwrap_or(serde_json::Value::Null) + } + PgColType::MacAddrArray => pg_macaddr_array_to_json_value(row, idx).unwrap_or(serde_json::Value::Null), + PgColType::BitStringArray => pg_bit_string_array_to_json_value(row, idx).unwrap_or(serde_json::Value::Null), + PgColType::GenericArray => pg_array_to_json_value(row, idx).unwrap_or(serde_json::Value::Null), + PgColType::Vector => pg_vector_value_to_json(row, idx), + PgColType::Geometry => { + if let Ok(PgRawBytes(raw)) = row.try_get::<_, PgRawBytes>(idx) { + return super::wkb::wkb_to_wkt(&raw) + .map(serde_json::Value::String) + .unwrap_or_else(|| super::binary_value_to_json(&raw)); } + serde_json::Value::Null } - return serde_json::Value::Null; + PgColType::Other => pg_fallback_value_to_json(row, idx), } +} - if upper == "GEOMETRY" || upper == "GEOGRAPHY" { - if let Ok(PgRawBytes(raw)) = row.try_get::<_, PgRawBytes>(idx) { - return super::wkb::wkb_to_wkt(&raw) - .map(serde_json::Value::String) - .unwrap_or_else(|| super::binary_value_to_json(&raw)); +fn pg_vector_value_to_json(row: &Row, idx: usize) -> serde_json::Value { + if let Ok(PgRawBytes(raw)) = row.try_get::<_, PgRawBytes>(idx) { + if let Some(floats) = decode_pgvector_bytes(&raw) { + return serde_json::Value::Array( + floats + .into_iter() + .map(|v| { + serde_json::Number::from_f64((v as f64 * 1_000_000.0).round() / 1_000_000.0) + .map(serde_json::Value::Number) + .unwrap_or(serde_json::Value::Null) + }) + .collect(), + ); } - return serde_json::Value::Null; } + serde_json::Value::Null +} +fn pg_fallback_value_to_json(row: &Row, idx: usize) -> serde_json::Value { row.try_get::<_, String>(idx) .map(serde_json::Value::String) .or_else(|e| pg_system_u32_to_json(row, idx).ok_or(e)) @@ -630,6 +702,7 @@ async fn execute_select_prepared( ); let columns: Vec = stmt.columns().iter().map(|c| c.name().to_string()).collect(); let column_types: Vec = stmt.columns().iter().map(|c| c.type_().name().to_string()).collect(); + let column_classes = classify_pg_column_types(&column_types); let params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = Vec::new(); let query_start = Instant::now(); @@ -653,7 +726,9 @@ async fn execute_select_prepared( let row = row_result?; result_rows.push( (0..row.columns().len()) - .map(|i| pg_value_to_json(&row, i, column_types.get(i).map(String::as_str).unwrap_or(""))) + .map(|i| { + pg_value_to_json_classified(&row, i, column_classes.get(i).copied().unwrap_or(PgColType::Other)) + }) .collect(), ); } @@ -785,6 +860,7 @@ async fn stream_select_query_prepared( client.prepare_cached(sql).await.map_err(|err| PostgresQueryStreamError::Postgres { err, emitted: false })?; let columns: Vec = stmt.columns().iter().map(|c| c.name().to_string()).collect(); let column_types: Vec = stmt.columns().iter().map(|c| c.type_().name().to_string()).collect(); + let column_classes = classify_pg_column_types(&column_types); let params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = Vec::new(); let stream = client @@ -806,7 +882,7 @@ async fn stream_select_query_prepared( columns_emitted = true; } let values = (0..row.columns().len()) - .map(|i| pg_value_to_json(&row, i, column_types.get(i).map(String::as_str).unwrap_or(""))) + .map(|i| pg_value_to_json_classified(&row, i, column_classes.get(i).copied().unwrap_or(PgColType::Other))) .collect(); on_item(PostgresQueryStreamItem::Row(values)).map_err(PostgresQueryStreamError::Export)?; rows_streamed += 1; @@ -915,6 +991,7 @@ async fn stream_query_rows_on_client( ) -> Result { let stmt = client.prepare_cached(sql).await.map_err(pg_error_to_string)?; let column_types: Vec = stmt.columns().iter().map(|c| c.type_().name().to_string()).collect(); + let column_classes = classify_pg_column_types(&column_types); let params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = Vec::new(); let stream = client.query_raw(&stmt, params).await.map_err(pg_error_to_string)?; tokio::pin!(stream); @@ -930,7 +1007,7 @@ async fn stream_query_rows_on_client( } let row = row_result.map_err(pg_error_to_string)?; let values: Vec = (0..row.columns().len()) - .map(|i| pg_value_to_json(&row, i, column_types.get(i).map(String::as_str).unwrap_or(""))) + .map(|i| pg_value_to_json_classified(&row, i, column_classes.get(i).copied().unwrap_or(PgColType::Other))) .collect(); on_row(&values)?; rows_exported += 1; @@ -3284,6 +3361,51 @@ mod tests { use std::time::Instant; use tokio_postgres::types::FromSql; + #[test] + fn classify_pg_type_covers_all_dispatch_branches() { + assert_eq!(classify_pg_type("bytea"), PgColType::Bytea); + assert_eq!(classify_pg_type("json"), PgColType::Json); + assert_eq!(classify_pg_type("JSONB"), PgColType::Json); + assert_eq!(classify_pg_type("bool"), PgColType::Bool); + assert_eq!(classify_pg_type("timestamp"), PgColType::Temporal { fallback: PgTemporalFallback::Probe }); + assert_eq!(classify_pg_type("timestamptz"), PgColType::Temporal { fallback: PgTemporalFallback::Probe }); + assert_eq!(classify_pg_type("date"), PgColType::Temporal { fallback: PgTemporalFallback::Probe }); + assert_eq!(classify_pg_type("time"), PgColType::Temporal { fallback: PgTemporalFallback::Probe }); + assert_eq!(classify_pg_type("timetz"), PgColType::Temporal { fallback: PgTemporalFallback::Probe }); + assert_eq!(classify_pg_type("interval"), PgColType::Temporal { fallback: PgTemporalFallback::Probe }); + // 时间数组类型名在原实现中先进时间分支、解码失败后落到通用数组分支 + assert_eq!(classify_pg_type("_timestamp"), PgColType::Temporal { fallback: PgTemporalFallback::GenericArray }); + assert_eq!(classify_pg_type("_interval"), PgColType::Temporal { fallback: PgTemporalFallback::GenericArray }); + // 同时命中时间关键字与 VECTOR( 前缀的类型名，原实现时间解码失败后走 vector 分支 + assert_eq!(classify_pg_type("vector(timestamp)"), PgColType::Temporal { fallback: PgTemporalFallback::Vector }); + assert_eq!(classify_pg_type("numeric"), PgColType::Numeric); + assert_eq!(classify_pg_type("money"), PgColType::Numeric); + assert_eq!(classify_pg_type("uuid"), PgColType::Uuid); + assert_eq!(classify_pg_type("inet"), PgColType::Inet { cidr: false }); + assert_eq!(classify_pg_type("cidr"), PgColType::Inet { cidr: true }); + assert_eq!(classify_pg_type("macaddr"), PgColType::MacAddr); + assert_eq!(classify_pg_type("macaddr8"), PgColType::MacAddr); + assert_eq!(classify_pg_type("bit"), PgColType::BitString); + assert_eq!(classify_pg_type("varbit"), PgColType::BitString); + assert_eq!(classify_pg_type("tsvector"), PgColType::TsVector); + assert_eq!(classify_pg_type("oid"), PgColType::SystemU32); + assert_eq!(classify_pg_type("xid"), PgColType::SystemU32); + assert_eq!(classify_pg_type("_inet"), PgColType::InetArray { cidr: false }); + assert_eq!(classify_pg_type("_cidr"), PgColType::InetArray { cidr: true }); + assert_eq!(classify_pg_type("_macaddr"), PgColType::MacAddrArray); + assert_eq!(classify_pg_type("_bit"), PgColType::BitStringArray); + assert_eq!(classify_pg_type("_varbit"), PgColType::BitStringArray); + assert_eq!(classify_pg_type("_int4"), PgColType::GenericArray); + assert_eq!(classify_pg_type("_time"), PgColType::GenericArray); + assert_eq!(classify_pg_type("vector"), PgColType::Vector); + assert_eq!(classify_pg_type("vector(3)"), PgColType::Vector); + assert_eq!(classify_pg_type("geometry"), PgColType::Geometry); + assert_eq!(classify_pg_type("geography"), PgColType::Geometry); + assert_eq!(classify_pg_type("int4"), PgColType::Other); + assert_eq!(classify_pg_type("varchar"), PgColType::Other); + assert_eq!(classify_pg_type(""), PgColType::Other); + } + struct DockerPostgres { name: String, port: u16, diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index e787b2cfd2..e05e14f354 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -3163,6 +3163,7 @@ where Ok(stmt) => { let column_types: Vec = stmt.columns().iter().map(|column| column.type_().name().to_string()).collect(); + let column_classes = db::postgres::classify_pg_column_types(&column_types); let params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = Vec::new(); match conn.query_raw(&stmt, params).await { Ok(stream) => { @@ -3175,10 +3176,13 @@ where Ok(row) => { let values = (0..row.columns().len()) .map(|index| { - db::postgres::pg_value_to_json( + db::postgres::pg_value_to_json_classified( &row, index, - column_types.get(index).map(String::as_str).unwrap_or(""), + column_classes + .get(index) + .copied() + .unwrap_or(db::postgres::PgColType::Other), ) }) .collect(); @@ -3335,6 +3339,7 @@ async fn execute_manual_txn_postgres_statement( let stmt = conn.prepare_cached(sql).await.map_err(|e| format!("Prepare failed: {e}"))?; let columns: Vec = stmt.columns().iter().map(|c| c.name().to_string()).collect(); let column_types: Vec = stmt.columns().iter().map(|c| c.type_().name().to_string()).collect(); + let column_classes = db::postgres::classify_pg_column_types(&column_types); let params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = Vec::new(); let stream = conn.query_raw(&stmt, params).await.map_err(|e| format!("Query failed: {e}"))?; tokio::pin!(stream); @@ -3347,7 +3352,13 @@ async fn execute_manual_txn_postgres_statement( } let row = row_result.map_err(|e| format!("Query failed: {e}"))?; let values: Vec = (0..row.columns().len()) - .map(|i| db::postgres::pg_value_to_json(&row, i, column_types.get(i).map(String::as_str).unwrap_or(""))) + .map(|i| { + db::postgres::pg_value_to_json_classified( + &row, + i, + column_classes.get(i).copied().unwrap_or(db::postgres::PgColType::Other), + ) + }) .collect(); data.push(values); }