feat: Adding support to return additional features from vector retrieval for Milvus db

[franciscojavierarceo]

What this PR does / why we need it:

This PR adds support for returning multiple fields when doing vector similarity search with Milvus.

The behavior is now as simple as:

# Define your FeatureView
document_embeddings = FeatureView(
    name="embedded_documents",
    entities=[item, author],
    schema=[
       # notice that we define key index and metric search information at the field level
        Field(
            name="vector",
            dtype=Array(Float32),
            vector_index=True,
            vector_search_metric="COSINE",
        ),
        Field(name="item_id", dtype=Int64),
        Field(name="author_id", dtype=String),
        Field(name="sentence_chunks", dtype=String),
        Field(name="created_timestamp", dtype=UnixTimestamp),
        Field(name="event_timestamp", dtype=UnixTimestamp),
    ],
    source=rag_documents_source,
    ttl=timedelta(hours=24),
)

# Get the documents
result = store.retrieve_online_documents_v2(
    features=[
        "embedded_documents:vector",
        "embedded_documents:item_id",
        "embedded_documents:author_id",
        "embedded_documents:sentence_chunks",
    ],
    query=query_embedding,
    top_k=3,
).to_dict()

1. Feature Store Enhancements

   • Added retrieve_online_documents_v2 method to FeatureStore for improved document retrieval using embeddings.
   • Modified retrieve_online_documents to enhance performance and support additional features.
   • Introduced _retrieve_from_online_store_v2 to support the new document retrieval method internally.
   • Added _get_feature_view_vector_field_metadata to fetch metadata for vector fields within a FeatureView.

2. Utility Functions

   • Added _get_unique_entities_from_values to obtain unique composite entities and their indexes.
   • Introduced _populate_response_from_feature_data_v2 to populate responses with feature data.
   • Enhanced _extract_proto_values_to_dict to handle serialization to string and vector columns.

3. Provider and Online Store Updates

   • Updated PassthroughProvider and Provider interfaces to include retrieve_online_documents_v2.
   • Enhanced MilvusOnlineStore and OnlineStore with support for the new document retrieval method and vector field handling.
   • Modified MilvusOnlineStoreConfig to default to COSINE metric type for vector searches.

4. Testing and Example Repositories

   • Updated test cases in test_online_retrieval.py to include scenarios for the new document retrieval method.
   • Enhanced example repositories to demonstrate the use of vector fields and the new retrieval method.

Key Modifications

• sdk/python/feast/feature_store.py

  • Added Field import.
  • Introduced retrieve_online_documents_v2 method for enhanced document retrieval.
  • Updated _retrieve_from_online_store to handle vector field metadata.
  • Added _retrieve_from_online_store_v2 for internal support of the new retrieval method.

• sdk/python/feast/infra/online_stores/milvus_online_store/milvus.py

  • Added support for deserializing entity keys.
  • Modified _get_collection to handle vector fields.
  • Enhanced online_write_batch to serialize vector fields and handle entity keys.
  • Introduced retrieve_online_documents_v2 for Milvus online store.

• sdk/python/feast/infra/online_stores/online_store.py

  • Added retrieve_online_documents_v2 method for online document retrieval using embeddings.

• sdk/python/feast/infra/passthrough_provider.py

  • Added retrieve_online_documents_v2 method to passthrough provider.

• sdk/python/feast/infra/provider.py

  • Added retrieve_online_documents_v2 abstract method to provider interface.

• sdk/python/feast/utils.py

  • Introduced _get_unique_entities_from_values utility function.
  • Added _populate_response_from_feature_data_v2 for populating response with feature data.
  • Enhanced _extract_proto_values_to_dict to handle vector columns and serialization.

• sdk/python/tests/example_repos/example_rag_feature_repo.py

  • Added author_id entity and updated schema to include vector fields.

• sdk/python/tests/foo_provider.py

  • Added retrieve_online_documents_v2 method to the provider.

• sdk/python/tests/unit/online_store/test_online_retrieval.py

  • Updated test cases to include scenarios for retrieve_online_documents_v2.

Which issue(s) this PR fixes:

Another one for #4364

Misc

N/A

[franciscojavierarceo]

@HaoXuAI take a look at this PR, I've moved the retrieval implementation to retrieve_online_documents_v2 instead of using retrieve_online_documents.

My recommendation is we eventually wipe both of these methods and embed retrieve_online_documents_v2 in get_online_features and have retrieve_online_documents call get_online_features.

Based on the Milvus documentation and allowing more features to be returned, it makes sense to just incorporate this into get_online_features in my opinion.

[HaoXuAI]

Might be better to make it an enum now since there are multiple values

[franciscojavierarceo]

Yeah sounds good, mind if I do that in a follow up PR?

[HaoXuAI]

yep, for sure

[shuchu]

lgtm