warren/momentry_core
1
0
Fork 0
Code Issues 2 Pull Requests Actions Packages Projects Releases 2 Wiki Activity

perf(tkg): pre-build node ID map to eliminate N+1 SQL in edge builders

Browse Source 
- Add build_node_id_map() to fetch all node IDs once into HashMap
- Pass node_id_map to all 6 edge builders
- build_co_occurrence_edges: uses map lookup instead of per-face SELECT
- build_speaker_face_edges: uses map lookup instead of per-trace SELECT
- Remove redundant Qdrant scroll in build_co_occurrence_edges
- Expected: 90%+ reduction in SQL queries during edge building
This commit is contained in:
Accusys
2026-07-09 23:44:06 +08:00
parent 356edb16be
commit 1ac0a144d5
1 changed files with 59 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/core/processor/tkg.rs
+59 -60
View File
@@ -104,6 +104,23 @@ fn find_trace_in_frame(
best.map(|(tid, _)| tid)
}
/// Pre-fetch all node IDs for a file into a HashMap for O(1) lookup during edge building
async fn build_node_id_map(pool: &PgPool, file_uuid: &str) -> HashMap<(String, String), i64> {
let nodes_table = t("tkg_nodes");
let rows: Vec<(String, String, i64)> = sqlx::query_as(&format!(
"SELECT node_type, external_id, id FROM {} WHERE file_uuid = $1",
nodes_table
))
.bind(file_uuid)
.fetch_all(pool)
.await
.unwrap_or_default();
rows.into_iter()
.map(|(node_type, external_id, id)| ((node_type, external_id), id))
.collect()
}
fn t(name: &str) -> String {
let schema = std::env::var("DATABASE_SCHEMA").unwrap_or_else(|_| "dev".to_string());
if schema == "public" {
@@ -854,6 +871,14 @@ pub async fn build_tkg(
publish_tkg_progress(r, file_uuid, &progress).await;
}
// Pre-fetch all node IDs for O(1) lookup during edge building
let node_id_map = build_node_id_map(pool, file_uuid).await;
tracing::info!(
"[TKG] Built node ID map with {} entries for {}",
node_id_map.len(),
file_uuid
);
// ── Edge builders ─────────────────────────────────────────────────
progress.update_phase(
TkgPhase::CoOccurrenceEdges,
@@ -863,7 +888,7 @@ pub async fn build_tkg(
if let Some(r) = redis.as_ref() {
publish_tkg_progress(r, file_uuid, &progress).await;
}
let e_co = build_co_occurrence_edges(pool, file_uuid, output_dir).await?;
let e_co = build_co_occurrence_edges(pool, file_uuid, output_dir, &face_points, &node_id_map).await?;
progress.stats.co_occurrence_edges = e_co as i64;
progress.update_phase(
TkgPhase::CoOccurrenceEdges,
@@ -882,7 +907,7 @@ pub async fn build_tkg(
if let Some(r) = redis.as_ref() {
publish_tkg_progress(r, file_uuid, &progress).await;
}
let e_sf = build_speaker_face_edges(pool, file_uuid, output_dir).await?;
let e_sf = build_speaker_face_edges(pool, file_uuid, output_dir, &node_id_map).await?;
progress.stats.speaker_face_edges = e_sf as i64;
progress.update_phase(
TkgPhase::SpeakerFaceEdges,
@@ -897,7 +922,7 @@ pub async fn build_tkg(
if let Some(r) = redis.as_ref() {
publish_tkg_progress(r, file_uuid, &progress).await;
}
let e_ff = build_face_face_edges(pool, file_uuid, &pose_data, &face_points).await?;
let e_ff = build_face_face_edges(pool, file_uuid, &pose_data, &face_points, &node_id_map).await?;
progress.stats.face_face_edges = e_ff as i64;
progress.update_phase(TkgPhase::FaceFaceEdges, 1.0, &format!("{} face-face", e_ff));
if let Some(r) = redis.as_ref() {
@@ -912,7 +937,7 @@ pub async fn build_tkg(
if let Some(r) = redis.as_ref() {
publish_tkg_progress(r, file_uuid, &progress).await;
}
let e_mg = build_mutual_gaze_edges(pool, file_uuid, &pose_data, &face_points).await?;
let e_mg = build_mutual_gaze_edges(pool, file_uuid, &pose_data, &face_points, &node_id_map).await?;
progress.stats.mutual_gaze_edges = e_mg as i64;
progress.update_phase(
TkgPhase::MutualGazeEdges,
@@ -927,7 +952,7 @@ pub async fn build_tkg(
if let Some(r) = redis.as_ref() {
publish_tkg_progress(r, file_uuid, &progress).await;
}
let e_ls = build_lip_sync_edges(pool, file_uuid, output_dir, &pose_data).await?;
let e_ls = build_lip_sync_edges(pool, file_uuid, output_dir, &pose_data, &node_id_map).await?;
progress.stats.lip_sync_edges = e_ls as i64;
progress.update_phase(TkgPhase::LipSyncEdges, 1.0, &format!("{} lip sync", e_ls));
if let Some(r) = redis.as_ref() {
@@ -942,7 +967,7 @@ pub async fn build_tkg(
if let Some(r) = redis.as_ref() {
publish_tkg_progress(r, file_uuid, &progress).await;
}
let e_ha = build_has_appearance_edges(pool, file_uuid).await?;
let e_ha = build_has_appearance_edges(pool, file_uuid, &node_id_map).await?;
progress.stats.has_appearance_edges = e_ha as i64;
progress.update_phase(
TkgPhase::HasAppearanceEdges,
@@ -1491,14 +1516,18 @@ async fn build_co_occurrence_edges(
pool: &PgPool,
file_uuid: &str,
output_dir: &str,
face_points: &[FacePoint],
node_id_map: &HashMap<(String, String), i64>,
) -> Result<usize> {
build_co_occurrence_edges_from_qdrant(pool, file_uuid, output_dir).await
build_co_occurrence_edges_from_qdrant(pool, file_uuid, output_dir, face_points, node_id_map).await
}
async fn build_co_occurrence_edges_from_qdrant(
pool: &PgPool,
file_uuid: &str,
output_dir: &str,
face_points: &[FacePoint],
node_id_map: &HashMap<(String, String), i64>,
) -> Result<usize> {
let yolo_path = Path::new(output_dir).join(format!("{}.yolo.json", file_uuid));
if !yolo_path.exists() {
@@ -1508,14 +1537,13 @@ async fn build_co_occurrence_edges_from_qdrant(
let content = std::fs::read_to_string(&yolo_path)?;
let yolo: YoloJson = serde_json::from_str(&content)?;
let nodes_table = t("tkg_nodes");
let edges_table = t("tkg_edges");
// Get face points from Qdrant
let points = scroll_face_points(file_uuid).await;
// Use shared face points (passed from build_tkg)
let points = face_points;
let mut edge_count = 0;
for face in &points {
for face in points {
let yolo_frame = match yolo.frames.get(&face.frame.to_string()) {
Some(f) => f,
None => continue,
@@ -1531,33 +1559,15 @@ async fn build_co_occurrence_edges_from_qdrant(
continue;
}
let external_id = format!("face_track_{}", face.trace_id);
let face_node: Option<(i64,)> = sqlx::query_as(&format!(
"SELECT id FROM {} WHERE file_uuid=$1 AND node_type='face_track' AND external_id=$2",
nodes_table
))
.bind(file_uuid)
.bind(&external_id)
.fetch_optional(pool)
.await?;
let face_node_id = match face_node {
Some((id,)) => id,
// Use pre-built node ID map (no SQL query)
let face_node_id = match node_id_map.get(&("face_track".to_string(), format!("face_track_{}", face.trace_id))) {
Some(id) => *id,
None => continue,
};
for det in dets {
let obj_node: Option<(i64,)> = sqlx::query_as(&format!(
"SELECT id FROM {} WHERE file_uuid=$1 AND node_type='object' AND external_id=$2",
nodes_table
))
.bind(file_uuid)
.bind(&det.class_name)
.fetch_optional(pool)
.await?;
let obj_node_id = match obj_node {
Some((id,)) => id,
let obj_node_id = match node_id_map.get(&("object".to_string(), det.class_name.clone())) {
Some(id) => *id,
None => continue,
};
@@ -1603,14 +1613,16 @@ async fn build_speaker_face_edges(
pool: &PgPool,
file_uuid: &str,
output_dir: &str,
node_id_map: &HashMap<(String, String), i64>,
) -> Result<usize> {
build_speaker_face_edges_from_pg(pool, file_uuid, output_dir).await
build_speaker_face_edges_from_pg(pool, file_uuid, output_dir, node_id_map).await
}
async fn build_speaker_face_edges_from_pg(
pool: &PgPool,
file_uuid: &str,
output_dir: &str,
node_id_map: &HashMap<(String, String), i64>,
) -> Result<usize> {
let asrx_path = Path::new(output_dir).join(format!("{}.asrx.json", file_uuid));
if !asrx_path.exists() {
@@ -1656,18 +1668,9 @@ async fn build_speaker_face_edges_from_pg(
let mut edge_count = 0;
for (tid, sf, ef) in &traces {
let face_ext_id = format!("face_track_{}", tid);
let face_node: Option<(i64,)> = sqlx::query_as(&format!(
"SELECT id FROM {} WHERE file_uuid=$1 AND node_type='face_track' AND external_id=$2",
nodes_table
))
.bind(file_uuid)
.bind(&face_ext_id)
.fetch_optional(pool)
.await?;
let face_node_id = match face_node {
Some((id,)) => id,
// Use pre-built node ID map (no SQL query)
let face_node_id = match node_id_map.get(&("face_track".to_string(), format!("face_track_{}", tid))) {
Some(id) => *id,
None => continue,
};
@@ -1695,17 +1698,9 @@ async fn build_speaker_face_edges_from_pg(
continue;
}
let speaker_node: Option<(i64,)> = sqlx::query_as(&format!(
"SELECT id FROM {} WHERE file_uuid=$1 AND node_type='speaker' AND external_id=$2",
nodes_table
))
.bind(file_uuid)
.bind(&seg.speaker_id)
.fetch_optional(pool)
.await?;
let speaker_node_id = match speaker_node {
Some((id,)) => id,
// Use pre-built node ID map (no SQL query)
let speaker_node_id = match node_id_map.get(&("speaker".to_string(), seg.speaker_id.clone())) {
Some(id) => *id,
None => continue,
};
@@ -1745,8 +1740,9 @@ async fn build_face_face_edges(
file_uuid: &str,
pose_data: &[FacePose],
face_points: &[FacePoint],
node_id_map: &HashMap<(String, String), i64>,
) -> Result<usize> {
build_face_face_edges_from_qdrant(pool, file_uuid, pose_data, face_points).await
build_face_face_edges_from_qdrant(pool, file_uuid, pose_data, face_points, node_id_map).await
}
async fn build_face_face_edges_from_qdrant(
@@ -1754,6 +1750,7 @@ async fn build_face_face_edges_from_qdrant(
file_uuid: &str,
pose_data: &[FacePose],
face_points: &[FacePoint],
node_id_map: &HashMap<(String, String), i64>,
) -> Result<usize> {
let nodes_table = t("tkg_nodes");
let edges_table = t("tkg_edges");
@@ -2084,6 +2081,7 @@ async fn build_mutual_gaze_edges(
file_uuid: &str,
pose_data: &[FacePose],
face_points: &[FacePoint],
_node_id_map: &HashMap<(String, String), i64>,
) -> Result<usize> {
let nodes_table = t("tkg_nodes");
let edges_table = t("tkg_edges");
@@ -2510,6 +2508,7 @@ async fn build_lip_sync_edges(
file_uuid: &str,
output_dir: &str,
pose_data: &[FacePose],
_node_id_map: &HashMap<(String, String), i64>,
) -> Result<usize> {
let nodes_table = t("tkg_nodes");
let edges_table = t("tkg_edges");
@@ -2824,7 +2823,7 @@ async fn build_accessory_nodes(pool: &PgPool, file_uuid: &str, output_dir: &str)
// ── HAS_APPEARANCE Edges ──────────────────────────────────────────
async fn build_has_appearance_edges(pool: &PgPool, file_uuid: &str) -> Result<usize> {
async fn build_has_appearance_edges(pool: &PgPool, file_uuid: &str, _node_id_map: &HashMap<(String, String), i64>) -> Result<usize> {
let nodes_table = t("tkg_nodes");
let edges_table = t("tkg_edges");