From 1ac0a144d51f1c1b36834995b192f4631088237e Mon Sep 17 00:00:00 2001 From: Accusys Date: Thu, 9 Jul 2026 23:44:06 +0800 Subject: [PATCH] perf(tkg): pre-build node ID map to eliminate N+1 SQL in edge builders - 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 --- src/core/processor/tkg.rs | 119 +++++++++++++++++++------------------- 1 file changed, 59 insertions(+), 60 deletions(-) diff --git a/src/core/processor/tkg.rs b/src/core/processor/tkg.rs index 2663b58..61dbc5a 100644 --- a/src/core/processor/tkg.rs +++ b/src/core/processor/tkg.rs @@ -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 { - 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 { 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 { - 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 { 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 { - 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 { 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 { 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 { 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 { +async fn build_has_appearance_edges(pool: &PgPool, file_uuid: &str, _node_id_map: &HashMap<(String, String), i64>) -> Result { let nodes_table = t("tkg_nodes"); let edges_table = t("tkg_edges");