diff --git a/py/vector_database/faiss/faiss_client.py b/py/vector_database/faiss/faiss_client.py index 91e6c57fe30..24ae4107832 100644 --- a/py/vector_database/faiss/faiss_client.py +++ b/py/vector_database/faiss/faiss_client.py @@ -891,105 +891,83 @@ def _vector_addDocument( "createdDocuments": [], } - # loop through and embed new docs + # Embed each input CSV in one call, then preserve the existing per-Source + # dataset/vector artifacts used by removal and master-index rebuilding. for document in documentFileLocation: - # Create the DataFrame for every file dataset = self._load_dataset(dataset_location=document) - - # Change the unit of work - # From being the document - # To the individual source inside the document sources = dataset["Source"].unique().tolist() + directory = os.path.dirname(document) + effective_columns = ( + list(dataset.columns) + if columns_to_index is None or len(columns_to_index) == 0 + else columns_to_index + ) + keyword_params = dict(keyword_search_params or {}) + keyword_search = keyword_params.pop("keywordSearch", None) is True + prepared_sources = [] + for source_name in sources: source_dataset = dataset[dataset["Source"] == source_name].reset_index( drop=True ) - - # Get the directory path and the base filename without extension - directory, base_filename = os.path.split(document) - dataset_extension = ".parquet" - vector_extension = ".npy" - - if columns_to_index == None or len(columns_to_index) == 0: - columns_to_index = list(source_dataset.columns) - - # save the dataset, this is for efficiency after removing docs - new_file_path = os.path.join( - directory, source_name + "_dataset" + dataset_extension + if len(source_dataset) == 0: + continue + parts = source_dataset[effective_columns].astype(str) + source_dataset[target_column] = parts.apply( + lambda row: separator.join(row) + separator, axis=1 ) - - # if applicable, create the concatenated columns - if len(source_dataset) > 0: - # vectorized equivalent of mapping `_concatenate_columns`: build a - # target_column whose value for each row is - # `...` (note trailing separator, - # preserved for parity with the previous Dataset.map behavior). - parts = source_dataset[columns_to_index].astype(str) - source_dataset[target_column] = parts.apply( - lambda row: separator.join(row) + separator, axis=1 + if keyword_search: + source_dataset[target_column] = ( + self.keyword_engine.keyword_extraction( + input=list(source_dataset[target_column]), + insight_id=insight_id, + param_dict=keyword_params, + ) ) - - # transform chunks into keywords - if ( - keyword_search_params != None - and keyword_search_params.pop("keywordSearch", None) is True - ): - keywords_for_target_col = ( - self.keyword_engine.keyword_extraction( - input=list(source_dataset[target_column]), - insight_id=insight_id, - param_dict=keyword_search_params, - ) + vectors = self._embed_and_validate( + list(source_dataset[target_column]), insight_id + ) + createDocumentsResponse["createdDocuments"].extend( + self._persist_source_artifacts( + directory, + source_name, + source_dataset, + vectors, + columns_to_remove, + target_column, ) - source_dataset[target_column] = keywords_for_target_col - - # get the embeddings for the document - vectors = self.embeddings_engine.embeddings( - strings_to_embed=list(source_dataset[target_column]), - insight_id=insight_id, ) - vectors = np.array(vectors[0]["response"], dtype=np.float32) - assert vectors.ndim == 2 + self._append_vectors(vectors) + else: + prepared_sources.append((source_name, source_dataset)) - columns_to_remove.append(target_column) - columns_to_drop = list( - set(columns_to_remove).intersection(set(source_dataset.columns)) + if prepared_sources: + all_text = [ + value + for _source_name, source_dataset in prepared_sources + for value in source_dataset[target_column].tolist() + ] + all_vectors = self._embed_and_validate(all_text, insight_id) + offset = 0 + for source_name, source_dataset in prepared_sources: + source_rows = len(source_dataset) + source_vectors = all_vectors[offset : offset + source_rows] + offset += source_rows + createDocumentsResponse["createdDocuments"].extend( + self._persist_source_artifacts( + directory, + source_name, + source_dataset, + source_vectors, + columns_to_remove, + target_column, + ) ) - source_dataset = source_dataset.drop(columns=columns_to_drop) - - # Keep per-source files dtype-consistent so subsequent - # `_validateEmbeddingFiles` concatenations don't produce - # mixed-type columns that fail pyarrow serialization. - self._enforce_canonical_schema(source_dataset) - source_dataset.to_parquet(new_file_path, index=False) - - # add the created source_dataset file path - createDocumentsResponse["createdDocuments"].append(new_file_path) - - # normalize the vectors if using huggingface - if type(self.tokenizer).__name__ == "HuggingfaceTokenizer": - faiss.normalize_L2(vectors) - - # write out the vectors as a .npy file (no pickle) - new_file_path = os.path.join( - directory, - source_name + "_vectors" + vector_extension, + if offset != len(all_vectors): + raise ValueError( + "Embedding rows could not be assigned to their Sources" ) - np.save(new_file_path, vectors, allow_pickle=False) - - # add the created embeddings file path - createDocumentsResponse["createdDocuments"].append(new_file_path) - - # TODO need to update the flow for how we instatiate - if self.encoded_vectors is None: - self.encoded_vectors = np.copy(vectors) - self.vector_dimensions = self.encoded_vectors.shape - else: - # make sure the dimensions are the same - assert self.vector_dimensions[1] == vectors.shape[1] - self.encoded_vectors = np.concatenate( - [self.encoded_vectors, vectors], axis=0 - ) + self._append_vectors(all_vectors) master_indexClass_files, corrupted_file_sets = self.createMasterFiles( path_to_files=os.path.dirname(os.path.dirname(documentFileLocation[0])) @@ -1003,6 +981,82 @@ def _vector_addDocument( return createDocumentsResponse + def _embed_and_validate( + self, strings_to_embed: List[str], insight_id: Optional[str] + ) -> np.ndarray: + """Embed one bounded CSV batch and validate it before writing artifacts.""" + response = self.embeddings_engine.embeddings( + strings_to_embed=strings_to_embed, + insight_id=insight_id, + ) + try: + vectors = np.asarray(response[0]["response"], dtype=np.float32) + except (IndexError, KeyError, TypeError, ValueError) as error: + raise ValueError( + "Embedding response does not contain a numeric response matrix" + ) from error + if ( + vectors.ndim != 2 + or vectors.shape[0] != len(strings_to_embed) + or vectors.shape[1] == 0 + ): + raise ValueError( + "Embedding response shape does not match the submitted rows: " + f"expected {len(strings_to_embed)} rows, received {vectors.shape}" + ) + if ( + self.vector_dimensions is not None + and self.vector_dimensions[1] != vectors.shape[1] + ): + raise ValueError( + "Embedding dimensions do not match the existing FAISS index: " + f"expected {self.vector_dimensions[1]}, received {vectors.shape[1]}" + ) + if type(self.tokenizer).__name__ == "HuggingfaceTokenizer": + faiss.normalize_L2(vectors) + return vectors + + def _persist_source_artifacts( + self, + directory: str, + source_name: str, + source_dataset: pd.DataFrame, + vectors: np.ndarray, + columns_to_remove: Optional[List[str]], + target_column: str, + ) -> List[str]: + """Write the stable dataset and vector pair for one Source.""" + if len(source_dataset) != len(vectors): + raise ValueError( + f"Source {source_name} has {len(source_dataset)} rows but {len(vectors)} vectors" + ) + columns_to_drop = list( + set([*(columns_to_remove or []), target_column]).intersection( + set(source_dataset.columns) + ) + ) + stored_dataset = source_dataset.drop(columns=columns_to_drop) + self._enforce_canonical_schema(stored_dataset) + dataset_path = os.path.join(directory, source_name + "_dataset.parquet") + vector_path = os.path.join(directory, source_name + "_vectors.npy") + stored_dataset.to_parquet(dataset_path, index=False) + np.save(vector_path, vectors, allow_pickle=False) + return [dataset_path, vector_path] + + def _append_vectors(self, vectors: np.ndarray) -> None: + """Maintain the in-memory vector shape until master files are rebuilt.""" + if self.encoded_vectors is None: + self.encoded_vectors = np.copy(vectors) + self.vector_dimensions = self.encoded_vectors.shape + return + if self.vector_dimensions[1] != vectors.shape[1]: + raise ValueError( + "Embedding dimensions do not match the existing FAISS index: " + f"expected {self.vector_dimensions[1]}, received {vectors.shape[1]}" + ) + self.encoded_vectors = np.concatenate([self.encoded_vectors, vectors], axis=0) + self.vector_dimensions = self.encoded_vectors.shape + def createMasterFiles(self, path_to_files: str) -> Tuple[str]: """ Create a master dataset and embeddings file based on the current documents. The main purpose of this is to improve startup runtime. diff --git a/py/vector_database/tests/test_faiss_batch_vector_csv.py b/py/vector_database/tests/test_faiss_batch_vector_csv.py new file mode 100644 index 00000000000..6b9196ebf77 --- /dev/null +++ b/py/vector_database/tests/test_faiss_batch_vector_csv.py @@ -0,0 +1,202 @@ +import os +import tempfile +import unittest +from pathlib import Path +from unittest.mock import Mock + +import numpy as np +import pandas as pd + +from vector_database.faiss.faiss_client import FAISSSearcher + + +class RecordingEmbedder: + def __init__(self): + self.calls = [] + + def embeddings(self, *, strings_to_embed, insight_id): + self.calls.append((list(strings_to_embed), insight_id)) + vectors = np.arange(len(strings_to_embed) * 3, dtype=np.float32).reshape(-1, 3) + return [{"response": vectors.tolist()}] + + +class RecordingKeywordEngine: + def __init__(self): + self.calls = [] + + def keyword_extraction(self, *, input, insight_id, param_dict): + self.calls.append((list(input), insight_id, dict(param_dict))) + return [f"keywords:{value}" for value in input] + + +class FaissBatchVectorCsvTests(unittest.TestCase): + def setUp(self): + self.temporary_directory = tempfile.TemporaryDirectory() + self.addCleanup(self.temporary_directory.cleanup) + self.indexed_files = ( + Path(self.temporary_directory.name) / "schema" / "default" / "indexed_files" + ) + self.indexed_files.mkdir(parents=True) + self.document = self.indexed_files / "batch.csv" + self.document.write_text("unused", encoding="utf-8") + + def searcher(self, dataset, embedder=None, keyword_engine=None): + searcher = object.__new__(FAISSSearcher) + searcher.embeddings_engine = embedder or RecordingEmbedder() + searcher.keyword_engine = keyword_engine or RecordingKeywordEngine() + searcher.tokenizer = object() + searcher.encoded_vectors = None + searcher.vector_dimensions = None + searcher._load_dataset = Mock(return_value=dataset.copy()) + searcher.createMasterFiles = Mock( + return_value=( + [ + str(self.indexed_files.parent / "vectors.npy"), + str(self.indexed_files.parent / "dataset.parquet"), + ], + [], + ) + ) + return searcher + + def test_embeds_one_csv_once_and_slices_stable_source_artifacts(self): + dataset = pd.DataFrame( + { + "Source": ["alpha.java", "alpha.java", "beta.java"], + "Modality": ["text", "text", "text"], + "Divider": ["a#L1-L1", "a#L2-L2", "b#L1-L1"], + "Part": [1, 2, 1], + "Tokens": [2, 2, 2], + "Content": ["alpha one", "alpha two", "beta one"], + } + ) + embedder = RecordingEmbedder() + searcher = self.searcher(dataset, embedder=embedder) + + response = searcher._vector_addDocument( + [str(self.document)], + ["Content"], + [], + "text", + "|", + {}, + "insight-1", + ) + + self.assertEqual(len(embedder.calls), 1) + self.assertEqual( + embedder.calls[0], + (["alpha one|", "alpha two|", "beta one|"], "insight-1"), + ) + alpha_vectors = np.load( + self.indexed_files / "alpha.java_vectors.npy", allow_pickle=False + ) + beta_vectors = np.load( + self.indexed_files / "beta.java_vectors.npy", allow_pickle=False + ) + np.testing.assert_array_equal(alpha_vectors, np.arange(6).reshape(2, 3)) + np.testing.assert_array_equal(beta_vectors, np.arange(6, 9).reshape(1, 3)) + self.assertTrue((self.indexed_files / "alpha.java_dataset.parquet").is_file()) + self.assertTrue((self.indexed_files / "beta.java_dataset.parquet").is_file()) + self.assertEqual(searcher.createMasterFiles.call_count, 1) + searcher.createMasterFiles.assert_called_once_with( + path_to_files=str(self.indexed_files.parent) + ) + self.assertIn( + str(self.indexed_files / "alpha.java_dataset.parquet"), + response["createdDocuments"], + ) + self.assertIn( + str(self.indexed_files / "alpha.java_vectors.npy"), + response["createdDocuments"], + ) + + def test_oversized_single_source_is_embedded_in_one_call(self): + row_count = 368 + dataset = pd.DataFrame( + { + "Source": ["Parser.java"] * row_count, + "Content": [f"chunk {index}" for index in range(row_count)], + } + ) + embedder = RecordingEmbedder() + searcher = self.searcher(dataset, embedder=embedder) + + searcher._vector_addDocument( + [str(self.document)], ["Content"], [], "text", "|", {}, "insight-1" + ) + + self.assertEqual(len(embedder.calls), 1) + self.assertEqual(len(embedder.calls[0][0]), row_count) + vectors = np.load( + self.indexed_files / "Parser.java_vectors.npy", allow_pickle=False + ) + self.assertEqual(vectors.shape, (row_count, 3)) + + def test_invalid_embedding_row_count_writes_no_source_artifacts(self): + class ShortEmbedder: + def embeddings(self, *, strings_to_embed, insight_id): + del insight_id + return [{"response": [[1.0, 2.0]] * (len(strings_to_embed) - 1)}] + + dataset = pd.DataFrame( + {"Source": ["alpha.java", "beta.java"], "Content": ["alpha", "beta"]} + ) + searcher = self.searcher(dataset, embedder=ShortEmbedder()) + + with self.assertRaisesRegex(ValueError, "shape does not match"): + searcher._vector_addDocument( + [str(self.document)], ["Content"], [], "text", "|", {}, "insight-1" + ) + + self.assertEqual(list(self.indexed_files.glob("*_dataset.parquet")), []) + self.assertEqual(list(self.indexed_files.glob("*_vectors.npy")), []) + searcher.createMasterFiles.assert_not_called() + + def test_keyword_search_remains_a_per_source_fallback(self): + dataset = pd.DataFrame( + { + "Source": ["alpha.java", "alpha.java", "beta.java"], + "Content": ["alpha one", "alpha two", "beta one"], + } + ) + embedder = RecordingEmbedder() + keyword_engine = RecordingKeywordEngine() + searcher = self.searcher( + dataset, embedder=embedder, keyword_engine=keyword_engine + ) + + searcher._vector_addDocument( + [str(self.document)], + ["Content"], + [], + "text", + "|", + {"keywordSearch": True, "limit": 4}, + "insight-1", + ) + + self.assertEqual(len(keyword_engine.calls), 2) + self.assertEqual(len(embedder.calls), 2) + self.assertEqual(keyword_engine.calls[0][2], {"limit": 4}) + self.assertEqual(keyword_engine.calls[1][2], {"limit": 4}) + + def test_per_source_artifact_names_remain_removal_compatible(self): + dataset = pd.DataFrame({"Source": ["alpha.java"], "Content": ["alpha one"]}) + searcher = self.searcher(dataset) + searcher._vector_addDocument( + [str(self.document)], ["Content"], [], "text", "|", {}, "insight-1" + ) + + dataset_path = self.indexed_files / "alpha.java_dataset.parquet" + vector_path = self.indexed_files / "alpha.java_vectors.npy" + self.assertTrue(dataset_path.is_file()) + self.assertTrue(vector_path.is_file()) + os.remove(dataset_path) + os.remove(vector_path) + self.assertFalse(dataset_path.exists()) + self.assertFalse(vector_path.exists()) + + +if __name__ == "__main__": + unittest.main()