xiaoxue_admin
/
dbx
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf(postgres): classify result types once per column

This commit is contained in:
vrustx 2026-07-17 13:48:19 +08:00 committed by GitHub
parent 7caf9b1feb
commit 1cd201d261
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 225 additions and 92 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

300
crates/dbx-core/src/db/postgres.rs
View File

@ -304,123 +304,195 @@ fn format_pg_timestamptz(value: DateTime<Local>) -> 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<u8>>(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<PgColType> {
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<u8>>(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<String> = stmt.columns().iter().map(|c| c.name().to_string()).collect();
let column_types: Vec<String> = 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<String> = stmt.columns().iter().map(|c| c.name().to_string()).collect();
let column_types: Vec<String> = 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<u64, String> {
let stmt = client.prepare_cached(sql).await.map_err(pg_error_to_string)?;
let column_types: Vec<String> = 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<serde_json::Value> = (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,

17
crates/dbx-core/src/query.rs
View File

@ -3163,6 +3163,7 @@ where
Ok(stmt) => {
let column_types: Vec<String> =
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<String> = stmt.columns().iter().map(|c| c.name().to_string()).collect();
let column_types: Vec<String> = 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<serde_json::Value> = (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);
}