Merge pull request ClickHouse#54103 from ClickHouse/ustweaks

Small usearch index improvements: metrics and configurable internal data type

docs/en/engines/table-engines/mergetree-family/annindexes.md

@@ -252,7 +252,7 @@ CREATE TABLE table_with_usearch_index
 (
  id Int64,
  vectors Array(Float32),
- INDEX [ann_index_name] vectors TYPE usearch([Distance]) [GRANULARITY N]
+ INDEX [ann_index_name] vectors TYPE usearch([Distance[, ScalarKind]]) [GRANULARITY N]
 )
 ENGINE = MergeTree
 ORDER BY id;
@@ -265,7 +265,7 @@ CREATE TABLE table_with_usearch_index
 (
  id Int64,
  vectors Tuple(Float32[, Float32[, ...]]),
- INDEX [ann_index_name] vectors TYPE usearch([Distance]) [GRANULARITY N]
+ INDEX [ann_index_name] vectors TYPE usearch([Distance[, ScalarKind]]) [GRANULARITY N]
 )
 ENGINE = MergeTree
 ORDER BY id;
@@ -277,5 +277,8 @@ USearch currently supports two distance functions:
 - `cosineDistance`, also called cosine similarity, is the cosine of the angle between two (non-zero) vectors
  ([Wikipedia](https://en.wikipedia.org/wiki/Cosine_similarity)).
 
+USearch allows storing the vectors in reduced precision formats. Supported scalar kinds are `f64`, `f32`, `f16` or `i8`. If no scalar kind
+was specified during index creation, `f16` is used as default.
+
 For normalized data, `L2Distance` is usually a better choice, otherwise `cosineDistance` is recommended to compensate for scale. If no
-distance function was specified during index creation, `L2Distance` is used as default.
+distance function was specified during index creation, `L2Distance` is used as default.

src/Common/ProfileEvents.cpp

@@ -253,6 +253,13 @@ The server successfully detected this situation and will download merged part fr
  M(PolygonsAddedToPool, "A polygon has been added to the cache (pool) for the 'pointInPolygon' function.") \
  M(PolygonsInPoolAllocatedBytes, "The number of bytes for polygons added to the cache (pool) for the 'pointInPolygon' function.") \
  \
+ M(USearchAddCount, "Number of vectors added to usearch indexes.") \
+ M(USearchAddVisitedMembers, "Number of nodes visited when adding vectors to usearch indexes.") \
+ M(USearchAddComputedDistances, "Number of times distance was computed when adding vectors to usearch indexes.") \
+ M(USearchSearchCount, "Number of search operations performed in usearch indexes.") \
+ M(USearchSearchVisitedMembers, "Number of nodes visited when searching in usearch indexes.") \
+ M(USearchSearchComputedDistances, "Number of times distance was computed when searching usearch indexes.") \
+ \
  M(RWLockAcquiredReadLocks, "Number of times a read lock was acquired (in a heavy RWLock).") \
  M(RWLockAcquiredWriteLocks, "Number of times a write lock was acquired (in a heavy RWLock).") \
  M(RWLockReadersWaitMilliseconds, "Total time spent waiting for a read lock to be acquired (in a heavy RWLock).") \

src/Storages/MergeTree/MergeTreeIndexAnnoy.cpp

@@ -143,6 +143,9 @@ void MergeTreeIndexAggregatorAnnoy<Distance>::update(const Block & block, size_t
  if (rows_read == 0)
  return;
 
+ if (rows_read > std::numeric_limits<uint32_t>::max())
+ throw Exception(ErrorCodes::INCORRECT_DATA, "Index granularity is too big: more than 4B rows per index granule.");
+
  if (index_sample_block.columns() > 1)
  throw Exception(ErrorCodes::LOGICAL_ERROR, "Expected block with single column");
 

src/Storages/MergeTree/MergeTreeIndexUSearch.cpp

@@ -15,6 +15,16 @@
 #include <Interpreters/Context.h>
 #include <Interpreters/castColumn.h>
 
+namespace ProfileEvents
+{
+ extern const Event USearchAddCount;
+ extern const Event USearchAddVisitedMembers;
+ extern const Event USearchAddComputedDistances;
+ extern const Event USearchSearchCount;
+ extern const Event USearchSearchVisitedMembers;
+ extern const Event USearchSearchComputedDistances;
+}
+
 namespace DB
 {
 
@@ -28,9 +38,20 @@ namespace ErrorCodes
  extern const int LOGICAL_ERROR;
 }
 
+namespace
+{
+
+std::unordered_map<String, unum::usearch::scalar_kind_t> nameToScalarKind = {
+ {"f64", unum::usearch::scalar_kind_t::f64_k},
+ {"f32", unum::usearch::scalar_kind_t::f32_k},
+ {"f16", unum::usearch::scalar_kind_t::f16_k},
+ {"i8", unum::usearch::scalar_kind_t::i8_k}};
+
+}
+
 template <unum::usearch::metric_kind_t Metric>
-USearchIndexWithSerialization<Metric>::USearchIndexWithSerialization(size_t dimensions)
- : Base(Base::make(unum::usearch::metric_punned_t(dimensions, Metric)))
+USearchIndexWithSerialization<Metric>::USearchIndexWithSerialization(size_t dimensions, unum::usearch::scalar_kind_t scalar_kind)
+ : Base(Base::make(unum::usearch::metric_punned_t(dimensions, Metric, scalar_kind)))
 {
 }
 
@@ -67,9 +88,11 @@ size_t USearchIndexWithSerialization<Metric>::getDimensions() const
 template <unum::usearch::metric_kind_t Metric>
 MergeTreeIndexGranuleUSearch<Metric>::MergeTreeIndexGranuleUSearch(
  const String & index_name_,
- const Block & index_sample_block_)
+ const Block & index_sample_block_,
+ unum::usearch::scalar_kind_t scalar_kind_)
  : index_name(index_name_)
  , index_sample_block(index_sample_block_)
+ , scalar_kind(scalar_kind_)
  , index(nullptr)
 {
 }
@@ -78,9 +101,11 @@ template <unum::usearch::metric_kind_t Metric>
 MergeTreeIndexGranuleUSearch<Metric>::MergeTreeIndexGranuleUSearch(
  const String & index_name_,
  const Block & index_sample_block_,
+ unum::usearch::scalar_kind_t scalar_kind_,
  USearchIndexWithSerializationPtr<Metric> index_)
  : index_name(index_name_)
  , index_sample_block(index_sample_block_)
+ , scalar_kind(scalar_kind_)
  , index(std::move(index_))
 {
 }
@@ -99,23 +124,25 @@ void MergeTreeIndexGranuleUSearch<Metric>::deserializeBinary(ReadBuffer & istr,
 {
  UInt64 dimension;
  readIntBinary(dimension, istr);
- index = std::make_shared<USearchIndexWithSerialization<Metric>>(dimension);
+ index = std::make_shared<USearchIndexWithSerialization<Metric>>(dimension, scalar_kind);
  index->deserialize(istr);
 }
 
 template <unum::usearch::metric_kind_t Metric>
 MergeTreeIndexAggregatorUSearch<Metric>::MergeTreeIndexAggregatorUSearch(
  const String & index_name_,
- const Block & index_sample_block_)
+ const Block & index_sample_block_,
+ unum::usearch::scalar_kind_t scalar_kind_)
  : index_name(index_name_)
  , index_sample_block(index_sample_block_)
+ , scalar_kind(scalar_kind_)
 {
 }
 
 template <unum::usearch::metric_kind_t Metric>
 MergeTreeIndexGranulePtr MergeTreeIndexAggregatorUSearch<Metric>::getGranuleAndReset()
 {
- auto granule = std::make_shared<MergeTreeIndexGranuleUSearch<Metric>>(index_name, index_sample_block, index);
+ auto granule = std::make_shared<MergeTreeIndexGranuleUSearch<Metric>>(index_name, index_sample_block, scalar_kind, index);
  index = nullptr;
  return granule;
 }
@@ -131,9 +158,13 @@ void MergeTreeIndexAggregatorUSearch<Metric>::update(const Block & block, size_t
  block.rows());
 
  size_t rows_read = std::min(limit, block.rows() - *pos);
+
  if (rows_read == 0)
  return;
 
+ if (rows_read > std::numeric_limits<uint32_t>::max())
+ throw Exception(ErrorCodes::INCORRECT_DATA, "Index granularity is too big: more than 4B rows per index granule.");
+
  if (index_sample_block.columns() > 1)
  throw Exception(ErrorCodes::LOGICAL_ERROR, "Expected block with single column");
 
@@ -151,26 +182,29 @@ void MergeTreeIndexAggregatorUSearch<Metric>::update(const Block & block, size_t
  const auto & offsets = column_array->getOffsets();
  const size_t num_rows = offsets.size();
 
-
  /// Check all sizes are the same
  size_t size = offsets[0];
  for (size_t i = 0; i < num_rows - 1; ++i)
  if (offsets[i + 1] - offsets[i] != size)
  throw Exception(ErrorCodes::INCORRECT_DATA, "All arrays in column {} must have equal length", index_column_name);
 
  if (!index)
- index = std::make_shared<USearchIndexWithSerialization<Metric>>(size);
+ index = std::make_shared<USearchIndexWithSerialization<Metric>>(size, scalar_kind);
 
  /// Add all rows of block
  if (!index->reserve(unum::usearch::ceil2(index->size() + num_rows)))
  throw Exception(ErrorCodes::CANNOT_ALLOCATE_MEMORY, "Could not reserve memory for usearch index");
 
- if (auto rc = index->add(index->size(), array.data()); !rc)
-  throw Exception(ErrorCodes::INCORRECT_DATA, rc.error.release());
- for (size_t current_row = 1; current_row < num_rows; ++current_row)
- if (auto rc = index->add(index->size(), &array[offsets[current_row - 1]]); !rc)
+ for (size_t current_row = 0; current_row < num_rows; ++current_row)
+ {
+  auto rc = index->add(static_cast<uint32_t>(index->size()), &array[offsets[current_row - 1]]);
+ if (!rc)
  throw Exception(ErrorCodes::INCORRECT_DATA, rc.error.release());
 
+ ProfileEvents::increment(ProfileEvents::USearchAddCount);
+ ProfileEvents::increment(ProfileEvents::USearchAddVisitedMembers, rc.visited_members);
+ ProfileEvents::increment(ProfileEvents::USearchAddComputedDistances, rc.computed_distances);
+ }
  }
  else if (const auto & column_tuple = typeid_cast<const ColumnTuple *>(column_cut.get()))
  {
@@ -187,14 +221,21 @@ void MergeTreeIndexAggregatorUSearch<Metric>::update(const Block & block, size_t
  throw Exception(ErrorCodes::LOGICAL_ERROR, "Tuple has 0 rows, {} rows expected", rows_read);
 
  if (!index)
- index = std::make_shared<USearchIndexWithSerialization<Metric>>(data[0].size());
+ index = std::make_shared<USearchIndexWithSerialization<Metric>>(data[0].size(), scalar_kind);
 
  if (!index->reserve(unum::usearch::ceil2(index->size() + data.size())))
  throw Exception(ErrorCodes::CANNOT_ALLOCATE_MEMORY, "Could not reserve memory for usearch index");
 
  for (const auto & item : data)
- if (auto rc = index->add(index->size(), item.data()); !rc)
+ {
+ auto rc = index->add(static_cast<uint32_t>(index->size()), item.data());
+ if (!rc)
  throw Exception(ErrorCodes::INCORRECT_DATA, rc.error.release());
+
+ ProfileEvents::increment(ProfileEvents::USearchAddCount);
+ ProfileEvents::increment(ProfileEvents::USearchAddVisitedMembers, rc.visited_members);
+ ProfileEvents::increment(ProfileEvents::USearchAddComputedDistances, rc.computed_distances);
+ }
  }
  else
  throw Exception(ErrorCodes::LOGICAL_ERROR, "Expected Array or Tuple column");
@@ -257,7 +298,12 @@ std::vector<size_t> MergeTreeIndexConditionUSearch::getUsefulRangesImpl(MergeTre
  ann_condition.getDimensions(), index->dimensions());
 
  auto result = index->search(reference_vector.data(), limit);
- std::vector<UInt64> neighbors(result.size()); /// indexes of dots which were closest to the reference vector
+
+ ProfileEvents::increment(ProfileEvents::USearchSearchCount);
+ ProfileEvents::increment(ProfileEvents::USearchSearchVisitedMembers, result.visited_members);
+ ProfileEvents::increment(ProfileEvents::USearchSearchComputedDistances, result.computed_distances);
+
+ std::vector<UInt32> neighbors(result.size()); /// indexes of dots which were closest to the reference vector
  std::vector<Float32> distances(result.size());
  result.dump_to(neighbors.data(), distances.data());
 
@@ -277,27 +323,28 @@ std::vector<size_t> MergeTreeIndexConditionUSearch::getUsefulRangesImpl(MergeTre
  return granule_numbers;
 }
 
-MergeTreeIndexUSearch::MergeTreeIndexUSearch(const IndexDescription & index_, const String & distance_function_)
+MergeTreeIndexUSearch::MergeTreeIndexUSearch(const IndexDescription & index_, const String & distance_function_, unum::usearch::scalar_kind_t scalar_kind_)
  : IMergeTreeIndex(index_)
  , distance_function(distance_function_)
+ , scalar_kind(scalar_kind_)
 {
 }
 
 MergeTreeIndexGranulePtr MergeTreeIndexUSearch::createIndexGranule() const
 {
  if (distance_function == DISTANCE_FUNCTION_L2)
- return std::make_shared<MergeTreeIndexGranuleUSearch<unum::usearch::metric_kind_t::l2sq_k>>(index.name, index.sample_block);
+ return std::make_shared<MergeTreeIndexGranuleUSearch<unum::usearch::metric_kind_t::l2sq_k>>(index.name, index.sample_block, scalar_kind);
  else if (distance_function == DISTANCE_FUNCTION_COSINE)
- return std::make_shared<MergeTreeIndexGranuleUSearch<unum::usearch::metric_kind_t::cos_k>>(index.name, index.sample_block);
+ return std::make_shared<MergeTreeIndexGranuleUSearch<unum::usearch::metric_kind_t::cos_k>>(index.name, index.sample_block, scalar_kind);
  std::unreachable();
 }
 
 MergeTreeIndexAggregatorPtr MergeTreeIndexUSearch::createIndexAggregator() const
 {
  if (distance_function == DISTANCE_FUNCTION_L2)
- return std::make_shared<MergeTreeIndexAggregatorUSearch<unum::usearch::metric_kind_t::l2sq_k>>(index.name, index.sample_block);
+ return std::make_shared<MergeTreeIndexAggregatorUSearch<unum::usearch::metric_kind_t::l2sq_k>>(index.name, index.sample_block, scalar_kind);
  else if (distance_function == DISTANCE_FUNCTION_COSINE)
- return std::make_shared<MergeTreeIndexAggregatorUSearch<unum::usearch::metric_kind_t::cos_k>>(index.name, index.sample_block);
+ return std::make_shared<MergeTreeIndexAggregatorUSearch<unum::usearch::metric_kind_t::cos_k>>(index.name, index.sample_block, scalar_kind);
  std::unreachable();
 }
 
@@ -313,18 +360,25 @@ MergeTreeIndexPtr usearchIndexCreator(const IndexDescription & index)
  if (!index.arguments.empty())
  distance_function = index.arguments[0].get<String>();
 
- return std::make_shared<MergeTreeIndexUSearch>(index, distance_function);
+ static constexpr auto default_scalar_kind = unum::usearch::scalar_kind_t::f16_k;
+ auto scalar_kind = default_scalar_kind;
+ if (index.arguments.size() > 1)
+ scalar_kind = nameToScalarKind.at(index.arguments[1].get<String>());
+
+ return std::make_shared<MergeTreeIndexUSearch>(index, distance_function, scalar_kind);
 }
 
 void usearchIndexValidator(const IndexDescription & index, bool /* attach */)
 {
  /// Check number and type of USearch index arguments:
 
- if (index.arguments.size() > 1)
+ if (index.arguments.size() > 2)
  throw Exception(ErrorCodes::INCORRECT_QUERY, "USearch index must not have more than one parameters");
 
  if (!index.arguments.empty() && index.arguments[0].getType() != Field::Types::String)
- throw Exception(ErrorCodes::INCORRECT_QUERY, "Distance function argument of USearch index must be of type String");
+ throw Exception(ErrorCodes::INCORRECT_QUERY, "First argument of USearch index (distance function) must be of type String");
+ if (index.arguments.size() > 1 && index.arguments[1].getType() != Field::Types::String)
+ throw Exception(ErrorCodes::INCORRECT_QUERY, "Second argument of USearch index (scalar type) must be of type String");
 
  /// Check that the index is created on a single column
 
@@ -340,6 +394,20 @@ void usearchIndexValidator(const IndexDescription & index, bool /* attach */)
  throw Exception(ErrorCodes::INCORRECT_DATA, "USearch index only supports distance functions '{}' and '{}'", DISTANCE_FUNCTION_L2, DISTANCE_FUNCTION_COSINE);
  }
 
+ /// Check that a supported kind was passed as a second argument
+
+ if (index.arguments.size() > 1 && !nameToScalarKind.contains(index.arguments[1].get<String>()))
+ {
+ String supported_kinds;
+ for (const auto & [name, kind] : nameToScalarKind)
+ {
+ if (!supported_kinds.empty())
+ supported_kinds += ", ";
+ supported_kinds += name;
+ }
+ throw Exception(ErrorCodes::INCORRECT_DATA, "Unrecognized scalar kind (second argument) for USearch index. Supported kinds are: {}", supported_kinds);
+ }
+
  /// Check data type of indexed column:
 
  auto throw_unsupported_underlying_column_exception = []()

src/Storages/MergeTree/MergeTreeIndexUSearch.h

@@ -12,13 +12,15 @@
 namespace DB
 {
 
+using USearchImplType = unum::usearch::index_dense_gt</* key_at */ uint32_t, /* compressed_slot_at */ uint32_t>;
+
 template <unum::usearch::metric_kind_t Metric>
-class USearchIndexWithSerialization : public unum::usearch::index_dense_t
+class USearchIndexWithSerialization : public USearchImplType
 {
- using Base = unum::usearch::index_dense_t;
+ using Base = USearchImplType;
 
 public:
- explicit USearchIndexWithSerialization(size_t dimensions);
+ USearchIndexWithSerialization(size_t dimensions, unum::usearch::scalar_kind_t scalar_kind);
  void serialize(WriteBuffer & ostr) const;
  void deserialize(ReadBuffer & istr);
  size_t getDimensions() const;
@@ -31,8 +33,8 @@ using USearchIndexWithSerializationPtr = std::shared_ptr<USearchIndexWithSeriali
 template <unum::usearch::metric_kind_t Metric>
 struct MergeTreeIndexGranuleUSearch final : public IMergeTreeIndexGranule
 {
- MergeTreeIndexGranuleUSearch(const String & index_name_, const Block & index_sample_block_);
- MergeTreeIndexGranuleUSearch(const String & index_name_, const Block & index_sample_block_, USearchIndexWithSerializationPtr<Metric> index_);
+ MergeTreeIndexGranuleUSearch(const String & index_name_, const Block & index_sample_block_, unum::usearch::scalar_kind_t scalar_kind_);
+ MergeTreeIndexGranuleUSearch(const String & index_name_, const Block & index_sample_block_, unum::usearch::scalar_kind_t scalar_kind_, USearchIndexWithSerializationPtr<Metric> index_);
 
  ~MergeTreeIndexGranuleUSearch() override = default;
 
@@ -43,14 +45,15 @@ struct MergeTreeIndexGranuleUSearch final : public IMergeTreeIndexGranule
 
  const String index_name;
  const Block index_sample_block;
+ const unum::usearch::scalar_kind_t scalar_kind;
  USearchIndexWithSerializationPtr<Metric> index;
 };
 
 
 template <unum::usearch::metric_kind_t Metric>
 struct MergeTreeIndexAggregatorUSearch final : IMergeTreeIndexAggregator
 {
- MergeTreeIndexAggregatorUSearch(const String & index_name_, const Block & index_sample_block);
+ MergeTreeIndexAggregatorUSearch(const String & index_name_, const Block & index_sample_block, unum::usearch::scalar_kind_t scalar_kind_);
  ~MergeTreeIndexAggregatorUSearch() override = default;
 
  bool empty() const override { return !index || index->size() == 0; }
@@ -59,6 +62,7 @@ struct MergeTreeIndexAggregatorUSearch final : IMergeTreeIndexAggregator
 
  const String index_name;
  const Block index_sample_block;
+ const unum::usearch::scalar_kind_t scalar_kind;
  USearchIndexWithSerializationPtr<Metric> index;
 };
 
@@ -90,7 +94,7 @@ class MergeTreeIndexConditionUSearch final : public IMergeTreeIndexConditionAppr
 class MergeTreeIndexUSearch : public IMergeTreeIndex
 {
 public:
- MergeTreeIndexUSearch(const IndexDescription & index_, const String & distance_function_);
+ MergeTreeIndexUSearch(const IndexDescription & index_, const String & distance_function_, unum::usearch::scalar_kind_t scalar_kind_);
 
  ~MergeTreeIndexUSearch() override = default;
 
@@ -102,6 +106,7 @@ class MergeTreeIndexUSearch : public IMergeTreeIndex
 
 private:
  const String distance_function;
+ const unum::usearch::scalar_kind_t scalar_kind;
 };
 
 }