Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion cmake/GetHNSW.cmake
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
set ( HNSW_GITHUB "https://github.com/manticoresoftware/hnswlib/archive/6568d3b.zip" )
set ( HNSW_GITHUB "https://github.com/manticoresoftware/hnswlib/archive/091f3dd.zip" )
set ( HNSW_BUNDLEZIP "${LIBS_BUNDLE}/hnswlib-0.7.0.tar.gz" )

cmake_minimum_required ( VERSION 3.17 FATAL_ERROR )
Expand Down
157 changes: 125 additions & 32 deletions knn/knn.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -310,6 +310,48 @@ class L2BinarySIMD16ResidualsDistFn_c : public DistFnDispatch_c<&L2BinaryFloatDi
};
#endif

// build-mode DistFn classes
using IPBinaryGenericBuildDistFn_c = DistFnDispatch_c<&IPBinaryFloatDistanceGenericBuild>;
using L2BinaryGenericBuildDistFn_c = DistFnDispatch_c<&L2BinaryFloatDistanceGenericBuild>;

#if !defined(USE_SIMDE)
class IPBinarySIMD16BuildDistFn_c : public DistFnDispatch_c<&IPBinaryFloatDistanceSIMD16Build>
{
public:
static void Eval2 ( const void * pVect1, const void * pVect2A, const void * pVect2B, size_t uRowID1, size_t uRowID2A, size_t uRowID2B, const void * pParam, float & fDistA, float & fDistB )
{
IPBinaryFloatDistanceSIMD16Batch2Build ( pVect1, pVect2A, pVect2B, uRowID1, uRowID2A, uRowID2B, pParam, fDistA, fDistB );
}
};

class IPBinarySIMD16ResidualsBuildDistFn_c : public DistFnDispatch_c<&IPBinaryFloatDistanceSIMD16ResidualsBuild>
{
public:
static void Eval2 ( const void * pVect1, const void * pVect2A, const void * pVect2B, size_t uRowID1, size_t uRowID2A, size_t uRowID2B, const void * pParam, float & fDistA, float & fDistB )
{
IPBinaryFloatDistanceSIMD16ResidualsBatch2Build ( pVect1, pVect2A, pVect2B, uRowID1, uRowID2A, uRowID2B, pParam, fDistA, fDistB );
}
};

class L2BinarySIMD16BuildDistFn_c : public DistFnDispatch_c<&L2BinaryFloatDistanceSIMD16Build>
{
public:
static void Eval2 ( const void * pVect1, const void * pVect2A, const void * pVect2B, size_t uRowID1, size_t uRowID2A, size_t uRowID2B, const void * pParam, float & fDistA, float & fDistB )
{
L2BinaryFloatDistanceSIMD16Batch2Build ( pVect1, pVect2A, pVect2B, uRowID1, uRowID2A, uRowID2B, pParam, fDistA, fDistB );
}
};

class L2BinarySIMD16ResidualsBuildDistFn_c : public DistFnDispatch_c<&L2BinaryFloatDistanceSIMD16ResidualsBuild>
{
public:
static void Eval2 ( const void * pVect1, const void * pVect2A, const void * pVect2B, size_t uRowID1, size_t uRowID2A, size_t uRowID2B, const void * pParam, float & fDistA, float & fDistB )
{
L2BinaryFloatDistanceSIMD16ResidualsBatch2Build ( pVect1, pVect2A, pVect2B, uRowID1, uRowID2A, uRowID2B, pParam, fDistA, fDistB );
}
};
#endif

