AI/dify
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/langgenius/dify.git synced 2025-04-21 21:29:48 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix segment query index not effect (#17704)

Browse Source
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
1 changed files with 21 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
api/core/rag/datasource/retrieval_service.py
View File

@ -5,7 +5,9 @@ from concurrent.futures import ThreadPoolExecutor
from typing import Optional
from flask import Flask, current_app
from sqlalchemy import and_, or_
from sqlalchemy.orm import load_only
from sqlalchemy.sql.expression import false
from configs import dify_config
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_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.id: segment
for segment in db.session.query(DocumentSegment)
.filter(
(
DocumentSegment.index_node_id.in_(index_node_ids)
| DocumentSegment.id.in_([chunk.segment_id for chunk in child_chunks])
),
and_(
filter_expr,
DocumentSegment.enabled == True,
DocumentSegment.status == "completed",
)
)
.options(
load_only(
DocumentSegment.id,