Fix test

TileDB-Inc · Jul 19, 2024 · f9bbdfb · f9bbdfb
1 parent 2d31da6
commit f9bbdfb
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 13 deletions.
diff --git a/apis/python/src/tiledb/vector_search/ingestion.py b/apis/python/src/tiledb/vector_search/ingestion.py
@@ -1713,21 +1713,19 @@ def ingest_vectors_udf(
             )
             index_array_rows_dim = tiledb.Dim(
                 name="rows",
-                domain=(0, partitions),
-                tile=partitions,
+                domain=(0, MAX_INT32),
+                tile=100000,
                 dtype=np.dtype(np.int32),
             )
             index_array_dom = tiledb.Domain(index_array_rows_dim)
             index_attr = tiledb.Attr(
                 name="values",
                 dtype=np.dtype(np.uint64),
-                filters=tiledb.FilterList([tiledb.ZstdFilter()]),
             )
             index_schema = tiledb.ArraySchema(
                 domain=index_array_dom,
                 sparse=False,
                 attrs=[index_attr],
-                capacity=partitions,
             )
             tiledb.Array.create(partial_write_array_index_tmp_uri, index_schema)
             partition_start = part_id * (partitions + 1)
@@ -1753,7 +1751,7 @@ def ingest_vectors_udf(
                     id_uri=partial_write_array_ids_uri,
                     start=part,
                     end=part_end,
-                    partition_start=part_id * (partitions + 1),
+                    # partition_start=part_id * (partitions + 1),
                     nthreads=threads,
                     **(
                         {"timestamp": index_timestamp}
@@ -1796,7 +1794,7 @@ def ingest_vectors_udf(
                     id_uri=partial_write_array_ids_uri,
                     start=part,
                     end=part_end,
-                    partition_start=part_id * (partitions + 1),
+                    # partition_start=part_id * (partitions + 1),
                     nthreads=threads,
                     **(
                         {"timestamp": index_timestamp}
@@ -1853,21 +1851,19 @@ def ingest_additions_udf(
         )
         index_array_rows_dim = tiledb.Dim(
             name="rows",
-            domain=(0, partitions),
-            tile=partitions,
+            domain=(0, MAX_INT32),
+            tile=100000,
             dtype=np.dtype(np.int32),
         )
         index_array_dom = tiledb.Domain(index_array_rows_dim)
         index_attr = tiledb.Attr(
             name="values",
             dtype=np.dtype(np.uint64),
-            filters=tiledb.FilterList([tiledb.ZstdFilter()]),
         )
         index_schema = tiledb.ArraySchema(
             domain=index_array_dom,
             sparse=False,
             attrs=[index_attr],
-            capacity=partitions,
         )
         tiledb.Array.create(partial_write_array_index_tmp_uri, index_schema)
 
@@ -1889,10 +1885,11 @@ def ingest_additions_udf(
             centroids_uri=centroids_uri,
             parts_uri=partial_write_array_parts_uri,
             index_array_uri=partial_write_array_index_tmp_uri,
+            # index_array_uri=partial_write_array_index_uri,
             id_uri=partial_write_array_ids_uri,
             start=write_offset,
             end=0,
-            partition_start=partition_start,
+            # partition_start=partition_start,
             nthreads=threads,
             **({"timestamp": index_timestamp} if index_timestamp is not None else {}),
             config=config,

diff --git a/apis/python/src/tiledb/vector_search/module.cc b/apis/python/src/tiledb/vector_search/module.cc
@@ -327,7 +327,8 @@ static void declare_ivf_index(py::module& m, const std::string& suffix) {
             start_pos,
             end_pos,
             nthreads,
-            timestamp);
+            timestamp,
+            partition_start);
       },
       py::keep_alive<1, 2>());
 }
@@ -363,7 +364,8 @@ static void declare_ivf_index_tdb(py::module& m, const std::string& suffix) {
             start_pos,
             end_pos,
             nthreads,
-            timestamp);
+            timestamp,
+            partition_start);
       },
       py::keep_alive<1, 2>());
 }