Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Improve os read #1065

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Improve os read #1065

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5c80ea5
Fetch min set of fields in read_records for document reconstruct case
austintlee Dec 6, 2024
b5ee277
Restore assert
austintlee Dec 6, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion lib/sycamore/sycamore/connectors/opensearch/opensearch_reader.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import copy
import logging

from sycamore.data import Document, Element
Expand Down Expand Up @@ -47,6 +48,9 @@ def read_records(self, query_params: BaseDBReader.QueryParams) -> "OpenSearchRea
if "size" not in query_params.query and "size" not in query_params.kwargs:
query_params.kwargs["size"] = 200
result = []
# We only fetch the minimum required fields for full document retrieval/reconstruction
if query_params.reconstruct_document:
query_params.kwargs["_source_includes"] = "doc_id,parent_id,properties"
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've looked into only fetching "parent_id", but the problem with only getting child documents here is that we also need parent docs (which do not have "parent_id") and their properties as during explode() we do not transfer all properties from parent to child docs:

sycamore/lib/sycamore/sycamore/transforms/explode.py

Lines 49 to 51 in 8471d43

for doc_property in parent.properties.keys():
if doc_property.startswith("_"):
cur.properties[doc_property] = parent.properties[doc_property]

Otherwise, we would have to make another round of calls to OpenSearch just to fetch parent docs which may be less efficient and more time consuming than just fetching doc_id, parent_id and properties here.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't understand why you need properties. At least in my testing; I was getting parent ids in the _id field.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we not need to reconstruct all of the properties? If that is the case, I think we can just reconstruct using only the child elements.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the test I ran to verify correctness:

sycamore/lib/sycamore/sycamore/tests/integration/connectors/opensearch/test_opensearch_read.py

Line 59 in 4b69ea9

def test_ingest_and_read(self, setup_index, exec_mode):

# No pagination needed for knn queries
if "query" in query_params.query and "knn" in query_params.query["query"]:
response = self._client.search(
Expand Down Expand Up @@ -148,6 +152,8 @@ def to_docs(self, query_params: "BaseDBReader.QueryParams") -> list[Document]:

# Batched retrieval of all elements belong to unique docs
doc_ids = list(unique_docs.keys())
# We can't safely exclude embeddings since we might need them for 'rerank', e.g.
# We will need the Planner to determine that and pass that info to the reader.
all_elements_for_docs = self._get_all_elements_for_doc_ids(doc_ids, query_params.index_name)

"""
Expand Down Expand Up @@ -183,7 +189,6 @@ def _get_all_elements_for_doc_ids(self, doc_ids: list[str], index: str) -> list[
"""
batch_size = 100
page_size = 500

all_elements = []
for i in range(0, len(doc_ids), batch_size):
doc_ids_batch = doc_ids[i : i + batch_size]
Expand Down
4 changes: 3 additions & 1 deletion lib/sycamore/sycamore/reader.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -224,7 +224,7 @@ def opensearch(
index_name: str,
query: Optional[Dict] = None,
reconstruct_document: bool = False,
query_kwargs: dict[str, Any] = {},
query_kwargs=None,
**kwargs,
) -> DocSet:
"""
Expand Down Expand Up @@ -287,6 +287,8 @@ def opensearch(
OpenSearchReaderQueryParams,
)

if query_kwargs is None:
query_kwargs = {}
client_params = OpenSearchReaderClientParams(os_client_args=os_client_args)
query_params = (
OpenSearchReaderQueryParams(
Expand Down
Loading