diff --git a/README.md b/README.md index 0a78b4b..e4d6200 100644 --- a/README.md +++ b/README.md @@ -44,6 +44,7 @@ See the [Contributing](#contributing) section to add support for your favorite v | DataStax Astra DB | ✅ | ✅ | | Chroma | ✅ | ✅ | | Turbopuffer | ✅ | ✅ | +| Weaviate | ✅ | ✅ | @@ -67,7 +68,6 @@ See the [Contributing](#contributing) section to add support for your favorite v | Vector Database | Import | Export | |--------------------------------|--------|--------| | Azure AI Search | ❌ | ❌ | -| Weaviate | ❌ | ❌ | | MongoDB Atlas | ❌ | ❌ | | OpenSearch | ❌ | ❌ | | Apache Cassandra | ❌ | ❌ | @@ -161,7 +161,7 @@ usage: export_vdf [-h] [-m MODEL_NAME] [--max_file_size MAX_FILE_SIZE] [--push_to_hub | --no-push_to_hub] [--public | --no-public] - {pinecone,qdrant,kdbai,milvus,vertexai_vectorsearch} + {pinecone,qdrant,kdbai,milvus,vertexai_vectorsearch,weaviate} ... Export data from various vector databases to the VDF format for vector datasets @@ -190,6 +190,7 @@ Vector Databases: vertexai_vectorsearch Export data from Vertex AI Vector Search + weaviate Export data from Weaviate ``` ## Import script @@ -198,7 +199,7 @@ Vector Databases: import_vdf --help usage: import_vdf [-h] [-d DIR] [-s | --subset | --no-subset] [--create_new | --no-create_new] - {milvus,pinecone,qdrant,vertexai_vectorsearch,kdbai} + {milvus,pinecone,qdrant,vertexai_vectorsearch,kdbai,weaviate} ... Import data from VDF to a vector database @@ -221,6 +222,7 @@ Vector Databases: vertexai_vectorsearch Import data to Vertex AI Vector Search kdbai Import data to KDB.AI + weaviate Import data to Weaviate ``` ## Re-embed script diff --git a/src/vdf_io/export_vdf/weaviate_export.py b/src/vdf_io/export_vdf/weaviate_export.py index 518dd3c..bb4b42e 100644 --- a/src/vdf_io/export_vdf/weaviate_export.py +++ b/src/vdf_io/export_vdf/weaviate_export.py @@ -1,15 +1,24 @@ +import json import os +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq from tqdm import tqdm -import weaviate +from vdf_io.constants import DEFAULT_BATCH_SIZE, ID_COLUMN from vdf_io.export_vdf.vdb_export_cls import ExportVDB from vdf_io.names import DBNames -from vdf_io.util import set_arg_from_input, set_arg_from_password - -# Set these environment variables -URL = os.getenv("YOUR_WCS_URL") -APIKEY = os.getenv("YOUR_WCS_API_KEY") +from vdf_io.util import set_arg_from_input +from vdf_io.weaviate_util import ( + collection_names, + connect_weaviate, + first_vector_dimension, + get_weaviate_distance, + make_weaviate_parser, + normalize_weaviate_vectors, + serialize_weaviate_config, +) class ExportWeaviate(ExportVDB): @@ -21,69 +30,134 @@ def make_parser(cls, subparsers): cls.DB_NAME_SLUG, help="Export data from Weaviate" ) - parser_weaviate.add_argument("--url", type=str, help="URL of Weaviate instance") - parser_weaviate.add_argument("--api_key", type=str, help="Weaviate API key") + make_weaviate_parser(parser_weaviate) + parser_weaviate.add_argument( + "--classes", + type=str, + help="Collections/classes to export (comma-separated)", + ) parser_weaviate.add_argument( - "--classes", type=str, help="Classes to export (comma-separated)" + "--batch_size", + type=int, + help="Batch size for exporting data", + default=DEFAULT_BATCH_SIZE, ) @classmethod def export_vdb(cls, args): - set_arg_from_input( - args, - "url", - "Enter the URL of Weaviate instance: ", - str, - ) - set_arg_from_password( - args, - "api_key", - "Enter the Weaviate API key: ", - "WEAVIATE_API_KEY", - ) weaviate_export = ExportWeaviate(args) - weaviate_export.all_classes = list( - weaviate_export.client.collections.list_all().keys() - ) + weaviate_export.all_classes = weaviate_export.get_all_index_names() set_arg_from_input( weaviate_export.args, "classes", - "Enter the name of the classes to export (comma-separated, all will be exported by default): ", + "Enter the name of the collections/classes to export (comma-separated, all will be exported by default): ", str, choices=weaviate_export.all_classes, ) + set_arg_from_input( + weaviate_export.args, + "batch_size", + f"Enter the batch size for exporting data (default: {DEFAULT_BATCH_SIZE}): ", + int, + DEFAULT_BATCH_SIZE, + ) weaviate_export.get_data() return weaviate_export - # Connect to a WCS instance def __init__(self, args): super().__init__(args) - self.client = weaviate.connect_to_wcs( - cluster_url=self.args["url"], - auth_credentials=weaviate.auth.AuthApiKey(self.args["api_key"]), - skip_init_checks=True, - ) + self.client = connect_weaviate(self.args) + + def get_all_index_names(self): + return collection_names(self.client) def get_index_names(self): if self.args.get("classes") is None: - return self.all_classes + return self.get_all_index_names() else: input_classes = self.args["classes"].split(",") - if set(input_classes) - set(self.all_classes): + all_classes = self.get_all_index_names() + if set(input_classes) - set(all_classes): tqdm.write( - f"These classes are not present in the Weaviate instance: {set(input_classes) - set(self.all_classes)}" + f"These collections/classes are not present in the Weaviate instance: {set(input_classes) - set(all_classes)}" ) - return [c for c in self.all_classes if c in input_classes] + return [c for c in all_classes if c in input_classes] def get_data(self): - # Get all objects of a class + index_metas = {} index_names = self.get_index_names() - for class_name in index_names: + for class_name in tqdm(index_names, desc="Exporting collections"): + rows = [] + total_exported = 0 + dimensions = -1 + vector_columns = [] collection = self.client.collections.get(class_name) + collection_config = collection.config.get() response = collection.aggregate.over_all(total_count=True) - print(f"{response.total_count=}") + total = response.total_count or 0 + vectors_directory = self.create_vec_dir(class_name) + batch_size = self.args.get("batch_size") or DEFAULT_BATCH_SIZE + + for item in tqdm( + collection.iterator(include_vector=True, cache_size=batch_size), + desc=f"Exporting {class_name}", + total=total, + ): + item_vectors = normalize_weaviate_vectors(item.vector) + if not item_vectors: + continue + if not vector_columns: + vector_columns = list(item_vectors.keys()) + dimensions = first_vector_dimension(item_vectors[vector_columns[0]]) + + row = {ID_COLUMN: str(item.uuid)} + row.update(item_vectors) + row.update(dict(item.properties or {})) + rows.append(row) + + if len(rows) >= batch_size: + total_exported += self.save_rows_to_parquet(rows, vectors_directory) + rows = [] - # objects = self.client.query.get( - # wvq.Objects(wvq.Class(class_name)).with_limit(1000) - # ) - # print(objects) + if rows: + total_exported += self.save_rows_to_parquet(rows, vectors_directory) + + if not vector_columns: + vector_columns = ["vector"] + + namespace_meta = self.get_namespace_meta( + class_name, + vectors_directory, + total=total, + num_vectors_exported=total_exported, + dim=dimensions, + index_config=serialize_weaviate_config(collection_config), + vector_columns=vector_columns, + distance=get_weaviate_distance(collection_config, vector_columns[0]), + ) + index_metas[class_name] = [namespace_meta] + self.args["exported_count"] += total_exported + + self.file_structure.append(os.path.join(self.vdf_directory, "VDF_META.json")) + internal_metadata = self.get_basic_vdf_meta(index_metas) + meta_text = json.dumps(internal_metadata.model_dump(), indent=4) + tqdm.write(meta_text) + with open(os.path.join(self.vdf_directory, "VDF_META.json"), "w") as json_file: + json_file.write(meta_text) + return True + + def save_rows_to_parquet(self, rows, vectors_directory): + if not rows: + return 0 + df = pd.DataFrame.from_records(rows) + parquet_file = os.path.join(vectors_directory, f"{self.file_ctr}.parquet") + df.to_parquet(parquet_file) + if not hasattr(self, "parquet_schema"): + self.parquet_schema = pq.read_schema(parquet_file) + else: + self.parquet_schema = pa.unify_schemas( + [self.parquet_schema, pq.read_schema(parquet_file)] + ) + self.file_structure.append(parquet_file) + self.file_ctr += 1 + return len(df) diff --git a/src/vdf_io/import_vdf/weaviate_import.py b/src/vdf_io/import_vdf/weaviate_import.py new file mode 100644 index 0000000..5ee0ac5 --- /dev/null +++ b/src/vdf_io/import_vdf/weaviate_import.py @@ -0,0 +1,177 @@ +from typing import Dict, List + +from dotenv import load_dotenv +from tqdm import tqdm + +from vdf_io.constants import DEFAULT_BATCH_SIZE, INT_MAX +from vdf_io.meta_types import NamespaceMeta +from vdf_io.names import DBNames +from vdf_io.util import cleanup_df, divide_into_batches, set_arg_from_input +from vdf_io.import_vdf.vdf_import_cls import ImportVDB +from vdf_io.weaviate_util import ( + collection_names, + compliant_collection_name, + connect_weaviate, + infer_weaviate_properties, + make_weaviate_parser, + row_to_properties, + uuid_for_id, + vector_config_for_columns, +) + + +load_dotenv() + + +class ImportWeaviate(ImportVDB): + DB_NAME_SLUG = DBNames.WEAVIATE + + @classmethod + def make_parser(cls, subparsers): + parser_weaviate = subparsers.add_parser( + cls.DB_NAME_SLUG, help="Import data to Weaviate" + ) + make_weaviate_parser(parser_weaviate) + + @classmethod + def import_vdb(cls, args): + set_arg_from_input( + args, + "batch_size", + f"Enter the batch size for importing data (default: {DEFAULT_BATCH_SIZE}): ", + int, + DEFAULT_BATCH_SIZE, + ) + weaviate_import = ImportWeaviate(args) + weaviate_import.upsert_data() + return weaviate_import + + def __init__(self, args): + super().__init__(args) + self.client = connect_weaviate(self.args) + + def upsert_data(self): + max_hit = False + self.total_imported_count = 0 + indexes_content: Dict[str, List[NamespaceMeta]] = self.vdf_meta["indexes"] + if len(indexes_content) == 0: + raise ValueError("No indexes found in VDF_META.json") + + collections = collection_names(self.client) + for index_name, index_meta in tqdm( + indexes_content.items(), desc="Importing indexes" + ): + for namespace_meta in tqdm(index_meta, desc="Importing namespaces"): + self.set_dims(namespace_meta, index_name) + data_path = namespace_meta["data_path"] + final_data_path = self.get_final_data_path(data_path) + parquet_files = self.get_parquet_files(final_data_path) + + collection_name = index_name + ( + f"_{namespace_meta['namespace']}" + if namespace_meta["namespace"] + else "" + ) + collection_name = compliant_collection_name(collection_name) + collection_name = self.create_new_name( + collection_name, collections, delimiter="_" + ) + collection_name = compliant_collection_name(collection_name) + + vector_column_names, _ = self.get_vector_column_name( + collection_name, namespace_meta, multi_vector_supported=True + ) + + if collection_name not in collections: + sample_df = self.load_sample_df(final_data_path, parquet_files) + properties = infer_weaviate_properties( + sample_df, vector_column_names, self.id_column + ) + self.client.collections.create( + collection_name, + vector_config=vector_config_for_columns( + vector_column_names, namespace_meta.get("metric") + ), + properties=properties, + ) + collections.append(collection_name) + + collection = self.client.collections.get(collection_name) + for file in tqdm(parquet_files, desc="Iterating parquet files"): + file_path = self.get_file_path(final_data_path, file) + df = self.read_parquet_progress( + file_path, + max_num_rows=( + (self.args.get("max_num_rows") or INT_MAX) + - self.total_imported_count + ), + ) + df = cleanup_df(df) + batch_size = self.args.get("batch_size") or DEFAULT_BATCH_SIZE + for batch_df in tqdm( + divide_into_batches(df, batch_size), + desc="Importing batches", + total=max(len(df) // batch_size, 1), + ): + with collection.batch.fixed_size( + batch_size=batch_size + ) as batch: + for _, row in batch_df.iterrows(): + vector = self.row_to_vector(row, vector_column_names) + if not vector: + continue + batch.add_object( + properties=row_to_properties( + row, vector_column_names, self.id_column + ), + uuid=uuid_for_id(row[self.id_column]), + vector=vector, + ) + self.total_imported_count += 1 + if self.total_imported_count >= ( + self.args.get("max_num_rows") or INT_MAX + ): + max_hit = True + break + failed_objects = collection.batch.failed_objects + if failed_objects: + raise RuntimeError( + f"Weaviate import failed for {len(failed_objects)} objects. " + f"First failure: {failed_objects[0]}" + ) + if max_hit: + break + if max_hit: + break + + tqdm.write( + f"Imported {self.total_imported_count} rows into {collection_name}" + ) + if max_hit: + break + if max_hit: + tqdm.write( + f"Max rows to be imported {self.args['max_num_rows']} hit. Exiting" + ) + break + + tqdm.write("Data import completed successfully.") + self.args["imported_count"] = self.total_imported_count + + def load_sample_df(self, final_data_path, parquet_files): + if not parquet_files: + raise ValueError("No parquet files found for Weaviate import") + first_file = self.get_file_path(final_data_path, parquet_files[0]) + return self.read_parquet_progress(first_file, max_num_rows=100) + + def row_to_vector(self, row, vector_column_names): + vectors = {} + for vector_column_name in vector_column_names: + vector_value = row.get(vector_column_name) + if vector_value is None: + continue + vectors[vector_column_name] = self.extract_vector(vector_value) + + if len(vector_column_names) == 1 and vector_column_names[0] == "vector": + return vectors.get("vector") + return vectors diff --git a/src/vdf_io/weaviate_util.py b/src/vdf_io/weaviate_util.py new file mode 100644 index 0000000..70d54cc --- /dev/null +++ b/src/vdf_io/weaviate_util.py @@ -0,0 +1,341 @@ +import argparse +import math +import re +from urllib.parse import urlparse +from uuid import UUID + +import numpy as np +import pandas as pd +import weaviate +from weaviate.classes.config import Configure, DataType, Property, VectorDistances +from weaviate.classes.init import Auth +from weaviate.util import generate_uuid5 + +from vdf_io.constants import ID_COLUMN +from vdf_io.names import DBNames +from vdf_io.util import standardize_metric_reverse + + +WEAVIATE_DISTANCE_METRICS = { + "cosine": VectorDistances.COSINE, + "l2-squared": VectorDistances.L2_SQUARED, + "dot": VectorDistances.DOT, + "manhattan": VectorDistances.MANHATTAN, +} + +VALID_PROPERTY_NAME_RE = re.compile(r"^[a-z][A-Za-z0-9_]*$") + + +def make_weaviate_parser(parser): + parser.add_argument( + "--url", + type=str, + help="Weaviate Cloud URL or custom HTTP(S) endpoint", + ) + parser.add_argument( + "--api_key", + type=str, + help="Weaviate API key. Defaults to WEAVIATE_API_KEY when set.", + ) + parser.add_argument( + "--local", + help="Connect to a local Weaviate instance", + default=False, + action=argparse.BooleanOptionalAction, + ) + parser.add_argument( + "--host", + type=str, + default="localhost", + help="Local/custom Weaviate host. Default: localhost", + ) + parser.add_argument( + "--port", + type=int, + default=8080, + help="Local/custom Weaviate HTTP port. Default: 8080", + ) + parser.add_argument( + "--grpc_host", + type=str, + help="Custom Weaviate gRPC host. Defaults to --host or --url host", + ) + parser.add_argument( + "--grpc_port", + type=int, + default=50051, + help="Local/custom Weaviate gRPC port. Default: 50051", + ) + parser.add_argument( + "--secure", + help="Use HTTPS/gRPC TLS for a custom endpoint", + default=None, + action=argparse.BooleanOptionalAction, + ) + parser.add_argument( + "--skip_init_checks", + help="Skip Weaviate client startup checks", + default=True, + action=argparse.BooleanOptionalAction, + ) + + +def connect_weaviate(args): + api_key = args.get("api_key") + if not api_key: + import os + + api_key = os.getenv("WEAVIATE_API_KEY") + auth_credentials = auth_api_key(api_key) if api_key else None + skip_init_checks = args.get("skip_init_checks", True) + url = args.get("url") + + if url: + parsed = urlparse(url if "://" in url else f"https://{url}") + cluster_url = parsed.geturl() + host = parsed.hostname or url + http_secure = ( + parsed.scheme == "https" + if args.get("secure") is None + else bool(args.get("secure")) + ) + http_port = parsed.port or (443 if http_secure else 80) + grpc_host = args.get("grpc_host") or host + grpc_port = args.get("grpc_port") or (443 if http_secure else 50051) + + if "weaviate.cloud" in host or "weaviate.network" in host: + if auth_credentials is None: + raise ValueError("Weaviate Cloud connections require --api_key") + connect_to_cloud = getattr( + weaviate, "connect_to_weaviate_cloud", None + ) or getattr(weaviate, "connect_to_wcs") + return connect_to_cloud( + cluster_url=cluster_url, + auth_credentials=auth_credentials, + skip_init_checks=skip_init_checks, + ) + + return weaviate.connect_to_custom( + http_host=host, + http_port=http_port, + http_secure=http_secure, + grpc_host=grpc_host, + grpc_port=grpc_port, + grpc_secure=http_secure, + auth_credentials=auth_credentials, + skip_init_checks=skip_init_checks, + ) + + return weaviate.connect_to_local( + host=args.get("host") or "localhost", + port=args.get("port") or 8080, + grpc_port=args.get("grpc_port") or 50051, + auth_credentials=auth_credentials, + skip_init_checks=skip_init_checks, + ) + + +def auth_api_key(api_key): + if hasattr(Auth, "api_key"): + return Auth.api_key(api_key) + return weaviate.auth.AuthApiKey(api_key) + + +def collection_names(client): + return list(client.collections.list_all().keys()) + + +def compliant_collection_name(name): + cleaned = re.sub(r"[^A-Za-z0-9_]", "_", str(name)) + cleaned = re.sub(r"_+", "_", cleaned).strip("_") + if not cleaned: + cleaned = "ImportedCollection" + if not cleaned[0].isalpha(): + cleaned = f"Collection_{cleaned}" + return cleaned[0].upper() + cleaned[1:] + + +def uuid_for_id(value): + value = str(value) + try: + return str(UUID(value)) + except ValueError: + return generate_uuid5(value) + + +def normalize_weaviate_vectors(raw_vector): + if raw_vector is None: + return {} + if isinstance(raw_vector, dict): + vectors = {} + for name, vector in raw_vector.items(): + if vector is None: + continue + column_name = "vector" if name in ("default", None) else str(name) + vectors[column_name] = normalize_vector_value(vector) + return vectors + return {"vector": normalize_vector_value(raw_vector)} + + +def normalize_vector_value(vector): + if isinstance(vector, np.ndarray): + return vector.tolist() + if hasattr(vector, "tolist"): + return vector.tolist() + return list(vector) if hasattr(vector, "__iter__") else vector + + +def first_vector_dimension(vector): + if vector is None: + return -1 + if isinstance(vector, np.ndarray): + vector = vector.tolist() + if isinstance(vector, list) and vector and isinstance(vector[0], list): + return len(vector[0]) + try: + return len(vector) + except TypeError: + return -1 + + +def get_weaviate_distance(collection_config, vector_column=None): + vector_index_config = getattr(collection_config, "vector_index_config", None) + vector_config = getattr(collection_config, "vector_config", None) or {} + if vector_column and vector_config: + config_key = "default" if vector_column == "vector" else vector_column + named_config = vector_config.get(config_key) + if named_config is not None: + vector_index_config = getattr(named_config, "vector_index_config", None) + + distance = getattr(vector_index_config, "distance_metric", None) + if hasattr(distance, "value"): + return distance.value + return distance + + +def serialize_weaviate_config(value): + if value is None: + return None + if hasattr(value, "model_dump"): + return value.model_dump() + if hasattr(value, "to_dict"): + return value.to_dict() + if isinstance(value, dict): + return {str(k): serialize_weaviate_config(v) for k, v in value.items()} + if isinstance(value, (list, tuple)): + return [serialize_weaviate_config(v) for v in value] + if hasattr(value, "value"): + return value.value + if hasattr(value, "__dict__"): + return { + k: serialize_weaviate_config(v) + for k, v in vars(value).items() + if not k.startswith("_") + } + return value + + +def vector_index_config_for_metric(metric): + metric_name = standardize_metric_reverse(metric, DBNames.WEAVIATE) + distance = WEAVIATE_DISTANCE_METRICS.get(metric_name, VectorDistances.COSINE) + return Configure.VectorIndex.hnsw(distance_metric=distance) + + +def vector_config_for_columns(vector_columns, metric): + vector_index_config = vector_index_config_for_metric(metric) + if len(vector_columns) == 1 and vector_columns[0] == "vector": + return Configure.Vectors.self_provided(vector_index_config=vector_index_config) + return [ + Configure.Vectors.self_provided( + name=column, vector_index_config=vector_index_config + ) + for column in vector_columns + ] + + +def infer_weaviate_properties(df, vector_columns, id_column=ID_COLUMN): + properties = [] + for column in df.columns: + if column == id_column or column in vector_columns: + continue + if not VALID_PROPERTY_NAME_RE.match(str(column)): + continue + data_type = infer_weaviate_data_type(df[column]) + if data_type is not None: + properties.append(Property(name=str(column), data_type=data_type)) + return properties + + +def infer_weaviate_data_type(series): + first = None + for value in series: + value = normalize_metadata_value(value) + if value is not None: + first = value + break + if first is None: + return None + if isinstance(first, bool): + return DataType.BOOL + if isinstance(first, int) and not isinstance(first, bool): + return DataType.INT + if isinstance(first, float): + return DataType.NUMBER + if isinstance(first, str): + return DataType.TEXT + if isinstance(first, list): + if all(isinstance(v, bool) for v in first): + return DataType.BOOL_ARRAY + if all(isinstance(v, int) and not isinstance(v, bool) for v in first): + return DataType.INT_ARRAY + if all(isinstance(v, (int, float)) and not isinstance(v, bool) for v in first): + return DataType.NUMBER_ARRAY + if all(isinstance(v, str) for v in first): + return DataType.TEXT_ARRAY + return None + + +def normalize_metadata_value(value): + if value is None: + return None + if isinstance(value, np.generic): + value = value.item() + if isinstance(value, float) and math.isnan(value): + return None + if isinstance(value, pd.Timestamp): + return value.isoformat() + if isinstance(value, np.ndarray): + value = value.tolist() + if isinstance(value, list): + return [normalize_metadata_value(v) for v in value] + if isinstance(value, dict): + normalized = {} + for k, v in value.items(): + normalized_value = normalize_metadata_value(v) + if normalized_value is not None: + normalized[str(k)] = normalized_value + return normalized + if isinstance(value, (str, int, float, bool)): + return value + return str(value) + + +def is_supported_property_value(value): + if isinstance(value, (str, int, float, bool)): + return True + if isinstance(value, list): + return all(isinstance(v, (str, int, float, bool)) for v in value) + return False + + +def row_to_properties(row, vector_columns, id_column=ID_COLUMN): + properties = {"vdf_id": str(row[id_column])} + for column, value in row.items(): + if column == id_column or column in vector_columns: + continue + if not VALID_PROPERTY_NAME_RE.match(str(column)): + continue + normalized = normalize_metadata_value(value) + if normalized is not None and is_supported_property_value(normalized): + properties[str(column)] = normalized + return properties