Add Weaviate import and export support

README.md

@@ -44,6 +44,7 @@ See the [Contributing](#contributing) section to add support for your favorite v
 | DataStax Astra DB              | ✅     | ✅     |
 | Chroma                         | ✅     | ✅     |
 | Turbopuffer                    | ✅     | ✅     |
+| Weaviate                       | ✅     | ✅     |
 
 </details>
 
@@ -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

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)