import inspect
import os
import warnings
from pathlib import Path
from typing import TYPE_CHECKING, Callable, Optional, Protocol, runtime_checkable
import pandas as pd
from feast.chunker import BaseChunker, TextChunker
from feast.embedder import BaseEmbedder, MultiModalEmbedder
if TYPE_CHECKING:
from feast.feature_store import FeatureStore
@runtime_checkable
class SchemaTransformFn(Protocol):
"""
Protocol defining the structure for schema transform functions.
The schema transform converts the output of Chunker + Embedder
into the format expected by the FeatureView schema.
"""
def __call__(self, df: pd.DataFrame) -> pd.DataFrame:
"""
Transform chunked + embedded DataFrame to FeatureView schema.
Args:
df: Input DataFrame with chunks and embeddings.
Returns:
DataFrame with columns matching FeatureView schema.
"""
...
def default_schema_transform_fn(df: pd.DataFrame) -> pd.DataFrame:
"""
Default schema transform function that transforms the output of Chunker + Embedder
into the format expected by the FeatureView schema.
"""
from datetime import datetime, timezone
return pd.DataFrame(
{
"passage_id": df["chunk_id"],
"text": df["text"],
"embedding": df["text_embedding"],
"event_timestamp": [datetime.now(timezone.utc)] * len(df),
"source_id": df["original_id"],
}
)
def generate_repo_file(
repo_path: str,
feature_view_name: str = "text_feature_view",
vector_length: int = 384,
) -> str:
"""
Generate a Python file with Entity and FeatureView definitions.
This file is compatible with `feast apply` CLI.
Args:
repo_path: Path to the feature repo directory.
feature_view_name: Name of the feature view to create.
vector_length: Dimension of the embedding vectors. Should match the
output dimension of the embedding model being used. Defaults to
384 (matching the default all-MiniLM-L6-v2 model).
Returns:
Path to generated file.
"""
from feast.repo_operations import is_valid_name
if not is_valid_name(feature_view_name) or not feature_view_name.isidentifier():
raise ValueError(
f"feature_view_name '{feature_view_name}' is invalid. "
"It should only contain alphanumeric characters, underscores, "
"and must not start with an underscore."
)
code = f'''"""
Auto-generated by DocEmbedder.
Compatible with `feast apply` CLI.
"""
from datetime import timedelta
from feast import Entity, FeatureView, Field, FileSource
from feast.types import Array, Float32, String, ValueType
# Entity
text_entity = Entity(
name="passage_id",
join_keys=["passage_id"],
description="Passage identifier",
value_type=ValueType.STRING,
)
# Source
{feature_view_name.replace(" ", "_").replace("-", "_")}_source = FileSource(
name="{feature_view_name}_source",
path="data/{feature_view_name}.parquet",
timestamp_field="event_timestamp",
)
# FeatureView
{feature_view_name.replace(" ", "_").replace("-", "_")} = FeatureView(
name="{feature_view_name}",
entities=[text_entity],
ttl=timedelta(days=1),
schema=[
Field(
name="text",
dtype=String,
description="Document text content",
),
Field(
name="embedding",
dtype=Array(Float32),
description="Vector embedding",
vector_index=True,
vector_length={vector_length},
vector_search_metric="COSINE",
),
Field(
name="source_id",
dtype=String,
description="Source ID",
),
],
source={feature_view_name.replace(" ", "_").replace("-", "_")}_source,
online=True,
)
'''
filepath = os.path.join(
repo_path, feature_view_name.replace(" ", "_").replace("-", "_") + ".py"
)
with open(filepath, "w") as f:
f.write(code)
return filepath
class DocEmbedder:
"""
DocEmbedder is a class that embeds documents and chunks them into a format expected by the FeatureView schema using a Logic Implementation By the user.
Args:
repo_path: Path to the feature repo (can be "." for current directory).
yaml_file: Name of the feature_store.yaml file inside repo_path.
Defaults to "feature_store.yaml".
feature_view_name: Name of the feature view to create.
chunker: Chunker to use for chunking the documents.
embedder: Embedder to use for embedding the documents.
schema_transform_fn: Schema transform function to use for transforming the output of the chunker and embedder into the format expected by the FeatureView schema.
create_feature_view: Whether to create a feature view in the feature repo. By default it will generate a Python file with the FeatureView definition.
vector_length: Explicit embedding dimension for the generated FeatureView schema.
If None (default), the dimension is auto-detected from the embedder
via ``get_embedding_dim("text")``. Falls back to 384 if detection
is not supported by the embedder.
auto_apply_repo: Whether to apply the repository automatically. By default it will apply the repository after creating the feature view.
"""
def __init__(
self,
repo_path: str,
yaml_file: str = "feature_store.yaml",
feature_view_name: str = "text_feature_view",
chunker: Optional[BaseChunker] = None,
embedder: Optional[BaseEmbedder] = None,
schema_transform_fn: SchemaTransformFn = default_schema_transform_fn,
create_feature_view: bool = True,
vector_length: Optional[int] = None,
auto_apply_repo: bool = True,
):
self.repo_path = repo_path
self.yaml_path = os.path.join(Path(repo_path).resolve(), yaml_file)
self.feature_view_name = feature_view_name
self.chunker = chunker or TextChunker()
self.embedder = embedder or MultiModalEmbedder()
self.store: Optional[FeatureStore] = None
sig = inspect.signature(schema_transform_fn)
params = list(sig.parameters.values())
if (
len(params) != 1
or params[0].annotation != pd.DataFrame
or sig.return_annotation != pd.DataFrame
):
raise ValueError(
"schema_transform_fn must be a function that takes a DataFrame and returns a DataFrame"
)
self.schema_transform_fn = schema_transform_fn
if create_feature_view:
resolved_vector_length = self._resolve_vector_length(vector_length, "text")
generate_repo_file(
repo_path=repo_path,
feature_view_name=feature_view_name,
vector_length=resolved_vector_length,
)
if auto_apply_repo:
self.apply_repo()
def _resolve_vector_length(
self, explicit_length: Optional[int], modality: str
) -> int:
"""
Determine the vector length to use for the generated FeatureView.
Priority:
1. Explicitly provided vector_length
2. Auto-detected from embedder via get_embedding_dim("text")
3. Default of 384 (matching all-MiniLM-L6-v2)
Args:
explicit_length: User-provided vector length, or None.
Returns:
The resolved vector length as an integer.
"""
_DEFAULT_VECTOR_LENGTH = 384
if explicit_length is not None:
return explicit_length
try:
dim = self.embedder.get_embedding_dim(modality)
if dim is not None:
return dim
except Exception:
pass
return _DEFAULT_VECTOR_LENGTH
def save_to_online_store(self, df: pd.DataFrame, feature_view_name: str) -> None:
"""
Save the embedded documents to the online store.
"""
from feast.feature_store import FeatureStore
if self.store is None:
self.store = FeatureStore(repo_path=self.repo_path)
self.store.write_to_online_store(
feature_view_name=feature_view_name,
df=df,
)
# TODO (Future scope): Implement save_to_offline_store to write embedded
# documents to the offline store. Currently blocked by DaskOfflineStore
# .offline_write_batch not creating the parquet file if it does not exist.
# Once that is fixed, add a method that calls:
# store = FeatureStore(repo_path=self.repo_path)
# store.write_to_offline_store(feature_view_name=feature_view_name, df=df)
def apply_repo(self) -> None:
"""
Apply the repository to register feature views in the registry.
"""
from feast.repo_config import load_repo_config
from feast.repo_operations import apply_total
original_cwd = None
try:
original_cwd = os.getcwd()
repo_path = Path(self.repo_path).resolve()
config = load_repo_config(
repo_path=repo_path,
fs_yaml_file=Path(self.yaml_path),
)
apply_total(
repo_config=config,
repo_path=repo_path,
skip_source_validation=True,
)
finally:
if original_cwd is not None:
os.chdir(original_cwd)
def embed_documents(
self,
documents: pd.DataFrame,
id_column: str,
source_column: str,
type_column: Optional[str] = None,
column_mapping: Optional[tuple[str, str]] = None,
custom_schema_transform_fn: Optional[
Callable[[pd.DataFrame], pd.DataFrame]
] = None,
) -> pd.DataFrame:
"""
Embed a list of documents and chunk them into a format expected by the FeatureView schema using a Logic Implementation By the user and save the DataFrame to the online store.
Args:
documents: DataFrame containing the documents to embed.
id_column: Column name containing the document IDs.
source_column: Column name containing the document sources.
type_column: Column name containing the document types.
column_mapping: Tuple mapping source columns to (modality, output column).
custom_schema_transform_fn: Custom schema transform function to use for transforming the output of the chunker and embedder into the format expected by the FeatureView schema.
Returns:
DataFrame with the embedded documents.
Example:
documents = pd.DataFrame({
"id": [1, 2, 3],
"source": ["source1", "source2", "source3"],
"type": ["type1", "type2", "type3"],
"text": ["text1", "text2", "text3"],
})
column_mapping = ("text", "text_embedding")
df = embed_documents(documents=documents, id_column="id", source_column="source", type_column="type", column_mapping=column_mapping)
"""
if custom_schema_transform_fn is not None:
sig = inspect.signature(custom_schema_transform_fn)
params = list(sig.parameters.values())
if (
len(params) != 1
or params[0].annotation != pd.DataFrame
or sig.return_annotation != pd.DataFrame
):
raise ValueError(
"custom_schema_transform_fn must be a function that takes a DataFrame and returns a DataFrame"
)
current_schema_transform_fn = (
custom_schema_transform_fn
if custom_schema_transform_fn is not None
else self.schema_transform_fn
)
if column_mapping is None:
column_mapping = ("text", "text_embedding")
if (
current_schema_transform_fn is default_schema_transform_fn
and column_mapping[0] == "text"
and (source_column != "text" or column_mapping[1] != "text_embedding")
):
raise ValueError(
f"source_column='{source_column}' with output column='{column_mapping[1]}' "
f"is not compatible with default_schema_transform_fn, which expects "
f"source_column='text' and column_mapping=('text', 'text_embedding'). "
f"Provide a custom schema_transform_fn."
)
if column_mapping[0] == "text":
df = self.chunker.chunk_dataframe(
df=documents,
id_column=id_column,
source_column=source_column,
type_column=type_column,
)
else:
df = documents
df = self.embedder.embed_dataframe(
df, column_mapping={source_column: column_mapping}
)
if (
column_mapping[0] == "text"
or current_schema_transform_fn is not default_schema_transform_fn
):
df = current_schema_transform_fn(df)
else:
warnings.warn(
f"Modality '{column_mapping[0]}' is not supported by the default schema transform function. "
f"The output DataFrame will be passed directly to the online store. "
f"Ensure your FeatureView schema matches the output columns. "
f"You can provide a custom schema transform function to handle this.",
UserWarning,
stacklevel=2,
)
self.save_to_online_store(df=df, feature_view_name=self.feature_view_name)
return df