diff --git a/.github/workflows/pack_publish.yml b/.github/workflows/pack_publish.yml index c3973063fe..d6e1be94a4 100644 --- a/.github/workflows/pack_publish.yml +++ b/.github/workflows/pack_publish.yml @@ -31,7 +31,39 @@ concurrency: cancel-in-progress: true jobs: - + check_branch: + name: Check branch existence + runs-on: ubuntu-22.04 + outputs: + columnar_locator: ${{ steps.set_locator.outputs.columnar_locator }} + steps: + - name: Check if branch exists in manticoresoftware/manticoresearch + id: check_branch + if: github.ref_name != 'master' + run: | + # Extract the actual branch name for pull requests + if [[ "${{ github.event_name }}" == "pull_request" ]]; then + BRANCH_NAME="${{ github.event.pull_request.head.ref }}" + else + BRANCH_NAME="${{ github.ref_name }}" + fi + + HTTP_STATUS=$(curl -s -o /dev/null -w "%{http_code}" https://api.github.com/repos/manticoresoftware/columnar/branches/$BRANCH_NAME) + if [ "$HTTP_STATUS" -eq "200" ]; then + echo "branch_exists=true" >> $GITHUB_OUTPUT + echo "branch_name=$BRANCH_NAME" >> $GITHUB_OUTPUT + else + echo "branch_exists=false" >> $GITHUB_OUTPUT + echo "branch_name=$BRANCH_NAME" >> $GITHUB_OUTPUT + fi + - name: Set Columnar Locator + id: set_locator + run: | + if [[ "${{ github.ref_name }}" != "master" && "${{ steps.check_branch.outputs.branch_exists }}" == "true" ]]; then + echo "columnar_locator=GIT_REPOSITORY https://github.com/manticoresoftware/columnar.git GIT_TAG ${{ steps.check_branch.outputs.branch_name }}" >> $GITHUB_OUTPUT + else + echo "columnar_locator=" >> $GITHUB_OUTPUT + fi pack: name: OK to pack? runs-on: ubuntu-22.04 @@ -156,13 +188,15 @@ jobs: pack_debian_ubuntu: name: Debian/Ubuntu packages uses: ./.github/workflows/build_template.yml - needs: [pack, check_deps] + needs: [pack, check_branch, check_deps] + if: needs.pack.outputs.should_continue == 'true' strategy: fail-fast: false matrix: DISTR: [bionic, focal, jammy, buster, bullseye, bookworm] arch: [x86_64, aarch64] with: + COLUMNAR_LOCATOR: ${{ needs.check_branch.outputs.columnar_locator }} DISTR: ${{ matrix.DISTR }} arch: ${{ matrix.arch }} cmake_command: | @@ -179,13 +213,15 @@ jobs: pack_rhel: name: RHEL packages uses: ./.github/workflows/build_template.yml - needs: [pack, check_deps] + needs: [pack, check_branch, check_deps] + if: needs.pack.outputs.should_continue == 'true' strategy: fail-fast: false matrix: DISTR: [rhel7, rhel8, rhel9] arch: [x86_64, aarch64] with: + COLUMNAR_LOCATOR: ${{ needs.check_branch.outputs.columnar_locator }} DISTR: ${{ matrix.DISTR }} arch: ${{ matrix.arch }} boost_url_key: boost_rhel_feb17 @@ -205,13 +241,15 @@ jobs: pack_macos: name: MacOS packages uses: ./.github/workflows/build_template.yml - needs: [pack, check_deps] + needs: [pack, check_branch, check_deps] + if: needs.pack.outputs.should_continue == 'true' strategy: fail-fast: false matrix: DISTR: [macos] arch: [x86_64, arm64] with: + COLUMNAR_LOCATOR: ${{ needs.check_branch.outputs.columnar_locator }} DISTR: ${{ matrix.DISTR }} arch: ${{ matrix.arch }} HOMEBREW_PREFIX: /opt/homebrew @@ -229,8 +267,10 @@ jobs: pack_windows: name: Windows x64 package uses: ./.github/workflows/build_template.yml - needs: [pack, check_deps] + needs: [pack, check_branch, check_deps] + if: needs.pack.outputs.should_continue == 'true' with: + COLUMNAR_LOCATOR: ${{ needs.check_branch.outputs.columnar_locator }} DISTR: windows arch: x64 sysroot_url_key: roots_mysql83_jan17 diff --git a/cmake/GetColumnar.cmake b/cmake/GetColumnar.cmake index 787c91d85f..138f88170a 100644 --- a/cmake/GetColumnar.cmake +++ b/cmake/GetColumnar.cmake @@ -40,6 +40,7 @@ endif() set ( LIB_MANTICORE_COLUMNAR "lib_manticore_columnar.${EXTENSION}" ) set ( LIB_MANTICORE_SECONDARY "lib_manticore_secondary.${EXTENSION}" ) set ( LIB_MANTICORE_KNN "lib_manticore_knn.${EXTENSION}" ) +set ( LIB_MANTICORE_KNN_EMBEDDINGS "libmanticore_knn_embeddings.${EXTENSION}" ) macro ( backup_paths ) set ( _CMAKE_FIND_ROOT_PATH "${CMAKE_FIND_ROOT_PATH}" ) diff --git a/config/config_cmake.h.in b/config/config_cmake.h.in index 1765ed68c1..a04da5dfb0 100644 --- a/config/config_cmake.h.in +++ b/config/config_cmake.h.in @@ -161,6 +161,9 @@ /* default name of knn search shared library */ #cmakedefine LIB_MANTICORE_KNN "${LIB_MANTICORE_KNN}" +/* default name of knn embeddings shared library */ +#cmakedefine LIB_MANTICORE_KNN_EMBEDDINGS "${LIB_MANTICORE_KNN_EMBEDDINGS}" + /* define to set custom directory for shared files. Absolute as /usr/local/share/manticore */ #define FULL_SHARE_DIR "${FULL_SHARE_DIR}" diff --git a/src/accumulator.cpp b/src/accumulator.cpp index 48442767bb..eebc207fe0 100644 --- a/src/accumulator.cpp +++ b/src/accumulator.cpp @@ -94,6 +94,223 @@ void RtAccum_t::Sort() } } + +static bool RepackBlob ( const CSphColumnInfo & tAttr, const CSphColumnInfo & tBlobLoc, int iBlobAttr, const CSphRowitem * pRow, const BYTE * pBlobPool, std::unique_ptr & pBlobWriter, CSphString & sError ) +{ + SphAttr_t tBlobRowOffset = sphGetRowAttr ( pRow, tBlobLoc.m_tLocator ); + const BYTE * pBlobRow = pBlobPool + tBlobRowOffset; + ByteBlob_t tBlob = sphGetBlobAttr ( pBlobRow, tAttr.m_tLocator ); + return pBlobWriter->SetAttr ( iBlobAttr, tBlob.first, tBlob.second, sError ); +} + + +static bool StoreEmbeddings ( const CSphSchema & tSchema, int iAttr, int iBlobAttr, int iColumnarAttr, DocstoreDoc_t & tDoc, std::unique_ptr & pNewBlobBuilder, std::unique_ptr & pNewColumnarBuilder, const IntVec_t & dDocstoreRemap, const std::vector & dEmbedding, CSphVector & dTmp, CSphString & sError ) +{ + const CSphColumnInfo & tAttr = tSchema.GetAttr(iAttr); + + if ( tAttr.IsColumnar() ) + { + dTmp.Resize ( dEmbedding.size() ); + ARRAY_FOREACH ( iEmb, dTmp ) + dTmp[iEmb] = sphF2DW ( dEmbedding[iEmb] ); + + pNewColumnarBuilder->SetAttr ( iColumnarAttr, dTmp.Begin(), dTmp.GetLength() ); + } + else + { + if ( !pNewBlobBuilder->SetAttr ( iBlobAttr, (const BYTE*)dEmbedding.data(), dEmbedding.size()*sizeof(float), sError ) ) + return false; + } + + if ( tAttr.IsStored() ) + { + int iId = dDocstoreRemap[iAttr]; + tDoc.m_dFields[iId].Resize ( dEmbedding.size()*sizeof(DWORD) ); + const BYTE * pStart = tDoc.m_dFields[iId].Begin(); + for ( auto i : dEmbedding ) + { + *(DWORD*)pStart = sphF2DW ( dEmbedding[i] ); + pStart += sizeof(DWORD); + } + } + + return true; +} + + +bool RtAccum_t::RebuildStoragesForEmbeddings ( RowID_t tRowID, CSphRowitem * pRow, const CSphVector & dAttrsWithModels, std::unique_ptr & pNewBlobBuilder, std::unique_ptr & pNewColumnarBuilder, std::unique_ptr & pNewDocstoreBuilder, CSphVector & dAllIterators, const IntVec_t & dDocstoreRemap, const CSphColumnInfo * pBlobLoc, std::vector>> & dAllEmbeddings, CSphVector & dTmp, CSphString & sError ) +{ + int iBlobAttr = 0; + int iColumnarAttr = 0; + DocstoreDoc_t tDoc; + + // fetch all fields and attributes without model from docstore (they will not be modified) + if ( pNewDocstoreBuilder ) + tDoc = m_pDocstore->GetDoc ( tRowID, nullptr, -1, false ); + + assert(m_pIndex); + const CSphSchema & tSchema = m_pIndex->GetInternalSchema(); + + ARRAY_FOREACH ( i, dAttrsWithModels ) + { + const CSphColumnInfo & tAttr = tSchema.GetAttr(i); + const AttrWithModel_t & tAttrWithModel = dAttrsWithModels[i]; + + if ( tAttrWithModel.m_pModel ) + { + if ( !StoreEmbeddings ( tSchema, i, iBlobAttr, iColumnarAttr, tDoc, pNewBlobBuilder, pNewColumnarBuilder, dDocstoreRemap, dAllEmbeddings[i][tRowID], dTmp, sError ) ) + return false; + } + else + { + if ( pNewBlobBuilder && !tAttr.IsColumnar() && sphIsBlobAttr(tAttr) ) + { + assert(pBlobLoc); + if ( !RepackBlob ( tAttr, *pBlobLoc, iBlobAttr, pRow, m_dBlobs.Begin(), pNewBlobBuilder, sError ) ) + return false; + } + + if ( pNewColumnarBuilder && tAttr.IsColumnar() ) + { + auto & tIt = dAllIterators[iColumnarAttr]; + SetColumnarAttr ( iColumnarAttr, tIt.second, pNewColumnarBuilder.get(), tIt.first, tRowID, dTmp ); + } + } + + if ( sphIsBlobAttr(tAttr) ) + iBlobAttr++; + + if ( tAttr.IsColumnar() ) + iColumnarAttr++; + } + + if ( pNewBlobBuilder ) + { + assert(pBlobLoc); + sphSetRowAttr ( pRow, pBlobLoc->m_tLocator, pNewBlobBuilder->Flush().first ); + } + + if ( pNewDocstoreBuilder ) + pNewDocstoreBuilder->AddDoc ( tRowID, tDoc ); + + return true; +} + + +bool RtAccum_t::GenerateEmbeddings ( int iAttr, int iAttrWithModel, const CSphVector & dAttrsWithModels, std::vector>> & dAllEmbeddings, CSphString & sError ) +{ + const AttrWithModel_t & tAttrWithModel = dAttrsWithModels[iAttr]; + if ( !tAttrWithModel.m_pModel ) + return true; + + assert(m_pIndex); + const CSphSchema & tSchema = m_pIndex->GetInternalSchema(); + + int iRowSize = tSchema.GetRowSize(); + std::vector> & dEmbeddingsForAttr = dAllEmbeddings[iAttr]; + std::vector dTexts; + CSphRowitem * pRow = m_dAccumRows.Begin(); + for ( RowID_t tRowID = 0; tRowID < m_uAccumDocs; ++tRowID, pRow += iRowSize ) + { + const auto & dConcat = m_pEmbeddingsSrc->Get ( tRowID, iAttrWithModel ); + dTexts.push_back ( { dConcat.Begin(), (size_t)dConcat.GetLength() } ); + } + + std::string sErrorSTL; + if ( !tAttrWithModel.m_pModel->Convert ( dTexts, dEmbeddingsForAttr, sErrorSTL ) ) + { + sError = sErrorSTL.c_str(); + return false; + } + + return true; +} + + +bool RtAccum_t::FetchEmbeddings ( TableEmbeddings_c * pEmbeddings, const CSphVector & dAttrsWithModels, CSphString & sError ) +{ + if ( !pEmbeddings ) + return true; + + assert(m_pIndex); + const CSphSchema & tSchema = m_pIndex->GetInternalSchema(); + + // we need to rebuild our blobs/columnar storage (to replace empty data with fetched floatvectors), but they are immutable by design + // so we'll have to rebuild the blobs/columnar storage fully + const CSphColumnInfo * pBlobLoc = tSchema.GetAttr ( sphGetBlobLocatorName() ); + bool bRebuildColumnar = false; + bool bRebuildBlobs = false; + bool bRebuildDocstore = false; + IntVec_t dDocstoreRemap; + dDocstoreRemap.Resize ( dAttrsWithModels.GetLength() ); + dDocstoreRemap.Fill(-1); + ARRAY_FOREACH ( i, dAttrsWithModels ) + { + if ( !dAttrsWithModels[i].m_pModel ) + continue; + + auto & tAttr = tSchema.GetAttr(i); + assert ( tAttr.m_eAttrType==SPH_ATTR_FLOAT_VECTOR ); + bRebuildColumnar |= tAttr.IsColumnar(); + bRebuildBlobs |= !tAttr.IsColumnar(); + bRebuildDocstore |= tAttr.IsStored(); + + dDocstoreRemap[i] = ((DocstoreBuilder_i*)m_pDocstore.get())->GetFieldId ( tAttr.m_sName, DOCSTORE_ATTR ); + } + + CSphTightVector dNewBlobPool; + std::unique_ptr pNewBlobBuilder = bRebuildBlobs ? sphCreateBlobRowBuilder ( tSchema, dNewBlobPool ) : nullptr; + std::unique_ptr pNewColumnarBuilder = bRebuildColumnar ? CreateColumnarBuilderRT(tSchema) : nullptr; + std::unique_ptr pColumnar = bRebuildColumnar ? CreateLightColumnarRT ( tSchema, m_pColumnarBuilder.get() ) : nullptr; + CSphVector dAllIterators; + if ( bRebuildColumnar ) + dAllIterators = CreateAllColumnarIterators ( pColumnar.get(), tSchema ); + + std::unique_ptr pNewDocstoreBuilder; + if ( bRebuildDocstore ) + { + pNewDocstoreBuilder = CreateDocstoreRT(); + SetupDocstoreFields ( *pNewDocstoreBuilder, tSchema ); + } + + // 1st pass - generate all embeddings for each attribute + int iRowSize = tSchema.GetRowSize(); + int iAttrWithModel = 0; + std::vector>> dAllEmbeddings; + dAllEmbeddings.resize ( dAttrsWithModels.GetLength() ); + ARRAY_FOREACH ( i, dAttrsWithModels ) + { + if ( !GenerateEmbeddings( i, iAttrWithModel, dAttrsWithModels, dAllEmbeddings, sError ) ) + return false; + + if ( dAttrsWithModels[i].m_pModel ) + iAttrWithModel++; + } + + // 2nd pass - rebuild attribute storages + std::vector dEmbedding; + CSphVector dTmp; + CSphRowitem * pRow = m_dAccumRows.Begin(); + for ( RowID_t tRowID = 0; tRowID < m_uAccumDocs; ++tRowID, pRow += iRowSize ) + if ( !RebuildStoragesForEmbeddings ( tRowID, pRow, dAttrsWithModels, pNewBlobBuilder, pNewColumnarBuilder, pNewDocstoreBuilder, dAllIterators, dDocstoreRemap, pBlobLoc, dAllEmbeddings, dTmp, sError ) ) + return false; + + if ( bRebuildBlobs ) + { + m_pBlobWriter = std::move(pNewBlobBuilder); + m_dBlobs.SwapData(dNewBlobPool); + } + + if ( bRebuildColumnar ) + m_pColumnarBuilder = std::move(pNewColumnarBuilder); + + if ( bRebuildDocstore ) + m_pDocstore = std::move(pNewDocstoreBuilder); + + return true; +} + + void RtAccum_t::CleanupPart() { m_dAccumRows.Resize ( 0 ); @@ -102,6 +319,7 @@ void RtAccum_t::CleanupPart() m_dPerDocHitsCount.Resize ( 0 ); m_dAccum.Resize ( 0 ); m_pDocstore.reset(); + m_pEmbeddingsSrc.reset(); ResetDict(); ResetRowID(); @@ -176,6 +394,81 @@ static void ResetTailHit ( CSphWordHit * pHit ) pHit->m_uWordPos = HITMAN::GetPosWithField ( pHit->m_uWordPos ); } + +static const char * FetchStringFromDoc ( int iAttr, const InsertDocData_c & tDoc, const ISphSchema & tSchema ) +{ + const char ** ppStr = tDoc.m_dStrings.Begin(); + int iStrAttr = 0; + + for ( int i=0; i < tSchema.GetAttrsCount(); ++i ) + { + const CSphColumnInfo & tAttr = tSchema.GetAttr(i); + + if ( tAttr.m_eAttrType!=SPH_ATTR_STRING ) + continue; + + if ( iAttr==i ) + return ppStr[iStrAttr]; + + iStrAttr++; + } + + return nullptr; +} + + +static CSphVector ConcatFromFields ( const InsertDocData_c & tDoc, const AttrWithModel_t & tAttr, const ISphSchema & tSchema ) +{ + CSphVector dTmp; + ARRAY_FOREACH ( i, tAttr.m_dFrom ) + { + auto tFrom = tAttr.m_dFrom[i]; + int iAttrFieldId = tFrom.first; + VecTraits_T dSrc; + if ( tFrom.second ) + dSrc = { tDoc.m_dFields[iAttrFieldId].Begin(), tDoc.m_dFields[iAttrFieldId].GetLength() }; + else + { + const char * szString = FetchStringFromDoc ( iAttrFieldId, tDoc, tSchema ); + dSrc = { szString, (int64_t)( szString ? strlen(szString) : 0 ) }; + } + + int iOldSize = dTmp.GetLength(); + dTmp.Resize ( iOldSize + dSrc.GetLength() + 1 ); + dTmp[iOldSize] = ' '; + memcpy ( dTmp.Begin() + iOldSize + 1, dSrc.Begin(), dSrc.GetLength() ); + } + + return dTmp; +} + + +void RtAccum_t::FetchEmbeddingsSrc ( InsertDocData_c & tDoc, const CSphVector & dAttrsWithModels ) +{ + if ( !m_pEmbeddingsSrc ) + { + int iAttrsWithModels = dAttrsWithModels.count_of ( []( auto & tData ){ return !!tData.m_pModel; } ); + m_pEmbeddingsSrc = std::make_unique(iAttrsWithModels); + } + + assert(m_pIndex); + const CSphSchema & tSchema = m_pIndex->GetInternalSchema(); + + CSphVector dTmp; + int iAttrWithModel = 0; + ARRAY_FOREACH ( i, dAttrsWithModels ) + { + const AttrWithModel_t & tAttrWithModel = dAttrsWithModels[i]; + if ( !tAttrWithModel.m_pModel ) + continue; + + auto dConcat = ConcatFromFields ( tDoc, tAttrWithModel, tSchema ); + m_pEmbeddingsSrc->Add ( iAttrWithModel, dConcat ); + iAttrWithModel++; + } +} + + void RtAccum_t::AddDocument ( ISphHits* pHits, const InsertDocData_c& tDoc, bool bReplace, int iRowSize, const DocstoreBuilder_i::Doc_t* pStoredDoc ) { MEMORY ( MEM_RT_ACCUM ); @@ -473,6 +766,8 @@ void RtAccum_t::CleanupDuplicates ( int iRowSize ) m_dAccum.Resize ( iDstRow ); RemoveColumnarDuplicates ( m_pColumnarBuilder, dRowMap, tSchema ); + if ( m_pEmbeddingsSrc ) + m_pEmbeddingsSrc->Remove(dRowMap); iDstRow = 0; ARRAY_FOREACH ( i, dRowMap ) diff --git a/src/accumulator.h b/src/accumulator.h index 63ab374db6..eae9caff5f 100644 --- a/src/accumulator.h +++ b/src/accumulator.h @@ -17,6 +17,7 @@ #include "sphinxint.h" #include "docstore.h" +#include "knnmisc.h" struct StoredQueryDesc_t { @@ -98,8 +99,15 @@ struct ReplicatedCommand_t std::unique_ptr MakeReplicationCommand ( ReplCmd_e eCommand, CSphString sIndex, CSphString sCluster = CSphString() ); +struct AttrWithModel_t +{ + knn::TextToEmbeddings_i * m_pModel = nullptr; + CSphVector> m_dFrom; +}; + class RtIndex_i; class ColumnarBuilderRT_i; +class TableEmbeddings_c; /// indexing accumulator class RtAccum_t @@ -122,6 +130,8 @@ class RtAccum_t public: void SetupDict ( const RtIndex_i * pIndex, const DictRefPtr_c& pDict, bool bKeywordDict ); void Sort(); + bool FetchEmbeddings ( TableEmbeddings_c * pEmbeddings, const CSphVector & dAttrsWithModels, CSphString & sError ); + void FetchEmbeddingsSrc ( InsertDocData_c & tDoc, const CSphVector & dAttrsWithModels ); void CleanupPart(); void Cleanup(); @@ -162,6 +172,7 @@ class RtAccum_t std::unique_ptr m_pBlobWriter; std::unique_ptr m_pDocstore; std::unique_ptr m_pColumnarBuilder; + std::unique_ptr m_pEmbeddingsSrc; RowID_t m_tNextRowID = 0; CSphFixedVector m_dPackedKeywords { 0 }; uint64_t m_uSchemaHash = 0; @@ -175,6 +186,9 @@ class RtAccum_t void ResetDict(); void SetupDocstore(); + bool RebuildStoragesForEmbeddings ( RowID_t tRowID, CSphRowitem * pRow, const CSphVector & dAttrsWithModels, std::unique_ptr & pNewBlobBuilder, std::unique_ptr & pNewColumnarBuilder, std::unique_ptr & pNewDocstoreBuilder, CSphVector & dAllIterators, const IntVec_t & dDocstoreRemap, const CSphColumnInfo * pBlobLoc, std::vector>> & dAllEmbeddings, CSphVector & dTmp, CSphString & sError ); + bool GenerateEmbeddings ( int iAttr, int iAttrWithModel, const CSphVector & dAttrsWithModels, std::vector>> & dAllEmbeddings, CSphString & sError ); + // defined in sphinxrt.cpp friend RtSegment_t* CreateSegment ( RtAccum_t*, int, ESphHitless, const VecTraits_T&, CSphString& ); }; diff --git a/src/ddl.l b/src/ddl.l index 90aafdbb8a..a7cfc80df6 100644 --- a/src/ddl.l +++ b/src/ddl.l @@ -52,9 +52,11 @@ "AS" { YYSTOREBOUNDS; return TOK_AS; } "AT" { YYSTOREBOUNDS; return TOK_AT; } "ATTRIBUTE" { YYSTOREBOUNDS; return TOK_ATTRIBUTE; } +"API_KEY" { YYSTOREBOUNDS; return TOK_API_KEY; } "BIGINT" { YYSTOREBOUNDS; return TOK_BIGINT; } "BIT" { YYSTOREBOUNDS; return TOK_BIT; } "BOOL" { YYSTOREBOUNDS; return TOK_BOOL; } +"CACHE_PATH" { YYSTOREBOUNDS; return TOK_CACHE_PATH; } "CLUSTER" { YYSTOREBOUNDS; return TOK_CLUSTER; } "COLUMN" { YYSTOREBOUNDS; return TOK_COLUMN; } "COLUMNAR" { YYSTOREBOUNDS; return TOK_COLUMNAR; } @@ -83,12 +85,13 @@ "KNN_DIMS" { YYSTOREBOUNDS; return TOK_KNN_DIMS; } "KNN_TYPE" { YYSTOREBOUNDS; return TOK_KNN_TYPE; } "LIKE" { YYSTOREBOUNDS; return TOK_LIKE; } -"OPTION" { YYSTOREBOUNDS; return TOK_OPTION; } +"MODEL_NAME" { YYSTOREBOUNDS; return TOK_MODEL_NAME; } "MODIFY" { YYSTOREBOUNDS; return TOK_MODIFY; } "MODIFY"[ \t\n\r]+"COLUMN" { return TOK_MODIFY_COLUMN; } "MULTI" { YYSTOREBOUNDS; return TOK_MULTI; } "MULTI64" { YYSTOREBOUNDS; return TOK_MULTI64; } "NOT" { YYSTOREBOUNDS; return TOK_NOT; } +"OPTION" { YYSTOREBOUNDS; return TOK_OPTION; } "PLUGIN" { YYSTOREBOUNDS; return TOK_PLUGIN; } "QUANTIZATION" { YYSTOREBOUNDS; return TOK_QUANTIZATION; } "REBUILD" { YYSTOREBOUNDS; return TOK_REBUILD; } @@ -106,6 +109,7 @@ "TYPE" { YYSTOREBOUNDS; return TOK_TYPE; } "UINT" { YYSTOREBOUNDS; return TOK_UINT; } "UPDATE" { YYSTOREBOUNDS; return TOK_UPDATE; } +"USE_GPU" { YYSTOREBOUNDS; return TOK_USE_GPU; } '([^'\\]|\\.|\\\\)*' { YYSTOREBOUNDS; pParser->m_pLastTokenStart = yytext; lvalp->SetValueInt ( ( (SphAttr_t)lvalp->m_iStart<<32 ) | ( lvalp->m_iEnd-lvalp->m_iStart ) ); return TOK_QUOTED_STRING; } diff --git a/src/ddl.y b/src/ddl.y index 7bb9cd0c82..3228a91e55 100644 --- a/src/ddl.y +++ b/src/ddl.y @@ -24,12 +24,14 @@ %token TOK_ADD %token TOK_ALTER +%token TOK_API_KEY %token TOK_AS %token TOK_AT %token TOK_ATTRIBUTE %token TOK_BIGINT %token TOK_BIT %token TOK_BOOL +%token TOK_CACHE_PATH %token TOK_CLUSTER %token TOK_COLUMN %token TOK_COLUMNAR @@ -58,12 +60,13 @@ %token TOK_KNN_DIMS %token TOK_KNN_TYPE %token TOK_LIKE -%token TOK_OPTION -%token TOK_MULTI -%token TOK_MULTI64 +%token TOK_MODEL_NAME %token TOK_MODIFY %token TOK_MODIFY_COLUMN +%token TOK_MULTI +%token TOK_MULTI64 %token TOK_NOT +%token TOK_OPTION %token TOK_PLUGIN %token TOK_QUANTIZATION %token TOK_REBUILD @@ -81,6 +84,7 @@ %token TOK_TYPE %token TOK_UINT %token TOK_UPDATE +%token TOK_USE_GPU %% @@ -326,6 +330,46 @@ item_option: YYERROR; } } + | TOK_MODEL_NAME '=' TOK_QUOTED_STRING + { + if ( !pParser->AddItemOptionModelName ( $3 ) ) + { + yyerror ( pParser, pParser->GetLastError() ); + YYERROR; + } + } + | TOK_FROM '=' TOK_QUOTED_STRING + { + if ( !pParser->AddItemOptionFrom ( $3 ) ) + { + yyerror ( pParser, pParser->GetLastError() ); + YYERROR; + } + } + | TOK_API_KEY '=' TOK_QUOTED_STRING + { + if ( !pParser->AddItemOptionAPIKey ( $3 ) ) + { + yyerror ( pParser, pParser->GetLastError() ); + YYERROR; + } + } + | TOK_CACHE_PATH '=' TOK_QUOTED_STRING + { + if ( !pParser->AddItemOptionCachePath ( $3 ) ) + { + yyerror ( pParser, pParser->GetLastError() ); + YYERROR; + } + } + | TOK_USE_GPU '=' TOK_QUOTED_STRING + { + if ( !pParser->AddItemOptionUseGPU ( $3 ) ) + { + yyerror ( pParser, pParser->GetLastError() ); + YYERROR; + } + } | TOK_QUANTIZATION '=' TOK_QUOTED_STRING { if ( !pParser->AddItemOptionQuantization ( $3 ) ) diff --git a/src/indexsettings.cpp b/src/indexsettings.cpp index ac7ccedaac..0cb8e99387 100644 --- a/src/indexsettings.cpp +++ b/src/indexsettings.cpp @@ -1491,7 +1491,9 @@ void IndexSettingsContainer_c::SetupKNNAttrs ( const CreateTableSettings_t & tCr { NamedKNNSettings_t & tNamedKNN = dKNNAttrs.Add(); (knn::IndexSettings_t&)tNamedKNN = i.m_tKNN; + (knn::ModelSettings_t&)tNamedKNN = i.m_tKNNModel; tNamedKNN.m_sName = i.m_tAttr.m_sName; + tNamedKNN.m_sFrom = i.m_sKNNFrom; } if ( !dKNNAttrs.GetLength() ) @@ -2214,7 +2216,7 @@ static void AddStorageSettings ( StringBuilder_c & sRes, const CSphColumnInfo & bool bColumnar = CombineEngines ( tIndex.GetSettings().m_eEngine, tAttr.m_eEngine )==AttrEngine_e::COLUMNAR; if ( bColumnar ) { - if ( tAttr.m_eAttrType!=SPH_ATTR_JSON && !(tAttr.m_uAttrFlags & CSphColumnInfo::ATTR_STORED) && iNumColumnar>1 ) + if ( tAttr.m_eAttrType!=SPH_ATTR_JSON && !tAttr.IsStored() && iNumColumnar>1 ) sRes << " fast_fetch='0'"; if ( tAttr.m_eAttrType==SPH_ATTR_STRING && !(tAttr.m_uAttrFlags & CSphColumnInfo::ATTR_COLUMNAR_HASHES) ) diff --git a/src/indexsettings.h b/src/indexsettings.h index 2cbb3fdec8..474d267a2c 100644 --- a/src/indexsettings.h +++ b/src/indexsettings.h @@ -400,6 +400,8 @@ struct CreateTableAttr_t bool m_bIndexed = false; bool m_bKNN = false; knn::IndexSettings_t m_tKNN; + knn::ModelSettings_t m_tKNNModel; + CSphString m_sKNNFrom; }; struct NameValueStr_t diff --git a/src/knnlib.cpp b/src/knnlib.cpp index ee3cd6f4e9..ef06298be2 100644 --- a/src/knnlib.cpp +++ b/src/knnlib.cpp @@ -17,13 +17,17 @@ using Create_fn = knn::KNN_i * (*) (); using CreateBuilder_fn = knn::Builder_i * (*) ( const knn::Schema_t & tSchema, int64_t iNumElements ); using CreateDistanceCalc_fn = knn::Distance_i * (*) ( const knn::IndexSettings_t & tSettings ); +using LoadEmbeddingsLib_fn = knn::EmbeddingsLib_i * (*) ( const std::string & sLibPath, std::string & sError ); using VersionStr_fn = const char * (*)(); using GetVersion_fn = int (*)(); static void * g_pKNNLib = nullptr; +std::unique_ptr g_pEmbeddingsLib; + static Create_fn g_fnCreate = nullptr; static CreateBuilder_fn g_fnCreateKNNBuilder = nullptr; static CreateDistanceCalc_fn g_fnCreateDistanceCalc = nullptr; +static LoadEmbeddingsLib_fn g_fnLoadEmbeddingsLib = nullptr; static VersionStr_fn g_fnVersionStr = nullptr; ///////////////////////////////////////////////////////////////////// @@ -79,12 +83,33 @@ std::unique_ptr CreateKNNBuilder ( const ISphSchema & tSchema, i } -std::unique_ptr CreateKNNDistanceCalc ( const knn::IndexSettings_t & tSettings ) +std::unique_ptr CreateKNNDistanceCalc ( const knn::IndexSettings_t & tSettings, CSphString & sError ) { + if ( !IsKNNLibLoaded() ) + { + sError = "knn library not loaded"; + return nullptr; + } + return std::unique_ptr ( g_fnCreateDistanceCalc(tSettings) ); } +std::unique_ptr CreateTextToEmbeddings ( const knn::ModelSettings_t & tSettings, CSphString & sError ) +{ + if ( !g_pEmbeddingsLib ) + { + sError = "embeddings library not loaded"; + return nullptr; + } + + std::string sErrorSTL; + auto pRes = std::unique_ptr ( g_pEmbeddingsLib->CreateTextToEmbeddings ( tSettings, sErrorSTL ) ); + sError = sErrorSTL.c_str(); + return pRes; +} + + #if HAVE_DLOPEN bool InitKNN ( CSphString & sError ) { @@ -124,8 +149,18 @@ bool InitKNN ( CSphString & sError ) if ( !LoadFunc ( g_fnCreate, tHandle.Get(), "CreateKNN", sLibfile, sError ) ) return false; if ( !LoadFunc ( g_fnCreateKNNBuilder, tHandle.Get(), "CreateKNNBuilder", sLibfile, sError ) ) return false; if ( !LoadFunc ( g_fnCreateDistanceCalc, tHandle.Get(), "CreateDistanceCalc", sLibfile, sError ) ) return false; + if ( !LoadFunc ( g_fnLoadEmbeddingsLib, tHandle.Get(), "LoadEmbeddingsLib", sLibfile, sError ) ) return false; if ( !LoadFunc ( g_fnVersionStr, tHandle.Get(), "GetKNNLibVersionStr", sLibfile, sError ) ) return false; + std::string sEmbeddingsLibFile = TryDifferentPaths ( LIB_MANTICORE_KNN_EMBEDDINGS, GetKNNEmbeddingsFullpath(), 0 ).cstr(); + if ( !sEmbeddingsLibFile.empty() ) + { + std::string sErrorSTL; + g_pEmbeddingsLib = std::unique_ptr ( g_fnLoadEmbeddingsLib ( sEmbeddingsLibFile, sErrorSTL ) ); + if ( !g_pEmbeddingsLib ) + return false; + } + g_pKNNLib = tHandle.Leak(); return true; @@ -134,7 +169,10 @@ bool InitKNN ( CSphString & sError ) void ShutdownKNN() { if ( g_pKNNLib ) + { + g_pEmbeddingsLib.reset(); dlclose(g_pKNNLib); + } } #else @@ -152,7 +190,23 @@ const char * GetKNNVersionStr() } +const char * GetKNNEmbeddingsVersionStr() +{ + if ( !g_pEmbeddingsLib ) + return nullptr; + + return g_pEmbeddingsLib->GetVersionStr().c_str(); +} + + bool IsKNNLibLoaded() { return !!g_pKNNLib; } + + +bool IsKNNEmbeddingsLibLoaded() +{ + return !!g_pEmbeddingsLib; +} + diff --git a/src/knnlib.h b/src/knnlib.h index 5db968f37c..d163a63db7 100644 --- a/src/knnlib.h +++ b/src/knnlib.h @@ -11,6 +11,7 @@ #pragma once #include "sphinxexpr.h" +#include "std/stringhash.h" #include "knn/knn.h" @@ -18,9 +19,12 @@ class ISphSchema; std::unique_ptr CreateKNN ( CSphString & sError ); std::unique_ptr CreateKNNBuilder ( const ISphSchema & tSchema, int64_t iNumElements, CSphString & sError ); -std::unique_ptr CreateKNNDistanceCalc ( const knn::IndexSettings_t & tSettings ); +std::unique_ptr CreateKNNDistanceCalc ( const knn::IndexSettings_t & tSettings, CSphString & sError ); +std::unique_ptr CreateTextToEmbeddings ( const knn::ModelSettings_t & tSettings, CSphString & sError ); bool InitKNN ( CSphString & sError ); void ShutdownKNN(); const char * GetKNNVersionStr(); +const char * GetKNNEmbeddingsVersionStr(); bool IsKNNLibLoaded(); +bool IsKNNEmbeddingsLibLoaded(); diff --git a/src/knnmisc.cpp b/src/knnmisc.cpp index 8acf39dd3e..9f165b170b 100644 --- a/src/knnmisc.cpp +++ b/src/knnmisc.cpp @@ -17,6 +17,60 @@ #include "sphinxjson.h" +bool TableEmbeddings_c::Load ( const CSphString & sAttr, const knn::ModelSettings_t & tSettings, CSphString & sError ) +{ + if ( m_hModels.Exists(sAttr) ) + { + sError.SetSprintf ( "Unable to load model '%s' for attribute '%s': already loaded", tSettings.m_sModelName.c_str(), sAttr.cstr() ); + return false; + } + + std::unique_ptr pModel = CreateTextToEmbeddings ( tSettings, sError ); + if ( !pModel ) + return false; + + m_hModels.Add ( std::move(pModel), sAttr ); + return true; +} + + +knn::TextToEmbeddings_i * TableEmbeddings_c::GetModel ( const CSphString & sAttr ) const +{ + std::unique_ptr * pFound = m_hModels(sAttr); + return pFound ? pFound->get() : nullptr; +} + +/////////////////////////////////////////////////////////////////////////////// + +EmbeddingsSrc_c::EmbeddingsSrc_c ( int iAttrs ) +{ + m_dStored.Resize(iAttrs); +} + + +void EmbeddingsSrc_c::Add ( int iAttr, CSphVector & dSrc ) +{ + auto & tNew = m_dStored[iAttr].Add(); + tNew.SwapData(dSrc); +} + + +void EmbeddingsSrc_c::Remove ( const CSphFixedVector & dRowMap ) +{ + for ( auto & i : m_dStored ) + for ( auto tRowID : dRowMap ) + if ( tRowID==INVALID_ROWID ) + i.Remove(tRowID); +} + + +const VecTraits_T EmbeddingsSrc_c::Get ( RowID_t tRowID, int iAttr ) const +{ + return m_dStored[iAttr][tRowID]; +} + +/////////////////////////////////////////////////////////////////////////////// + void NormalizeVec ( VecTraits_T & dData ) { float fNorm = 0.0f; @@ -62,7 +116,8 @@ Expr_KNNDist_c::Expr_KNNDist_c ( const CSphVector & dAnchor, const CSphCo { knn::IndexSettings_t tDistSettings = tAttr.m_tKNN; tDistSettings.m_eQuantization = knn::Quantization_e::NONE; // we operate on non-quantized data - m_pDistCalc = CreateKNNDistanceCalc(tDistSettings); + CSphString sError; // fixme! report it + m_pDistCalc = CreateKNNDistanceCalc ( tDistSettings, sError ); if ( tAttr.m_tKNN.m_eHNSWSimilarity==knn::HNSWSimilarity_e::COSINE ) NormalizeVec(m_dAnchor); @@ -264,36 +319,59 @@ void AddKNNSettings ( StringBuilder_c & sRes, const CSphColumnInfo & tAttr ) if ( tKNN.m_iHNSWEFConstruction!=tDefault.m_iHNSWEFConstruction ) sRes << " hnsw_ef_construction='" << tKNN.m_iHNSWEFConstruction << "'"; + const auto & tKNNModel = tAttr.m_tKNNModel; + knn::ModelSettings_t tDefaultModel; + if ( !tKNNModel.m_sModelName.empty() ) + sRes << " model_name='" << tKNNModel.m_sModelName.c_str() << "'"; + + if ( !tKNNModel.m_sCachePath.empty() ) + sRes << " cache_path='" << tKNNModel.m_sCachePath.c_str() << "'"; + + if ( tKNNModel.m_bUseGPU!=tDefaultModel.m_bUseGPU ) + sRes << " use_gpu='" << ( tKNNModel.m_bUseGPU ? 1 : 0 ) << "'"; + if ( tKNN.m_eQuantization!=tDefault.m_eQuantization ) sRes << " quantization='" << Quantization2Str ( tKNN.m_eQuantization ) << "'"; } -knn::IndexSettings_t ReadKNNJson ( bson::Bson_c tRoot ) +void ReadKNNJson ( bson::Bson_c tRoot, knn::IndexSettings_t & tIS, knn::ModelSettings_t & tMS, CSphString & sKNNFrom ) { - knn::IndexSettings_t tRes; - tRes.m_iDims = (int) bson::Int ( tRoot.ChildByName ( "knn_dims" ) ); - Str2HNSWSimilarity ( bson::String ( tRoot.ChildByName ( "hnsw_similarity" ) ), tRes.m_eHNSWSimilarity ); - tRes.m_iHNSWM = (int) bson::Int ( tRoot.ChildByName ( "hnsw_m" ), tRes.m_iHNSWM ); - tRes.m_iHNSWEFConstruction = (int) bson::Int ( tRoot.ChildByName ( "hnsw_ef_construction" ), tRes.m_iHNSWEFConstruction ); - Str2Quantization ( bson::String ( tRoot.ChildByName ( "quantization" ) ), tRes.m_eQuantization ); - - return tRes; + tIS.m_iDims = (int) bson::Int ( tRoot.ChildByName ( "knn_dims" ) ); + Str2HNSWSimilarity ( bson::String ( tRoot.ChildByName ( "hnsw_similarity" ) ), tIS.m_eHNSWSimilarity ); + tIS.m_iHNSWM = (int) bson::Int ( tRoot.ChildByName ( "hnsw_m" ), tIS.m_iHNSWM ); + tIS.m_iHNSWEFConstruction = (int) bson::Int ( tRoot.ChildByName ( "hnsw_ef_construction" ), tIS.m_iHNSWEFConstruction ); + Str2Quantization ( bson::String ( tRoot.ChildByName ( "quantization" ) ), tIS.m_eQuantization ); + + tMS.m_sModelName = bson::String ( tRoot.ChildByName ( "model_name" ) ).cstr(); + tMS.m_sAPIKey = bson::String ( tRoot.ChildByName ( "api_key" ) ).cstr(); + tMS.m_sCachePath = bson::String ( tRoot.ChildByName ( "cache_path" ) ).cstr(); + tMS.m_bUseGPU = bson::Bool ( tRoot.ChildByName ( "use_gpu" ), tMS.m_bUseGPU ); + sKNNFrom = bson::String ( tRoot.ChildByName ( "from" ) ); } -void operator << ( JsonEscapedBuilder & tOut, const knn::IndexSettings_t & tSettings ) +void FormatKNNSettings ( JsonEscapedBuilder & tOut, const knn::IndexSettings_t & tIS, const knn::ModelSettings_t & tMS, const CSphString & sKNNFrom ) { auto _ = tOut.Object(); - knn::IndexSettings_t tDefault; + knn::IndexSettings_t tDefaultIS; tOut.NamedString ( "knn_type", "hnsw" ); - tOut.NamedVal ( "knn_dims", tSettings.m_iDims ); - tOut.NamedString ( "hnsw_similarity", HNSWSimilarity2Str ( tSettings.m_eHNSWSimilarity ) ); - tOut.NamedValNonDefault ( "hnsw_m", tSettings.m_iHNSWM, tDefault.m_iHNSWM ); - tOut.NamedValNonDefault ( "hnsw_ef_construction", tSettings.m_iHNSWEFConstruction, tDefault.m_iHNSWEFConstruction ); - tOut.NamedString ( "quantization", Quantization2Str ( tSettings.m_eQuantization ) ); + tOut.NamedVal ( "knn_dims", tIS.m_iDims ); + tOut.NamedString ( "hnsw_similarity", HNSWSimilarity2Str ( tIS.m_eHNSWSimilarity ) ); + tOut.NamedValNonDefault ( "hnsw_m", tIS.m_iHNSWM, tDefaultIS.m_iHNSWM ); + tOut.NamedValNonDefault ( "hnsw_ef_construction", tIS.m_iHNSWEFConstruction, tDefaultIS.m_iHNSWEFConstruction ); + tOut.NamedString ( "quantization", Quantization2Str ( tIS.m_eQuantization ) ); + + if ( !tMS.m_sModelName.empty() ) + { + tOut.NamedString ( "model_name", tMS.m_sModelName.c_str() ); + tOut.NamedString ( "from", sKNNFrom ); + tOut.NamedString ( "cache_path", tMS.m_sCachePath.c_str() ); + tOut.NamedString ( "api_key", tMS.m_sAPIKey.c_str() ); + tOut.NamedVal ( "use_gpu", tMS.m_bUseGPU ); + } } @@ -312,6 +390,16 @@ CSphString FormatKNNConfigStr ( const CSphVector & dAttrs ) tObj.AddInt ( "hnsw_m", i.m_iHNSWM ); tObj.AddInt ( "hnsw_ef_construction", i.m_iHNSWEFConstruction ); tObj.AddStr ( "quantization", Quantization2Str ( i.m_eQuantization ) ); + + if ( !i.m_sModelName.empty() ) + { + tObj.AddStr ( "model_name", i.m_sModelName.c_str() ); + tObj.AddStr ( "from", i.m_sFrom.cstr() ); + tObj.AddStr ( "cache_path", i.m_sCachePath.c_str() ); + tObj.AddStr ( "api_key", i.m_sAPIKey.c_str() ); + tObj.AddBool ( "use_gpu", i.m_bUseGPU ); + } + tArray.AddItem(tObj); } @@ -369,6 +457,16 @@ bool ParseKNNConfigStr ( const CSphString & sStr, CSphVector if ( !Str2Quantization ( sQuantization.cstr(), tParsed.m_eQuantization, &sError ) ) return false; } + + if ( !i.FetchStrItem ( tParsed.m_sModelName, "model_name", sError, true ) ) return false; + + if ( !tParsed.m_sModelName.empty() ) + { + if ( !i.FetchStrItem ( tParsed.m_sFrom, "from", sError, true ) ) return false; + if ( !i.FetchStrItem ( tParsed.m_sAPIKey, "api_key", sError, true ) ) return false; + if ( !i.FetchStrItem ( tParsed.m_sCachePath, "cache_path", sError, true ) ) return false; + if ( !i.FetchBoolItem ( tParsed.m_bUseGPU, "use_gpu", sError, true ) ) return false; + } } return true; diff --git a/src/knnmisc.h b/src/knnmisc.h index d753d9bd18..c26f4bb6c0 100644 --- a/src/knnmisc.h +++ b/src/knnmisc.h @@ -16,15 +16,38 @@ #include "indexsettings.h" #include "secondaryindex.h" +class TableEmbeddings_c +{ +public: + bool Load ( const CSphString & sAttr, const knn::ModelSettings_t & tSettings, CSphString & sError ); + knn::TextToEmbeddings_i * GetModel ( const CSphString & sAttr ) const; + +private: + SmallStringHash_T> m_hModels; +}; + +class EmbeddingsSrc_c +{ +public: + EmbeddingsSrc_c ( int iAttrs ); + + void Add ( int iAttr, CSphVector & dSrc ); + void Remove ( const CSphFixedVector & dRowMap ); + const VecTraits_T Get ( RowID_t tRowID, int iAttr ) const; + +private: + CSphVector>> m_dStored; +}; + const char * GetKnnDistAttrName(); ISphExpr * CreateExpr_KNNDist ( const CSphVector & dAnchor, const CSphColumnInfo & tAttr ); void NormalizeVec ( VecTraits_T & dData ); -void operator << ( JsonEscapedBuilder & tOut, const knn::IndexSettings_t & tSettings ); void AddKNNSettings ( StringBuilder_c & sRes, const CSphColumnInfo & tAttr ); -knn::IndexSettings_t ReadKNNJson ( bson::Bson_c tRoot ); +void ReadKNNJson ( bson::Bson_c tRoot, knn::IndexSettings_t & tIS, knn::ModelSettings_t & tMS, CSphString & sKNNFrom ); CSphString FormatKNNConfigStr ( const CSphVector & dAttrs ); bool ParseKNNConfigStr ( const CSphString & sStr, CSphVector & dParsed, CSphString & sError ); +void FormatKNNSettings ( JsonEscapedBuilder & tOut, const knn::IndexSettings_t & tIndexSettings, const knn::ModelSettings_t & tModelSettings, const CSphString & sKNNFrom ); bool Str2HNSWSimilarity ( const CSphString & sSimilarity, knn::HNSWSimilarity_e & eSimilarity, CSphString * pError = nullptr ); bool Str2Quantization ( const CSphString & sQuantization, knn::Quantization_e & eQuantization, CSphString * pError = nullptr ); diff --git a/src/queuecreator.cpp b/src/queuecreator.cpp index 60cc03b0f3..758bc7878b 100644 --- a/src/queuecreator.cpp +++ b/src/queuecreator.cpp @@ -1944,7 +1944,7 @@ void QueueCreator_c::ReplaceStaticStringsWithExprs ( CSphMatchComparatorState & CSphColumnInfo tRemapCol ( sAttrName.cstr(), SPH_ATTR_STRINGPTR ); tRemapCol.m_eStage = SPH_EVAL_PRESORT; - tRemapCol.m_pExpr = CreateExpr_GetColumnarString ( sAttrName, tAttr.m_uAttrFlags & CSphColumnInfo::ATTR_STORED ); + tRemapCol.m_pExpr = CreateExpr_GetColumnarString ( sAttrName, tAttr.IsStored() ); tSorterSchema.AddAttr ( tRemapCol, true ); iRemap = tSorterSchema.GetAttrIndex ( sAttrName.cstr() ); diff --git a/src/schema/columninfo.cpp b/src/schema/columninfo.cpp index 1b9e297205..085278550b 100644 --- a/src/schema/columninfo.cpp +++ b/src/schema/columninfo.cpp @@ -83,6 +83,12 @@ bool CSphColumnInfo::IsIndexedSI() const } +bool CSphColumnInfo::IsStored() const +{ + return m_uAttrFlags & CSphColumnInfo::ATTR_STORED; +} + + CSphString sphDumpAttr ( const CSphColumnInfo & tAttr ) { CSphString sRes; diff --git a/src/schema/columninfo.h b/src/schema/columninfo.h index 2b44824e33..4a85b9a721 100644 --- a/src/schema/columninfo.h +++ b/src/schema/columninfo.h @@ -22,9 +22,10 @@ class CSphWriter; class CSphReader; -struct NamedKNNSettings_t : public knn::IndexSettings_t +struct NamedKNNSettings_t : public knn::IndexSettings_t, public knn::ModelSettings_t { - CSphString m_sName; + CSphString m_sName; + CSphString m_sFrom; }; /// source column info @@ -72,7 +73,9 @@ struct CSphColumnInfo DWORD m_uAttrFlags = ATTR_NONE; ///< attribute storage spec AttrEngine_e m_eEngine = AttrEngine_e::DEFAULT; ///< used together with per-table engine specs to determine attribute storage - knn::IndexSettings_t m_tKNN; ///< knn index settings + knn::IndexSettings_t m_tKNN; ///< knn index settings + knn::ModelSettings_t m_tKNNModel; ///< knn model settings + CSphString m_sKNNFrom; ///< fields/attrs used by the model WORD m_uNext = 0xFFFF; ///< next in linked list for hash in CSphSchema @@ -92,6 +95,7 @@ struct CSphColumnInfo bool IsIndexedKNN() const; bool IsJoined() const; bool IsIndexedSI() const; + bool IsStored() const; }; diff --git a/src/schema/schema.cpp b/src/schema/schema.cpp index 1e60c418b5..9aa11d73fb 100644 --- a/src/schema/schema.cpp +++ b/src/schema/schema.cpp @@ -570,7 +570,10 @@ void CSphSchema::SetupKNNFlags ( const CSphSourceSettings & tSettings ) continue; int iId = hKNN[tAttr.m_sName]; - tAttr.m_tKNN = tSettings.m_dKNN[iId]; + const auto & tKNN = tSettings.m_dKNN[iId]; + tAttr.m_tKNN = tKNN; + tAttr.m_tKNNModel = tKNN; + tAttr.m_sKNNFrom = tKNN.m_sFrom; } } @@ -662,7 +665,7 @@ bool CSphSchema::HasStoredFields() const bool CSphSchema::HasStoredAttrs() const { - return m_dAttrs.any_of ( []( const CSphColumnInfo & tAttr ){ return tAttr.m_uAttrFlags & CSphColumnInfo::ATTR_STORED; } ); + return m_dAttrs.any_of ( []( const CSphColumnInfo & tAttr ){ return tAttr.IsStored(); } ); } @@ -698,7 +701,7 @@ bool CSphSchema::IsFieldStored ( int iField ) const bool CSphSchema::IsAttrStored ( int iAttr ) const { - return !!( m_dAttrs[iAttr].m_uAttrFlags & CSphColumnInfo::ATTR_STORED ); + return m_dAttrs[iAttr].IsStored(); } /////////////////////////////////////////////////////////////////////////////// diff --git a/src/searchd.cpp b/src/searchd.cpp index d1cbc0ebed..afc8c9e9a0 100644 --- a/src/searchd.cpp +++ b/src/searchd.cpp @@ -5962,7 +5962,7 @@ static CSphString DescribeAttributeProperties ( const CSphColumnInfo & tAttr ) if ( tAttr.IsIndexedSI() ) sProps << "indexed"; - if ( tAttr.m_uAttrFlags & CSphColumnInfo::ATTR_STORED ) + if ( tAttr.IsStored() ) sProps << "fast_fetch"; if ( tAttr.IsColumnar() && tAttr.m_eAttrType==SPH_ATTR_STRING && !(tAttr.m_uAttrFlags & CSphColumnInfo::ATTR_COLUMNAR_HASHES) ) @@ -6350,6 +6350,14 @@ static bool CheckCreateTable ( const SqlStmt_t & tStmt, CSphString & sError ) if ( !CheckAttrs ( tStmt.m_tCreateTable.m_dFields, []( const CSphColumnInfo & tAttr ) { return tAttr.m_sName; }, sError ) ) return false; + for ( auto & i : tStmt.m_tCreateTable.m_dAttrs ) + if ( i.m_bKNN && !i.m_tKNNModel.m_sModelName.empty() && !IsKNNEmbeddingsLibLoaded() ) + { + sError.SetSprintf ( "model_name specified for '%s', but embeddings library is not loded", i.m_tAttr.m_sName.cstr() ); + return false; + } + + // cross-checks attrs and fields for ( const auto & i : tStmt.m_tCreateTable.m_dAttrs ) for ( const auto & j : tStmt.m_tCreateTable.m_dFields ) @@ -14125,10 +14133,15 @@ static void InitBanner() if ( szKNNVer ) sKNN.SetSprintf ( " (knn %s)", szKNNVer ); - g_sBannerVersion.SetSprintf ( "%s%s%s%s", szMANTICORE_NAME, sColumnar.cstr(), sSi.cstr(), sKNN.cstr() ); + const char * szKNNEmbVer = GetKNNEmbeddingsVersionStr(); + CSphString sKNNEmb = ""; + if ( szKNNEmbVer ) + sKNNEmb.SetSprintf ( " (embeddings %s)", szKNNEmbVer ); + + g_sBannerVersion.SetSprintf ( "%s%s%s%s%s", szMANTICORE_NAME, sColumnar.cstr(), sSi.cstr(), sKNN.cstr(), sKNNEmb.cstr() ); g_sBanner.SetSprintf ( "%s%s", g_sBannerVersion.cstr(), szMANTICORE_BANNER_TEXT ); - g_sMySQLVersion.SetSprintf ( "%s%s%s%s", szMANTICORE_VERSION, sColumnar.cstr(), sSi.cstr(), sKNN.cstr() ); - g_sStatusVersion.SetSprintf ( "%s%s%s%s", szMANTICORE_VERSION, sColumnar.cstr(), sSi.cstr(), sKNN.cstr() ); + g_sMySQLVersion.SetSprintf ( "%s%s%s%s%s", szMANTICORE_VERSION, sColumnar.cstr(), sSi.cstr(), sKNN.cstr(), sKNNEmb.cstr() ); + g_sStatusVersion.SetSprintf ( "%s%s%s%s%s", szMANTICORE_VERSION, sColumnar.cstr(), sSi.cstr(), sKNN.cstr(), sKNNEmb.cstr() ); } diff --git a/src/searchdddl.cpp b/src/searchdddl.cpp index af3e7f3c05..21aea41daa 100644 --- a/src/searchdddl.cpp +++ b/src/searchdddl.cpp @@ -44,10 +44,16 @@ class DdlParser_c : public SqlParserTraits_c knn::Quantization_e m_eQuantization = knn::Quantization_e::NONE; bool m_bKNNDimsSpecified = false; bool m_bHNSWSimilaritySpecified = false; + CSphString m_sModelName; + CSphString m_sAPIKey; + CSphString m_sCachePath; + CSphString m_sFrom; + bool m_bUseGPU = false; void Reset() { *this = ItemOptions_t(); } DWORD ToFlags() const; knn::IndexSettings_t ToKNN() const; + knn::ModelSettings_t ToKNNModel() const; void CopyOptionsTo ( CreateTableAttr_t & tAttr ) const; }; @@ -69,6 +75,11 @@ class DdlParser_c : public SqlParserTraits_c bool AddItemOptionHNSWSimilarity ( const SqlNode_t & tOption ); bool AddItemOptionHNSWM ( const SqlNode_t & tOption ); bool AddItemOptionHNSWEfConstruction ( const SqlNode_t & tOption ); + bool AddItemOptionModelName ( const SqlNode_t & tOption ); + bool AddItemOptionFrom ( const SqlNode_t & tOption ); + bool AddItemOptionCachePath ( const SqlNode_t & tOption ); + bool AddItemOptionAPIKey ( const SqlNode_t & tOption ); + bool AddItemOptionUseGPU ( const SqlNode_t & tOption ); bool AddItemOptionQuantization ( const SqlNode_t & tOption ); void AddCreateTableOption ( const SqlNode_t & tName, const SqlNode_t & tValue ); @@ -157,6 +168,19 @@ knn::IndexSettings_t DdlParser_c::ItemOptions_t::ToKNN() const } +knn::ModelSettings_t DdlParser_c::ItemOptions_t::ToKNNModel() const +{ + knn::ModelSettings_t tModel; + + tModel.m_sModelName = m_sModelName.scstr(); + tModel.m_sAPIKey = m_sAPIKey.scstr(); + tModel.m_sCachePath = m_sCachePath.scstr(); + tModel.m_bUseGPU = m_bUseGPU; + + return tModel; +} + + void DdlParser_c::ItemOptions_t::CopyOptionsTo ( CreateTableAttr_t & tAttr ) const { tAttr.m_tAttr.m_eEngine = m_eEngine; @@ -232,10 +256,19 @@ bool DdlParser_c::CheckFieldFlags ( ESphAttr eAttrType, int iFlags, const CSphSt } else if ( eAttrType==SPH_ATTR_FLOAT_VECTOR ) { - if ( !tOpts.m_sKNNType.IsEmpty() && ( !tOpts.m_bKNNDimsSpecified || !tOpts.m_bHNSWSimilaritySpecified ) ) + if ( !tOpts.m_sKNNType.IsEmpty() ) { - sError = "knn_dims and hnsw_similarity are required if knn_type='hnsw'"; - return false; + if ( ( !tOpts.m_bKNNDimsSpecified || !tOpts.m_bHNSWSimilaritySpecified ) && tOpts.m_sModelName.IsEmpty() ) + { + sError = "knn_dims and hnsw_similarity are required if knn_type='hnsw'"; + return false; + } + + if ( tOpts.m_bKNNDimsSpecified && !tOpts.m_sModelName.IsEmpty() ) + { + sError = "knn_dims can't be used together with model_name"; + return false; + } } } else @@ -306,11 +339,13 @@ bool DdlParser_c::AddCreateTableCol ( const SqlNode_t & tName, const SqlNode_t & if ( eAttrType!=SPH_ATTR_STRING ) { CreateTableAttr_t & tAttr = m_pStmt->m_tCreateTable.m_dAttrs.Add(); - tAttr.m_tAttr.m_sName = sName; - tAttr.m_tAttr.m_eAttrType = eAttrType; + tAttr.m_tAttr.m_sName = sName; + tAttr.m_tAttr.m_eAttrType = eAttrType; tOpts.CopyOptionsTo(tAttr); - tAttr.m_bKNN = !tOpts.m_sKNNType.IsEmpty(); - tAttr.m_tKNN = tOpts.ToKNN(); + tAttr.m_bKNN = !tOpts.m_sKNNType.IsEmpty(); + tAttr.m_tKNN = tOpts.ToKNN(); + tAttr.m_tKNNModel = tOpts.ToKNNModel(); + tAttr.m_sKNNFrom = tOpts.m_sFrom; return true; } @@ -457,6 +492,42 @@ bool DdlParser_c::AddItemOptionHNSWEfConstruction ( const SqlNode_t & tOption ) } +bool DdlParser_c::AddItemOptionModelName ( const SqlNode_t & tOption ) +{ + m_tItemOptions.m_sModelName = ToStringUnescape(tOption); + return true; +} + + +bool DdlParser_c::AddItemOptionFrom ( const SqlNode_t & tOption ) +{ + m_tItemOptions.m_sFrom= ToStringUnescape(tOption); + return true; +} + + +bool DdlParser_c::AddItemOptionAPIKey ( const SqlNode_t & tOption ) +{ + m_tItemOptions.m_sAPIKey = ToStringUnescape(tOption); + return true; +} + + +bool DdlParser_c::AddItemOptionCachePath ( const SqlNode_t & tOption ) +{ + m_tItemOptions.m_sCachePath = ToStringUnescape(tOption); + return true; +} + + +bool DdlParser_c::AddItemOptionUseGPU ( const SqlNode_t & tOption ) +{ + CSphString sValue = ToStringUnescape(tOption); + m_tItemOptions.m_bUseGPU = !!strtoull ( sValue.cstr(), NULL, 10 ); + return true; +} + + bool DdlParser_c::AddItemOptionQuantization ( const SqlNode_t & tOption ) { return Str2Quantization ( ToStringUnescape(tOption), m_tItemOptions.m_eQuantization, &m_sError ); diff --git a/src/sphinx.cpp b/src/sphinx.cpp index fbff50dec8..ff1fb217da 100644 --- a/src/sphinx.cpp +++ b/src/sphinx.cpp @@ -3932,7 +3932,7 @@ static void ReadSchemaColumnJson ( bson::Bson_c tNode, CSphColumnInfo & tCol ) NodeHandle_t tKNN = tNode.ChildByName ("knn"); if ( tKNN!=nullnode ) - tCol.m_tKNN = ReadKNNJson(tKNN); + ReadKNNJson ( tKNN, tCol.m_tKNN, tCol.m_tKNNModel, tCol.m_sKNNFrom ); } @@ -4027,7 +4027,10 @@ void DumpAttrToJson ( JsonEscapedBuilder& tOut, const CSphColumnInfo& tCol ) tOut.NamedVal ( "locator", tCol.m_tLocator ); if ( tCol.IsIndexedKNN() ) - tOut.NamedVal ( "knn", tCol.m_tKNN ); + { + tOut.Named ( "knn" ); + FormatKNNSettings ( tOut, tCol.m_tKNN, tCol.m_tKNNModel, tCol.m_sKNNFrom ); + } } } // namespace diff --git a/src/sphinxexpr.cpp b/src/sphinxexpr.cpp index 3ed429a6d7..ceb7908f46 100644 --- a/src/sphinxexpr.cpp +++ b/src/sphinxexpr.cpp @@ -5975,7 +5975,7 @@ ISphExpr * ExprParser_t::CreateExistNode ( const ExprNode_t & tNode ) } bool bColumnar = tCol.IsColumnar(); - bool bStored = tCol.m_uAttrFlags & CSphColumnInfo::ATTR_STORED; + bool bStored = tCol.IsStored(); const CSphAttrLocator & tLoc = tCol.m_tLocator; if ( tNode.m_eRetType==SPH_ATTR_FLOAT ) { @@ -6637,28 +6637,28 @@ ISphExpr * ExprParser_t::CreateFieldNode ( int iField ) ISphExpr * ExprParser_t::CreateColumnarIntNode ( int iAttr, ESphAttr eAttrType ) { const CSphColumnInfo & tAttr = m_pSchema->GetAttr(iAttr); - return CreateExpr_GetColumnarInt ( tAttr.m_sName, tAttr.m_uAttrFlags & CSphColumnInfo::ATTR_STORED ); + return CreateExpr_GetColumnarInt ( tAttr.m_sName, tAttr.IsStored() ); } ISphExpr * ExprParser_t::CreateColumnarFloatNode ( int iAttr ) { const CSphColumnInfo & tAttr = m_pSchema->GetAttr(iAttr); - return CreateExpr_GetColumnarFloat ( tAttr.m_sName, tAttr.m_uAttrFlags & CSphColumnInfo::ATTR_STORED ); + return CreateExpr_GetColumnarFloat ( tAttr.m_sName, tAttr.IsStored() ); } ISphExpr * ExprParser_t::CreateColumnarStringNode ( int iAttr ) { const CSphColumnInfo & tAttr = m_pSchema->GetAttr(iAttr); - return CreateExpr_GetColumnarString ( tAttr.m_sName, tAttr.m_uAttrFlags & CSphColumnInfo::ATTR_STORED ); + return CreateExpr_GetColumnarString ( tAttr.m_sName, tAttr.IsStored() ); } ISphExpr * ExprParser_t::CreateColumnarMvaNode ( int iAttr ) { const CSphColumnInfo & tAttr = m_pSchema->GetAttr(iAttr); - return CreateExpr_GetColumnarMva ( tAttr.m_sName, tAttr.m_uAttrFlags & CSphColumnInfo::ATTR_STORED ); + return CreateExpr_GetColumnarMva ( tAttr.m_sName, tAttr.IsStored() ); } ////////////////////////////////////////////////////////////////////////// diff --git a/src/sphinxjson.cpp b/src/sphinxjson.cpp index c45892b20c..31256ee76e 100644 --- a/src/sphinxjson.cpp +++ b/src/sphinxjson.cpp @@ -1969,6 +1969,15 @@ bool JsonObj_c::FetchStrItem ( CSphString & sValue, const char * szName, CSphStr } +bool JsonObj_c::FetchStrItem ( std::string & sValue, const char * szName, CSphString & sError, bool bIgnoreMissing ) const +{ + CSphString sTmp; + bool bRes = FetchStrItem ( sTmp, szName, sError, bIgnoreMissing ); + sValue = sTmp.cstr(); + return bRes; +} + + bool JsonObj_c::HasItem ( const char * szItem ) const { assert ( m_pRoot ); diff --git a/src/sphinxjson.h b/src/sphinxjson.h index f33d5d44da..b582ac6e4d 100644 --- a/src/sphinxjson.h +++ b/src/sphinxjson.h @@ -498,6 +498,7 @@ class JsonObj_c bool FetchBoolItem ( bool & bValue, const char * szName, CSphString & sError, bool bIgnoreMissing=false ) const; bool FetchFltItem ( float & fValue, const char * szName, CSphString & sError, bool bIgnoreMissing=false ) const; bool FetchStrItem ( CSphString & sValue, const char * szName, CSphString & sError, bool bIgnoreMissing=false ) const; + bool FetchStrItem ( std::string & sValue, const char * szName, CSphString & sError, bool bIgnoreMissing=false ) const; bool HasItem ( const char * szName ) const; static JsonObj_c CreateInt ( int64_t iInt ); diff --git a/src/sphinxrt.cpp b/src/sphinxrt.cpp index 2a50ff9d1c..eb32cd3d27 100644 --- a/src/sphinxrt.cpp +++ b/src/sphinxrt.cpp @@ -41,6 +41,7 @@ #include "tracer.h" #include "pseudosharding.h" #include "knnmisc.h" +#include "knnlib.h" #include "jsonsi.h" #include "std/sys.h" #include "dict/infix/infix_builder.h" @@ -1521,6 +1522,9 @@ class RtIndex_c final : public RtIndex_i, public ISphNoncopyable, public ISphWor mutable int m_iTrackFailedRamActions; int m_iAlterGeneration = 0; // increased every time index altered + std::unique_ptr m_pEmbeddings; + CSphVector m_dAttrsWithModels; + bool BindAccum ( RtAccum_t * pAccExt, CSphString* pError = nullptr ) final; int CompareWords ( const RtWord_t * pWord1, const RtWord_t * pWord2 ) const; @@ -1590,6 +1594,8 @@ class RtIndex_c final : public RtIndex_i, public ISphNoncopyable, public ISphWor void StopMergeSegmentsWorker(); bool NeedStoreWordID () const override; int64_t GetMemLimit() const final { return m_iRtMemLimit; } + + bool LoadEmbeddingModels ( CSphString & sError ); bool VerifyKNN ( InsertDocData_c & tDoc, CSphString & sError ) const; template @@ -2048,7 +2054,15 @@ bool RtIndex_c::VerifyKNN ( InsertDocData_c & tDoc, CSphString & sError ) const if ( tAttr.m_eAttrType!=SPH_ATTR_FLOAT_VECTOR || !tAttr.IsIndexedKNN() ) continue; - if ( !bDefault && iNumValues!=tAttr.m_tKNN.m_iDims ) + if ( m_dAttrsWithModels[i].m_pModel ) + { + if ( iNumValues!=0 ) + { + sError.SetSprintf ( "attribute '%s' has model_name=%s specified, but vector contents with %d values is provided", tAttr.m_sName.cstr(), tAttr.m_tKNNModel.m_sModelName.c_str(), iNumValues ); + return false; + } + } + else if ( !bDefault && iNumValues!=tAttr.m_tKNN.m_iDims ) { sError.SetSprintf ( "KNN error: data has %d values, index '%s' needs %d values", iNumValues, tAttr.m_sName.cstr(), tAttr.m_tKNN.m_iDims ); return false; @@ -2156,6 +2170,9 @@ bool RtIndex_c::AddDocument ( InsertDocData_c & tDoc, bool bReplace, const CSphS ISphHits * pHits = tSrc.IterateHits ( sError ); pAcc->GrabLastWarning ( sWarning ); + if ( m_pEmbeddings ) + pAcc->FetchEmbeddingsSrc ( tDoc, m_dAttrsWithModels ); + if ( !VerifyKNN ( tDoc, sError ) ) return false; @@ -3022,6 +3039,15 @@ bool RtIndex_c::Commit ( int * pDeleted, RtAccum_t * pAcc, CSphString * pError ) pAcc->Sort(); CleanupHitDuplicates ( pAcc->m_dAccum ); + CSphString sError; + if ( !pAcc->FetchEmbeddings ( m_pEmbeddings.get(), m_dAttrsWithModels, sError ) ) + { + if ( pError ) + *pError = sError; + + return false; + } + CSphString sCreateError; RtSegmentRefPtf_t pNewSeg { CreateSegment ( pAcc, m_iWordsCheckpoint, m_tSettings.m_eHitless, m_dHitlessWords, sCreateError ) }; if ( !pNewSeg && !sCreateError.IsEmpty() ) @@ -3046,7 +3072,6 @@ bool RtIndex_c::Commit ( int * pDeleted, RtAccum_t * pAcc, CSphString * pError ) // now on to the stuff that needs locking and recovery int iKilled = 0; - CSphString sError; if ( !CommitReplayable ( pNewSeg, pAcc->m_dAccumKlist, pAcc->m_iAccumBytes, iKilled, sError ) ) { if ( pError ) @@ -4979,6 +5004,83 @@ bool RtIndex_c::PreallocDiskChunks ( FilenameBuilder_i * pFilenameBuilder, StrVe } +static bool ParseKNNFrom ( AttrWithModel_t & tAttrWithModel, const CSphString & sFrom, const ISphSchema & tSchema, CSphString & sError ) +{ + StrVec_t dFrom; + sphSplit ( dFrom, sFrom.cstr(), " \t," ); + + for ( const auto & i : dFrom ) + { + int iAttrId = tSchema.GetAttrIndex ( i.cstr() ); + int iFieldId = tSchema.GetFieldIndex ( i.cstr() ); + + if ( iFieldId==-1 && iAttrId==-1 ) + { + sError.SetSprintf ( "embedding source '%s' not found", i.cstr() ); + return false; + } + + if ( iAttrId!=-1 && tSchema.GetAttr(iAttrId).m_eAttrType!=SPH_ATTR_STRING ) + { + sError.SetSprintf ( "embedding source attribute '%s' is not a string", i.cstr() ); + return false; + } + + tAttrWithModel.m_dFrom.Add ( { iFieldId==-1 ? iAttrId : iFieldId, iFieldId!=-1 } ); + } + + return true; +} + + +bool RtIndex_c::LoadEmbeddingModels ( CSphString & sError ) +{ + if ( m_pEmbeddings ) + return true; + + bool bHaveModels = false; + for ( int i = 0 ; i < m_tSchema.GetAttrsCount(); i++ ) + bHaveModels |= !m_tSchema.GetAttr(i).m_tKNNModel.m_sModelName.empty(); + + if ( !bHaveModels ) + return true; + + m_dAttrsWithModels.Resize ( m_tSchema.GetAttrsCount() ); + + m_pEmbeddings = std::make_unique(); + for ( int i = 0; i < m_tSchema.GetAttrsCount(); i++ ) + { + const auto & tAttr = m_tSchema.GetAttr(i); + m_dAttrsWithModels[i].m_pModel = nullptr; + + if ( tAttr.m_tKNNModel.m_sModelName.empty() ) + continue; + + if ( !ParseKNNFrom ( m_dAttrsWithModels[i], tAttr.m_sKNNFrom, m_tSchema, sError ) ) + { + m_pEmbeddings.reset(); + return false; + } + + if ( !m_pEmbeddings->Load ( tAttr.m_sName, tAttr.m_tKNNModel, sError ) ) + { + m_pEmbeddings.reset(); + return false; + } + + auto pModel = m_pEmbeddings->GetModel ( tAttr.m_sName ); + assert(pModel); + + m_dAttrsWithModels[i].m_pModel = pModel; + + // fixme! modifying the schema + const_cast(tAttr).m_tKNN.m_iDims = pModel->GetDims(); + } + + return true; +} + + bool RtIndex_c::Prealloc ( bool bStripPath, FilenameBuilder_i * pFilenameBuilder, StrVec_t & dWarnings ) { MEMORY ( MEM_INDEX_RT ); @@ -5013,7 +5115,6 @@ bool RtIndex_c::Prealloc ( bool bStripPath, FilenameBuilder_i * pFilenameBuilder if ( !LoadMeta ( pFilenameBuilder, bStripPath, uVersion, bRebuildInfixes, dWarnings ) ) return false; - CSphString sMutableFile = GetFilename ( SPH_EXT_SETTINGS ); m_tMutableSettings.m_iMemLimit = m_iRtMemLimit; // to avoid overriding value from meta by default value, if no settings provided if ( !m_tMutableSettings.Load ( sMutableFile.cstr(), GetName() ) ) @@ -5030,6 +5131,12 @@ bool RtIndex_c::Prealloc ( bool bStripPath, FilenameBuilder_i * pFilenameBuilder return false; } + if ( m_tSchema.HasKNNAttrs() && !IsKNNLibLoaded() ) + { + m_sLastError.SetSprintf ( "failed to load table with knn attributes without knn library" ); + return false; + } + if ( m_bDebugCheck ) { // load ram chunk @@ -5051,6 +5158,9 @@ bool RtIndex_c::Prealloc ( bool bStripPath, FilenameBuilder_i * pFilenameBuilder ARRAY_FOREACH ( i, m_dFieldLens ) m_dFieldLens[i] = m_dFieldLensDisk[i] + m_dFieldLensRam[i]; + if ( !LoadEmbeddingModels(m_sLastError) ) + return false; + // set up values for on timer save m_iSavedTID = m_iTID; m_tmSaved = sphMicroTimer(); diff --git a/src/sphinxrt.h b/src/sphinxrt.h index 9f13ddf85d..111636a2a0 100644 --- a/src/sphinxrt.h +++ b/src/sphinxrt.h @@ -52,6 +52,7 @@ class InsertDocData_c const int64_t * GetMVA ( int iMVA ) const { return m_dMvas.Begin()+iMVA; } void FixParsedMVAs ( const CSphVector & dParsed, int iCount ); static std::pair ReadMVALength ( const int64_t * & pMVA ); + void SwapMVAs ( InsertDocData_c & tSrc ) { Swap ( m_dMvas, tSrc.m_dMvas ); } private: static const uint64_t DEFAULT_FLAG = 1ULL << 63; diff --git a/src/std/env.cpp b/src/std/env.cpp index b4d401e09b..94a7fa0a3e 100644 --- a/src/std/env.cpp +++ b/src/std/env.cpp @@ -114,6 +114,7 @@ CSphString GetSecondaryFullpath() return sResult; } + CSphString GetKNNFullpath() { CSphString sResult; @@ -127,6 +128,19 @@ CSphString GetKNNFullpath() } +CSphString GetKNNEmbeddingsFullpath() +{ + CSphString sResult; + const char* szEnv = getenv ( "LIB_MANTICORE_KNN_EMBEDDINGS" ); + if ( szEnv ) + sResult = szEnv; + else + sResult.SetSprintf ( "%s/" LIB_MANTICORE_KNN_EMBEDDINGS, GET_MANTICORE_MODULES() ); + + return sResult; +} + + static CSphString GetSegmentationDataPath ( const char * szEnvVar, const char * szDir, const char * szMask ) { const char* szEnv = getenv(szEnvVar); diff --git a/src/std/env.h b/src/std/env.h index a57d2a54a7..dc77c05672 100644 --- a/src/std/env.h +++ b/src/std/env.h @@ -34,6 +34,7 @@ CSphString GET_GALERA_FULLPATH(); CSphString GetColumnarFullpath(); CSphString GetSecondaryFullpath(); CSphString GetKNNFullpath(); +CSphString GetKNNEmbeddingsFullpath(); // return value of asked ENV, or default. // note, default determines the type which to return