Setting REDIS_KEEP_DOCUMENTS=1 as default

Author: filipecosta90

engine/clients/redis/config.py

@@ -5,7 +5,7 @@
 REDIS_USER = os.getenv("REDIS_USER", None)
 REDIS_CLUSTER = bool(int(os.getenv("REDIS_CLUSTER", 0)))
 REDIS_HYBRID_POLICY = os.getenv("REDIS_HYBRID_POLICY", None)
-REDIS_KEEP_DOCUMENTS = bool(os.getenv("REDIS_KEEP_DOCUMENTS", 0))
+REDIS_KEEP_DOCUMENTS = bool(os.getenv("REDIS_KEEP_DOCUMENTS", 1))
 GPU_STATS = bool(int(os.getenv("GPU_STATS", 0)))
 GPU_STATS_ENDPOINT = os.getenv("GPU_STATS_ENDPOINT", None)
 

engine/clients/redis/upload.py

@@ -55,48 +55,41 @@ def init_client(cls, host, distance, connection_params, upload_params):
     def upload_batch(
         cls, ids: List[int], vectors: List[list], metadata: Optional[List[dict]]
     ):
-        # if we don't delete the docs we can skip sending them again
-        # By default we always send the docs
-        if REDIS_KEEP_DOCUMENTS is False:
-            #p = cls.client.pipeline(transaction=False)
-            for i in range(len(ids)):
-                idx = ids[i]
-                vector_key = str(idx) 
-                if cls.client.exists(vector_key) is False:
-                    print(f"Setting missing key: {vector_key}")
-                    vec = vectors[i]
-                    meta = metadata[i] if metadata else {}
-                    geopoints = {}
-                    payload = {}
-                    if meta is not None:
-                        for k, v in meta.items():
-                            # This is a patch for arxiv-titles dataset where we have a list of "labels", and
-                            # we want to index all of them under the same TAG field (whose separator is ';').
-                            if k == "labels":
-                                payload[k] = ";".join(v)
-                            if (
-                                v is not None
-                                and not isinstance(v, dict)
-                                and not isinstance(v, list)
-                            ):
-                                payload[k] = v
-                        # Redis treats geopoints differently and requires putting them as
-                        # a comma-separated string with lat and lon coordinates
-                        geopoints = {
-                            k: ",".join(map(str, convert_to_redis_coords(v["lon"], v["lat"])))
-                            for k, v in meta.items()
-                            if isinstance(v, dict)
-                        }
-                    cls.client.hset(
-                        vector_key,
-                        mapping={
-                            "vector": np.array(vec).astype(cls.np_data_type).tobytes(),
-                            **payload,
-                            **geopoints,
-                        },
-                    )
-                    
-            #p.execute()
+        for i in range(len(ids)):
+            idx = ids[i]
+            vector_key = str(idx)
+            vec = vectors[i]
+            meta = metadata[i] if metadata else {}
+            geopoints = {}
+            payload = {}
+            if meta is not None:
+                for k, v in meta.items():
+                    # This is a patch for arxiv-titles dataset where we have a list of "labels", and
+                    # we want to index all of them under the same TAG field (whose separator is ';').
+                    if k == "labels":
+                        payload[k] = ";".join(v)
+                    if (
+                        v is not None
+                        and not isinstance(v, dict)
+                        and not isinstance(v, list)
+                    ):
+                        payload[k] = v
+                # Redis treats geopoints differently and requires putting them as
+                # a comma-separated string with lat and lon coordinates
+                geopoints = {
+                    k: ",".join(map(str, convert_to_redis_coords(v["lon"], v["lat"])))
+                    for k, v in meta.items()
+                    if isinstance(v, dict)
+                }
+            cls.client.hset(
+                vector_key,
+                mapping={
+                    "vector": np.array(vec).astype(cls.np_data_type).tobytes(),
+                    **payload,
+                    **geopoints,
+                },
+            )
+                
 
     @classmethod
     def post_upload(cls, _distance):