Skip to content

Commit 36978a6

Browse files
authored
Add start, stop, and temporal_policy to FeatureVector and FeatureVectorArray (#549)
1 parent 849b6e3 commit 36978a6

10 files changed

Lines changed: 163 additions & 65 deletions

File tree

apis/python/src/tiledb/vector_search/ingestion.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1661,7 +1661,12 @@ def ingest_type_erased(
16611661

16621662
ctx = vspy.Ctx(config)
16631663
data = vspy.FeatureVectorArray(
1664-
ctx, parts_array_uri, ids_array_uri, 0, to_temporal_policy(index_timestamp)
1664+
ctx,
1665+
parts_array_uri,
1666+
ids_array_uri,
1667+
0, # first_col
1668+
0, # last_col
1669+
to_temporal_policy(index_timestamp),
16651670
)
16661671
if index_type == "VAMANA":
16671672
index = vspy.IndexVamana(ctx, index_group_uri)

apis/python/src/tiledb/vector_search/type_erased_module.cc

Lines changed: 21 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -182,9 +182,22 @@ void init_type_erased_module(py::module_& m) {
182182

183183
py::class_<FeatureVector>(m, "FeatureVector", py::buffer_protocol())
184184
.def(
185-
py::init<const tiledb::Context&, const std::string&>(),
186-
py::keep_alive<1, 2>() // IndexIVFFlat should keep ctx alive.
187-
)
185+
"__init__",
186+
[](FeatureVector& instance,
187+
const tiledb::Context& ctx,
188+
const std::string& uri,
189+
size_t first_col,
190+
size_t last_col,
191+
std::optional<TemporalPolicy> temporal_policy) {
192+
new (&instance)
193+
FeatureVector(ctx, uri, first_col, last_col, temporal_policy);
194+
},
195+
py::keep_alive<1, 2>(),
196+
py::arg("ctx"),
197+
py::arg("uri"),
198+
py::arg("first_col") = 0,
199+
py::arg("last_col") = 0,
200+
py::arg("temporal_policy") = std::nullopt)
188201
.def(py::init<size_t, const std::string&>())
189202
.def(py::init<size_t, void*, const std::string&>())
190203
.def("dimensions", &FeatureVector::dimensions)
@@ -244,16 +257,18 @@ void init_type_erased_module(py::module_& m) {
244257
const tiledb::Context& ctx,
245258
const std::string& uri,
246259
const std::string& ids_uri,
247-
size_t num_vectors,
260+
size_t first_col,
261+
size_t last_col,
248262
std::optional<TemporalPolicy> temporal_policy) {
249263
new (&instance) FeatureVectorArray(
250-
ctx, uri, ids_uri, num_vectors, temporal_policy);
264+
ctx, uri, ids_uri, first_col, last_col, temporal_policy);
251265
},
252266
py::keep_alive<1, 2>(), // FeatureVectorArray should keep ctx alive.
253267
py::arg("ctx"),
254268
py::arg("uri"),
255269
py::arg("ids_uri") = "",
256-
py::arg("num_vectors") = 0,
270+
py::arg("first_col") = 0,
271+
py::arg("last_col") = 0,
257272
py::arg("temporal_policy") = std::nullopt)
258273
.def(py::init<size_t, size_t, const std::string&, const std::string&>())
259274
.def("dimensions", &FeatureVectorArray::dimensions)

src/include/api/feature_vector.h

Lines changed: 34 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -86,13 +86,20 @@ class FeatureVector {
8686
* @param ctx
8787
* @param uri
8888
*/
89-
FeatureVector(const tiledb::Context& ctx, const std::string& uri) {
90-
auto array = tiledb_helpers::open_array(tdb_func__, ctx, uri, TILEDB_READ);
89+
FeatureVector(
90+
const tiledb::Context& ctx,
91+
const std::string& uri,
92+
size_t start = 0,
93+
size_t end = 0,
94+
std::optional<TemporalPolicy> temporal_policy_input = std::nullopt) {
95+
auto temporal_policy = temporal_policy_input.value_or(TemporalPolicy{});
96+
auto array = tiledb_helpers::open_array(
97+
tdb_func__, ctx, uri, TILEDB_READ, temporal_policy);
9198

9299
feature_type_ = get_array_datatype(*array);
93100
array->close();
94101

95-
tdb_vector_from_datatype(ctx, uri);
102+
tdb_vector_from_datatype(ctx, uri, start, end, temporal_policy);
96103
}
97104

98105
/*
@@ -126,25 +133,35 @@ class FeatureVector {
126133
* Dispatch to the appropriate concrete class based on the datatype.
127134
*/
128135
void tdb_vector_from_datatype(
129-
const tiledb::Context& ctx, const std::string& uri) {
136+
const tiledb::Context& ctx,
137+
const std::string& uri,
138+
size_t start,
139+
size_t end,
140+
TemporalPolicy temporal_policy) {
130141
switch (feature_type_) {
131142
case TILEDB_FLOAT32:
132-
vector_ = std::make_unique<vector_impl<tdbVector<float>>>(ctx, uri);
143+
vector_ = std::make_unique<vector_impl<tdbVector<float>>>(
144+
ctx, uri, start, end, temporal_policy);
133145
break;
134146
case TILEDB_INT8:
135-
vector_ = std::make_unique<vector_impl<tdbVector<int8_t>>>(ctx, uri);
147+
vector_ = std::make_unique<vector_impl<tdbVector<int8_t>>>(
148+
ctx, uri, start, end, temporal_policy);
136149
break;
137150
case TILEDB_UINT8:
138-
vector_ = std::make_unique<vector_impl<tdbVector<uint8_t>>>(ctx, uri);
151+
vector_ = std::make_unique<vector_impl<tdbVector<uint8_t>>>(
152+
ctx, uri, start, end, temporal_policy);
139153
break;
140154
case TILEDB_INT32:
141-
vector_ = std::make_unique<vector_impl<tdbVector<int32_t>>>(ctx, uri);
155+
vector_ = std::make_unique<vector_impl<tdbVector<int32_t>>>(
156+
ctx, uri, start, end, temporal_policy);
142157
break;
143158
case TILEDB_UINT32:
144-
vector_ = std::make_unique<vector_impl<tdbVector<uint32_t>>>(ctx, uri);
159+
vector_ = std::make_unique<vector_impl<tdbVector<uint32_t>>>(
160+
ctx, uri, start, end, temporal_policy);
145161
break;
146162
case TILEDB_UINT64:
147-
vector_ = std::make_unique<vector_impl<tdbVector<uint64_t>>>(ctx, uri);
163+
vector_ = std::make_unique<vector_impl<tdbVector<uint64_t>>>(
164+
ctx, uri, start, end, temporal_policy);
148165
break;
149166
default:
150167
throw std::runtime_error("Unsupported attribute type");
@@ -202,8 +219,13 @@ class FeatureVector {
202219
explicit vector_impl(size_t size)
203220
: vector_(size) {
204221
}
205-
vector_impl(const tiledb::Context& ctx, const std::string& uri)
206-
: vector_(ctx, uri) {
222+
vector_impl(
223+
const tiledb::Context& ctx,
224+
const std::string& uri,
225+
size_t start,
226+
size_t end,
227+
TemporalPolicy temporal_policy)
228+
: vector_(ctx, uri, start, end, temporal_policy) {
207229
}
208230
//[[nodiscard]] void* data() override {
209231
// return _cpo::data(vector_);

0 commit comments

Comments
 (0)