template <typename DistFn = void>
static void RunSearchPath ( const hnswlib::HierarchicalNSW<float> & tAlg, std::vector<DocDist_t> & dResults, const void * pData, int64_t iResults, HNSWFilterWrapper_c * pFilter, size_t * pSearchEf, int iSearchPath )
{
Expand Down Expand Up @@ -359,7 +401,8 @@ void HNSWIndex_c::Search ( std::vector<DocDist_t> & dResults, const Span_T<float
const void * pData = dData.begin();
if ( m_pQuantizer )
{
m_pQuantizer->Encode ( 0, dData, dQuantized );
std::vector<uint8_t> dUnusedQuantizedForQuery;
m_pQuantizer->Encode ( 0, dData, dQuantized, dUnusedQuantizedForQuery );
pData = dQuantized.data();
}

Expand Down Expand Up @@ -536,7 +579,8 @@ class HNSWIndexBuilder_i
virtual ~HNSWIndexBuilder_i() = default;

virtual void Train ( const util::Span_T<float> & dData ) = 0;
virtual bool AddDoc ( uint32_t uRowID, const util::Span_T<float> & dData, std::string & sError ) = 0;
virtual bool FinalizeTraining ( std::string & sError ) = 0;
virtual bool AddDoc ( uint32_t uRowID, const util::Span_T<float> & dData, BuildContext_t & tBuildCtx, std::string & sError ) = 0;
virtual void Save ( FileWriter_c & tWriter ) = 0;
virtual const AttrWithSettings_t & GetAttr() const = 0;
virtual const QuantizationSettings_t & GetQuantizationSettings() const = 0;
Expand All @@ -549,28 +593,56 @@ class HNSWIndexBuilder_c : public HNSWIndexBuilder_i, public HNSWDist_c
HNSWIndexBuilder_c ( const AttrWithSettings_t & tAttr, int64_t iNumElements, ScalarQuantizer_i * pQuantizer );

void Train ( const util::Span_T<float> & dData ) override;
bool AddDoc ( uint32_t uRowID, const util::Span_T<float> & dData, std::string & sError ) override;
bool FinalizeTraining ( std::string & sError ) override;
bool AddDoc ( uint32_t uRowID, const util::Span_T<float> & dData, BuildContext_t & tBuildCtx, std::string & sError ) override;
void Save ( FileWriter_c & tWriter ) override;
const AttrWithSettings_t & GetAttr() const override { return m_tAttr; }
const QuantizationSettings_t & GetQuantizationSettings() const override { return m_pQuantizer->GetSettings(); }

private:
AttrWithSettings_t m_tAttr;
bool m_bFirstDoc = true;
SpanResizeable_T<float> m_dNormalized;
std::vector<uint8_t> m_dQuantized;
std::unique_ptr<ScalarQuantizer_i> m_pQuantizer;
std::unique_ptr<hnswlib::HierarchicalNSW<float>> m_pAlg;
using AddPoint_fn = void (*) ( hnswlib::HierarchicalNSW<float> &, const void *, uint32_t );

template <typename DistFn>
static void AddPointTyped ( hnswlib::HierarchicalNSW<float> & tAlg, const void * pVec, uint32_t uRowID ) { tAlg.template addPoint<DistFn, false> ( pVec, (size_t)uRowID, -1 ); }
static void AddPointFallback ( hnswlib::HierarchicalNSW<float> & tAlg, const void * pVec, uint32_t uRowID ) { tAlg.addPoint ( pVec, (size_t)uRowID ); }
AddPoint_fn SelectAddPointFn() const;

AttrWithSettings_t m_tAttr;
std::unique_ptr<ScalarQuantizer_i> m_pQuantizer;
std::unique_ptr<hnswlib::HierarchicalNSW<float>> m_pAlg;
AddPoint_fn m_fnAddPoint = AddPointFallback;
};


HNSWIndexBuilder_c::AddPoint_fn HNSWIndexBuilder_c::SelectAddPointFn() const
{
switch ( m_pSpace->GetDistFuncId() )
{
case DistFuncId_e::IP_FLOAT32: return AddPointTyped<IPFloatDistFn_c>;
case DistFuncId_e::L2_FLOAT32: return AddPointTyped<L2FloatDistFn_c>;
case DistFuncId_e::IP_BINARY_GENERIC: return AddPointTyped<IPBinaryGenericBuildDistFn_c>;
case DistFuncId_e::L2_BINARY_GENERIC: return AddPointTyped<L2BinaryGenericBuildDistFn_c>;

#if !defined(USE_SIMDE)
case DistFuncId_e::IP_BINARY_SIMD16: return AddPointTyped<IPBinarySIMD16BuildDistFn_c>;
case DistFuncId_e::IP_BINARY_SIMD16_RESIDUALS: return AddPointTyped<IPBinarySIMD16ResidualsBuildDistFn_c>;
case DistFuncId_e::L2_BINARY_SIMD16: return AddPointTyped<L2BinarySIMD16BuildDistFn_c>;
case DistFuncId_e::L2_BINARY_SIMD16_RESIDUALS: return AddPointTyped<L2BinarySIMD16ResidualsBuildDistFn_c>;
#endif

default:
return AddPointFallback;
}
}


HNSWIndexBuilder_c::HNSWIndexBuilder_c ( const AttrWithSettings_t & tAttr, int64_t iNumElements, ScalarQuantizer_i * pQuantizer )
: HNSWDist_c ( tAttr.m_iDims, tAttr.m_eHNSWSimilarity, tAttr.m_eQuantization, true )
, m_tAttr ( tAttr )
, m_pQuantizer ( pQuantizer )
{
m_pAlg = std::make_unique<hnswlib::HierarchicalNSW<float>>( m_pSpace.get(), iNumElements, m_tAttr.m_iHNSWM, m_tAttr.m_iHNSWEFConstruction );
m_dNormalized.resize ( tAttr.m_iDims );
m_fnAddPoint = SelectAddPointFn();
}


Expand All @@ -581,39 +653,51 @@ void HNSWIndexBuilder_c::Train ( const util::Span_T<float> & dData )
}


