feat(flat): accept contiguous external vector batches

src/core/algorithm/flat/flat_streamer.cc

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include "flat_streamer.h"
+#include <limits>
 #include <zvec/core/framework/index_factory.h>
 #include "flat_streamer_context.h"
 #include "flat_streamer_dumper.h"
@@ -517,27 +518,61 @@ int FlatStreamer<BATCH_SIZE>::search_external(
     auto *heap = ctx->result_heap();
     heap->clear();
     // Destroy the outstanding lease before its cursor, including on errors.
-    VectorSource::Iterator::Pointer cursor;
-    VectorSource::Batch batch;
+    VectorIterator::Pointer cursor;
+    VectorBatch batch;
     if (!keys) cursor = source->create_iterator();
+    auto clear_batch = [&]() {
+      if (auto *pointers = std::get_if<PointerBatch>(&batch)) {
+        pointers->ids.clear();
+        pointers->vectors.clear();
+        pointers->lease.reset();
+      } else {
+        std::get<ContiguousBatch>(batch) = {};
+      }
+    };
     auto consume = [&]() -> int {
-      if (batch.ids.size() != batch.vectors.size() ||
-          batch.ids.size() > batch_limit)
+      const auto *contiguous = std::get_if<ContiguousBatch>(&batch);
+      const auto *pointers = std::get_if<PointerBatch>(&batch);
+      if (!contiguous && !pointers) return IndexError_ReadData;
+      const size_t batch_size =
+          contiguous ? contiguous->count : pointers->ids.size();
+      if (batch_size > batch_limit) return IndexError_ReadData;
+      if (contiguous) {
+        if (batch_size &&
+            (!contiguous->data || contiguous->stride < meta_.element_size() ||
+             batch_size - 1 >
+                 (std::numeric_limits<size_t>::max() - meta_.element_size()) /
+                     contiguous->stride ||
+             contiguous->first_id >= kInvalidNodeId ||
+             batch_size > kInvalidNodeId - contiguous->first_id))
+          return IndexError_ReadData;
+      } else if (pointers->ids.size() != pointers->vectors.size()) {
         return IndexError_ReadData;
+      }
+      const auto *ids = pointers ? pointers->ids.data() : nullptr;
+      const uint32_t first_id = contiguous ? contiguous->first_id : 0;
       scratch->vector_ptrs.clear();
       scratch->vector_keys.clear();
-      if (cursor) entity_->key_presence(batch.ids, &scratch->key_presence);
-      for (size_t i = 0; i < batch.ids.size(); ++i) {
-        const auto id = batch.ids[i];
+      scratch->vector_ptrs.reserve(batch_size);
+      scratch->vector_keys.reserve(batch_size);
+      if (cursor)
+        entity_->key_presence(ids, batch_size, first_id,
+                              &scratch->key_presence);
+      for (size_t i = 0; i < batch_size; ++i) {
+        const auto id = ids ? ids[i] : first_id + static_cast<uint32_t>(i);
         if (id == kInvalidNodeId) return IndexError_ReadData;
         if (cursor && (!scratch->key_presence[i] ||
                        (ctx->filter().is_valid() && ctx->filter()(id)))) {
           ++*ctx->mutable_stats(q)->mutable_filtered_count();
           continue;
         }
         // Invisible or unregistered rows need not have a readable vector.
-        if (!batch.vectors[i]) return IndexError_ReadData;
-        scratch->vector_ptrs.push_back(batch.vectors[i]);
+        const void *vector = contiguous
+                                 ? static_cast<const char *>(contiguous->data) +
+                                       i * contiguous->stride
+                                 : pointers->vectors[i];
+        if (!vector) return IndexError_ReadData;
+        scratch->vector_ptrs.push_back(vector);
         scratch->vector_keys.push_back(id);
       }
       const size_t size = scratch->vector_ptrs.size();
@@ -568,22 +603,26 @@ int FlatStreamer<BATCH_SIZE>::search_external(
 
     if (cursor) {
       while (true) {
-        batch.clear();
+        clear_batch();
         int ret = cursor->next_batch(batch_limit, &batch);
         if (ret != 0) return ret;
         ret = consume();
         if (ret != 0) return ret;
-        if (batch.ids.empty()) break;
+        const auto *contiguous = std::get_if<ContiguousBatch>(&batch);
+        if (contiguous ? contiguous->count == 0
+                       : std::get<PointerBatch>(batch).ids.empty())
+          break;
       }
     } else {
       std::vector<uint32_t> ids;
       ids.reserve(batch_limit);
       auto read_batch = [&]() -> int {
         if (ids.empty()) return 0;
-        batch.clear();
-        batch.ids = ids;
-        batch.vectors.resize(ids.size());
-        source->get_vectors(ids.data(), ids.size(), batch.vectors.data());
+        clear_batch();
+        auto &pointers = std::get<PointerBatch>(batch);
+        pointers.ids = ids;
+        pointers.vectors.resize(ids.size());
+        source->get_vectors(ids.data(), ids.size(), pointers.vectors.data());
         int ret = consume();
         ids.clear();
         return ret;
@@ -622,7 +661,7 @@ int FlatStreamer<BATCH_SIZE>::search_external(
         }
       }
     }
-    batch.clear();
+    clear_batch();
     ctx->take_topk_result(q);
     if (ctx->fetch_vector()) {
       // Match HNSW's borrowed random-access result contract. Scan pages have

src/core/algorithm/flat/flat_streamer_entity.h

@@ -180,13 +180,15 @@ class FlatStreamerEntity {
     return key_info_map_.find(key) != key_info_map_.end();
   }
 
-  void key_presence(const std::vector<uint32_t> &keys,
+  // A null keys pointer denotes consecutive IDs starting at first_key.
+  void key_presence(const uint32_t *keys, size_t count, uint32_t first_key,
                     std::vector<uint8_t> *present) const {
-    present->resize(keys.size());
+    present->resize(count);
     ailego::ReadLock lock(*key_info_map_lock_);
     std::lock_guard<ailego::ReadLock> guard(lock);
-    for (size_t i = 0; i < keys.size(); ++i) {
-      (*present)[i] = key_info_map_.find(keys[i]) != key_info_map_.end();
+    for (size_t i = 0; i < count; ++i) {
+      const auto key = keys ? keys[i] : first_key + i;
+      (*present)[i] = key_info_map_.find(key) != key_info_map_.end();
     }
   }
 

src/core/interface/vector_source.cc

@@ -12,16 +12,65 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <algorithm>
+#include <limits>
+#include <utility>
+#include <zvec/core/framework/index_error.h>
 #include <zvec/core/interface/vector_source.h>
 
 namespace zvec {
 namespace core {
 
+namespace {
+
+bool ValidContiguousView(const ContiguousBatch &view) {
+  if (!view.count) return true;
+  constexpr auto kInvalidId = std::numeric_limits<uint32_t>::max();
+  if (!view.data || !view.stride || view.first_id >= kInvalidId ||
+      view.count > kInvalidId - view.first_id)
+    return false;
+  const auto max_span = std::numeric_limits<uintptr_t>::max() -
+                        reinterpret_cast<uintptr_t>(view.data);
+  return view.count <= max_span / view.stride;
+}
+
+const void *GetContiguousVector(const ContiguousBatch &view, uint32_t id) {
+  if (id < view.first_id || id - view.first_id >= view.count) return nullptr;
+  return static_cast<const char *>(view.data) +
+         size_t(id - view.first_id) * view.stride;
+}
+
+class ContiguousVectorIterator final : public VectorIterator {
+ public:
+  explicit ContiguousVectorIterator(ContiguousBatch view)
+      : view_(std::move(view)), valid_(ValidContiguousView(view_)) {}
+
+  int next_batch(uint32_t max_count, VectorBatch *out) override {
+    if (!out || !max_count) return IndexError_InvalidArgument;
+    auto &batch = out->emplace<ContiguousBatch>();
+    if (!valid_) return IndexError_InvalidArgument;
+    if (offset_ == view_.count) return 0;
+    const auto count = std::min(max_count, view_.count - offset_);
+    batch = {
+        static_cast<const char *>(view_.data) + size_t(offset_) * view_.stride,
+        view_.first_id + offset_, count, view_.stride, view_.lease};
+    offset_ += count;
+    return 0;
+  }
+
+ private:
+  const ContiguousBatch view_;
+  const bool valid_;
+  uint32_t offset_{0};
+};
+
+}  // namespace
+
 VectorSource::VectorSource() = default;
 
 VectorSource::~VectorSource() = default;
 
-VectorSource::Iterator::Pointer VectorSource::create_iterator() const {
+VectorIterator::Pointer VectorSource::create_iterator() const {
   return nullptr;
 }
 
@@ -32,5 +81,24 @@ void VectorSource::get_vectors(const uint32_t *ids, uint32_t count,
   }
 }
 
+const void *ContiguousVectorSource::get_vector(uint32_t id) const {
+  const auto view = get_contiguous_view();
+  return ValidContiguousView(view) ? GetContiguousVector(view, id) : nullptr;
+}
+
+void ContiguousVectorSource::get_vectors(const uint32_t *ids, uint32_t count,
+                                         const void **out) const {
+  if (!count) return;
+  const auto view = get_contiguous_view();
+  const bool valid = ValidContiguousView(view);
+  for (uint32_t i = 0; i < count; ++i) {
+    out[i] = valid ? GetContiguousVector(view, ids[i]) : nullptr;
+  }
+}
+
+VectorIterator::Pointer ContiguousVectorSource::create_iterator() const {
+  return std::make_unique<ContiguousVectorIterator>(get_contiguous_view());
+}
+
 }  // namespace core
 }  // namespace zvec

src/include/zvec/core/interface/vector_source.h

@@ -14,43 +14,46 @@
 
 #pragma once
 
+#include <cstddef>
 #include <cstdint>
 #include <memory>
+#include <variant>
 #include <vector>
 #include <zvec/export.h>
 
 namespace zvec {
 namespace core {
 
-class ZVEC_CORE_API VectorSource {
+// Borrowed raw FP32 rows in get_vector()'s ID space. Leases may pin storage.
+struct PointerBatch {
+  std::vector<uint32_t> ids;
+  // Only filtered or unregistered rows may have null pointers.
+  std::vector<const void *> vectors;
+  std::shared_ptr<const void> lease;
+};
+
+// Row i has ID first_id + i; stride is in bytes, at least one vector's size.
+struct ContiguousBatch {
+  const void *data{nullptr};
+  uint32_t first_id{0};
+  uint32_t count{0};  // Includes filtered rows.
+  size_t stride{0};
+  std::shared_ptr<const void> lease;
+};
+
+using VectorBatch = std::variant<PointerBatch, ContiguousBatch>;
+
+class ZVEC_CORE_API VectorIterator {
  public:
-  // Logical row-major vectors in the same ID namespace as get_vector().
-  // Pointers need not be contiguous. External Flat currently expects raw FP32.
-  struct Batch {
-    std::vector<uint32_t> ids;
-    // Rows excluded by Flat's registered-ID set or filter may be null.
-    // Every row selected for distance computation must have a valid pointer.
-    std::vector<const void *> vectors;
-    // Optional backing allocation or page pin. Data must remain valid until
-    // the batch is released, or (for scans) the next next_batch() call.
-    std::shared_ptr<const void> lease;
-
-    void clear() {
-      vectors.clear();
-      ids.clear();
-      lease.reset();
-    }
-  };
-
-  class Iterator {
-   public:
-    using Pointer = std::unique_ptr<Iterator>;
-    virtual ~Iterator() = default;
-    // Return 0 with 1..max_count rows, 0 with an empty batch at EOF, or an
-    // IndexError on failure. Enumerate every visible ID exactly once.
-    virtual int next_batch(uint32_t max_count, Batch *out) = 0;
-  };
+  using Pointer = std::unique_ptr<VectorIterator>;
+  virtual ~VectorIterator() = default;
+  // Return 0 with 1..max_count rows, 0 with an empty batch at EOF, or an
+  // IndexError on failure. Enumerate every visible ID exactly once.
+  virtual int next_batch(uint32_t max_count, VectorBatch *out) = 0;
+};
 
+class ZVEC_CORE_API VectorSource {
+ public:
   VectorSource();
   virtual ~VectorSource();
 
@@ -66,7 +69,21 @@ class ZVEC_CORE_API VectorSource {
   // The caller binds the source, scan domain and visibility filter to the same
   // stable snapshot for the whole request, including subsequent random reads.
   // A disk source can scan physical blocks behind this interface.
-  virtual Iterator::Pointer create_iterator() const;
+  virtual VectorIterator::Pointer create_iterator() const;
+};
+
+class ZVEC_CORE_API ContiguousVectorSource : public VectorSource {
+ public:
+  // Iterators retain this view and its lease. The caller keeps storage alive
+  // for random reads and any borrowed search results.
+  virtual ContiguousBatch get_contiguous_view() const = 0;
+
+  const void *get_vector(uint32_t id) const override;
+
+  void get_vectors(const uint32_t *ids, uint32_t count,
+                   const void **out) const override;
+
+  VectorIterator::Pointer create_iterator() const override;
 };
 
 }  // namespace core