fix segment query index not effect (#17704)

This commit is contained in:
Jyong 2025-04-09 19:09:08 +08:00 committed by GitHub
parent 6cf1ada36e
commit abfcd9d3b6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -5,7 +5,9 @@ from concurrent.futures import ThreadPoolExecutor
from typing import Optional from typing import Optional
from flask import Flask, current_app from flask import Flask, current_app
from sqlalchemy import and_, or_
from sqlalchemy.orm import load_only from sqlalchemy.orm import load_only
from sqlalchemy.sql.expression import false
from configs import dify_config from configs import dify_config
from core.rag.data_post_processor.data_post_processor import DataPostProcessor from core.rag.data_post_processor.data_post_processor import DataPostProcessor
@ -315,18 +317,30 @@ class RetrievalService:
child_chunks = db.session.query(ChildChunk).filter(ChildChunk.index_node_id.in_(child_index_node_ids)).all() child_chunks = db.session.query(ChildChunk).filter(ChildChunk.index_node_id.in_(child_index_node_ids)).all()
child_chunk_map = {chunk.index_node_id: chunk for chunk in child_chunks} child_chunk_map = {chunk.index_node_id: chunk for chunk in child_chunks}
# Batch query DocumentSegment with unified conditions segment_ids_from_child = [chunk.segment_id for chunk in child_chunks]
segment_conditions = []
if index_node_ids:
segment_conditions.append(DocumentSegment.index_node_id.in_(index_node_ids))
if segment_ids_from_child:
segment_conditions.append(DocumentSegment.id.in_(segment_ids_from_child))
if segment_conditions:
filter_expr = or_(*segment_conditions)
else:
filter_expr = false()
segment_map = { segment_map = {
segment.id: segment segment.id: segment
for segment in db.session.query(DocumentSegment) for segment in db.session.query(DocumentSegment)
.filter( .filter(
( and_(
DocumentSegment.index_node_id.in_(index_node_ids) filter_expr,
| DocumentSegment.id.in_([chunk.segment_id for chunk in child_chunks])
),
DocumentSegment.enabled == True, DocumentSegment.enabled == True,
DocumentSegment.status == "completed", DocumentSegment.status == "completed",
) )
)
.options( .options(
load_only( load_only(
DocumentSegment.id, DocumentSegment.id,