bool HNSWIndexBuilder_c::AddDoc ( uint32_t uRowID, const util::Span_T<float> & dData, std::string & sError )
bool HNSWIndexBuilder_c::FinalizeTraining ( std::string & sError )
{
if ( dData.size()!=m_tAttr.m_iDims )
if ( !m_pQuantizer )
return true;

if ( m_pQuantizer->IsFinalized() )
return true;

if ( !m_pQuantizer->FinalizeTraining ( sError ) )
return false;

m_pSpace->SetQuantizationSettings ( *m_pQuantizer );
return true;
}


bool HNSWIndexBuilder_c::AddDoc ( uint32_t uRowID, const util::Span_T<float> & dData, BuildContext_t & tBuildCtx, std::string & sError )
{
if ( dData.size()!=(size_t)m_tAttr.m_iDims )
{
sError = FormatStr ( "HNSW error: data has %llu values, index '%s' needs %d values", dData.size(), m_tAttr.m_sName.c_str(), m_tAttr.m_iDims );
return false;
}

assert ( !m_pQuantizer || m_pQuantizer->IsFinalized() );

Span_T<float> dToAdd = dData;
if ( m_tAttr.m_eHNSWSimilarity==HNSWSimilarity_e::COSINE )
{
memcpy ( m_dNormalized.data(), dData.data(), dData.size()*sizeof(dData[0] ) );
VecNormalize(m_dNormalized);
dToAdd = m_dNormalized;
tBuildCtx.m_dNormalized.resize ( dData.size() );
memcpy ( tBuildCtx.m_dNormalized.data(), dData.data(), dData.size()*sizeof(dData[0] ) );
VecNormalize ( tBuildCtx.m_dNormalized );
dToAdd = tBuildCtx.m_dNormalized;
}

const void * pVec = nullptr;
if ( m_pQuantizer )
{
if ( m_bFirstDoc )
{
m_bFirstDoc = false;

if ( !m_pQuantizer->FinalizeTraining(sError) )
return false;

m_pSpace->SetQuantizationSettings ( *m_pQuantizer );
}

m_pQuantizer->Encode ( uRowID, dToAdd, m_dQuantized );
m_pAlg->addPoint ( (void*)m_dQuantized.data(), (size_t)uRowID );
m_pQuantizer->Encode ( uRowID, dToAdd, tBuildCtx.m_dQuantized, tBuildCtx.m_dQuantizedForQuery );
pVec = (void*)tBuildCtx.m_dQuantized.data();
}
else
m_pAlg->addPoint ( (void*)dToAdd.data(), (size_t)uRowID );
pVec = (void*)dToAdd.data();

m_fnAddPoint ( *m_pAlg, pVec, uRowID );

return true;
}
Expand All @@ -634,14 +718,13 @@ class HNSWBuilder_c : public Builder_i
public:
HNSWBuilder_c ( const Schema_t & tSchema, int64_t iNumElements, const std::string & sTmpFilename );

