pyflink.dataframe.ai.llm.LLMAccessor.vector_search#

LLMAccessor.vector_search(search_source: DataFrame, column_to_search: str, column_to_query: Union[str, Expression], top_k: int, *, agg: bool = False, output_columns: Optional[Union[str, List[str]]] = None, config: Optional[Mapping[str, str]] = None, ignore_empty_and_null: bool = False) → DataFrame[source]#

Join this DataFrame with vector search results from another source.

Parameters

search_source – Search source as a DataFrame.
column_to_search – Vector column in the search source.
column_to_query – Query vector column or direct column reference from this DataFrame.
top_k – Number of nearest matches to return per input row.
agg – Whether to aggregate matches into one array column.
output_columns – Optional output column name or names.
config – Runtime vector search config.
ignore_empty_and_null – Whether to drop rows with no search output.

Returns

A new DataFrame with vector search outputs appended.

Note

Let N be the number of columns in search_source.

In non-aggregation mode (agg=False), the result may contain up to top_k rows for every input row and appends N + 1 columns: the N columns from search_source with their original types, followed by one score column of type DOUBLE. If output_columns is specified, it must contain exactly N + 1 names and is used as the names of these appended columns in order.

In aggregation mode (agg=True), the result contains one row for every input row and appends one array column. The array element is a row with N + 1 nested columns: the N columns from search_source with their original types, followed by one score field of type DOUBLE. If output_columns is specified, it must contain exactly one name and is used as the name of the appended array column.

Example:

>>> import pyflink.dataframe as pf
>>> from pyflink.table import DataTypes
>>> queries = pf.from_records(
...     [(1, [0.1, 0.2]), (2, [0.3, 0.4])],
...     schema=["query_id", "query_embedding"],
... )
>>> documents = pf.read_milvus(
...     endpoint="http://milvus.example.com",
...     username="${secret_values.milvus_user}",
...     password="${secret_values.milvus_password}",
...     database_name="my_db",
...     collection_name="documents",
...     schema={
...         "doc_id": DataTypes.BIGINT(),
...         "embedding": DataTypes.ARRAY(DataTypes.FLOAT()),
...         "title": DataTypes.STRING(),
...     },
...     columns=["doc_id", "embedding", "title"],
... )
>>> result = queries.llm.vector_search(
...     documents,
...     column_to_search="embedding",
...     column_to_query="query_embedding",
...     top_k=5,
...     output_columns=[
...         "match_doc_id",
...         "match_embedding",
...         "match_title",
...         "score",
...     ],
... )