void Train ( int iAttr, uint32_t uRowID, const util::Span_T<float> & dData ) override { m_dIndexes[iAttr]->Train(dData); }
bool SetAttr ( int iAttr, uint32_t uRowID, const util::Span_T<float> & dData ) override { return m_dIndexes[iAttr]->AddDoc ( uRowID, dData, m_sError ); }
void Train ( int iAttr, uint32_t uRowID, const util::Span_T<float> & dData ) override { m_dIndexes[iAttr]->Train(dData); }
bool SetAttr ( int iAttr, uint32_t uRowID, const util::Span_T<float> & dData, BuildContext_t & tBuildCtx ) override { return m_dIndexes[iAttr]->AddDoc ( uRowID, dData, tBuildCtx, tBuildCtx.m_sError ); }
bool FinalizeTraining ( std::string & sError ) override;
bool Save ( const std::string & sFilename, size_t tBufferSize, std::string & sError ) override;
const std::string & GetError() const override { return m_sError; }

private:
std::vector<std::unique_ptr<HNSWIndexBuilder_i>> m_dIndexes;
std::string m_sError;
};


Expand All @@ -653,6 +736,16 @@ HNSWBuilder_c::HNSWBuilder_c ( const Schema_t & tSchema, int64_t iNumElements, c
}


bool HNSWBuilder_c::FinalizeTraining ( std::string & sError )
{
for ( auto & i : m_dIndexes )
if ( !i->FinalizeTraining(sError) )
return false;

return true;
}


bool HNSWBuilder_c::Save ( const std::string & sFilename, size_t tBufferSize, std::string & sError )
{
FileWriter_c tWriter;
Expand Down
15 changes: 12 additions & 3 deletions knn/knn.h
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@
namespace knn
{

static const int LIB_VERSION = 13;
static const int LIB_VERSION = 14;
static const uint32_t STORAGE_VERSION = 3;

enum class HNSWSimilarity_e
Expand Down Expand Up @@ -122,15 +122,24 @@ class KNN_i
virtual bool ShouldUseFullscan ( const std::string & sName, int64_t iResults, int iEf, int64_t iFilterCount ) = 0;
};

// passed via SetAttr so the builder itself holds no per-row mutable state
struct BuildContext_t
{
util::SpanResizeable_T<float> m_dNormalized;
std::vector<uint8_t> m_dQuantized;
std::vector<uint8_t> m_dQuantizedForQuery; // 4-bit transposed representation, produced only by the BIT1 binary quantizer during BUILD mode
std::string m_sError;
};

class Builder_i
{
public:
virtual ~Builder_i() = default;

virtual void Train ( int iAttr, uint32_t uRowID, const util::Span_T<float> & dData ) = 0;
virtual bool SetAttr ( int iAttr, uint32_t uRowID, const util::Span_T<float> & dData ) = 0;
virtual bool SetAttr ( int iAttr, uint32_t uRowID, const util::Span_T<float> & dData, BuildContext_t & tBuildCtx ) = 0;
virtual bool FinalizeTraining ( std::string & sError ) = 0;
virtual bool Save ( const std::string & sFilename, size_t tBufferSize, std::string & sError ) = 0;
virtual const std::string & GetError() const = 0;
};

class TextToEmbeddings_i
Expand Down
Loading
Loading