milvus-io--milvus/internal/util/initcore/query_node.go

// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
//	http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// 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.
package initcore

/*
#cgo pkg-config: milvus_core

#include <stdlib.h>
#include <stdint.h>
#include "common/init_c.h"
#include "segcore/segcore_init_c.h"
#include "storage/storage_c.h"
#include "segcore/arrow_fs_c.h"
#include "common/type_c.h"
#include "segcore/collection_c.h"
#include "segcore/segment_c.h"
#include "exec/expression/function/init_c.h"
*/
import "C"

import (
	"context"
	"path"
	"sync"
	"unsafe"

	"github.com/milvus-io/milvus/internal/util/pathutil"
	"github.com/milvus-io/milvus/pkg/v3/mlog"
	"github.com/milvus-io/milvus/pkg/v3/util/hardware"
	"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
	"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)

var initQueryNodeOnce sync.Once

// InitQueryNode initializes query node once.
func InitQueryNode(ctx context.Context) error {
	var err error
	initQueryNodeOnce.Do(func() {
		err = doInitQueryNodeOnce(ctx)
	})
	return err
}

// doInitQueryNodeOnce initializes query node once.
func doInitQueryNodeOnce(ctx context.Context) error {
	nodeID := paramtable.GetNodeID()

	cGlogConf := C.CString(path.Join(paramtable.GetBaseTable().GetConfigDir(), paramtable.DefaultGlogConf))
	C.SegcoreInit(cGlogConf)
	C.free(unsafe.Pointer(cGlogConf))

	C.LogOpenSSLFIPSStatus()

	// update log level based on current setup
	UpdateLogLevel(paramtable.Get().LogCfg.Level.GetValue())

	// override segcore chunk size
	cChunkRows := C.int64_t(paramtable.Get().QueryNodeCfg.ChunkRows.GetAsInt64())
	C.SegcoreSetChunkRows(cChunkRows)

	cMaxGroupByGroups := C.int64_t(paramtable.Get().CommonCfg.GroupByMaxGroups.GetAsInt64())
	C.SegcoreSetMaxGroupByGroups(cMaxGroupByGroups)

	visibilityEnabled := paramtable.Get().CommonCfg.VisibilityFilterEnabled.GetAsBool()
	bloomEnabled := paramtable.Get().CommonCfg.BloomFilterEnabled.GetAsBool()
	C.SegcoreSetVisibilityFilterEnabled(C.bool(visibilityEnabled))
	if !visibilityEnabled && bloomEnabled {
		mlog.Warn(ctx, "visibilityFilterEnabled=false with bloomFilterEnabled=true: deletes are forwarded via bloom filter but never applied — consider disabling bloom filter to save memory")
	}

	SyncPreferFieldDataWhenIndexHasRawData(ctx, paramtable.Get())
	SyncEnableGrowingSourceFlush(ctx, paramtable.Get())

	cKnowhereThreadPoolSize := C.uint32_t(paramtable.Get().QueryNodeCfg.KnowhereThreadPoolSize.GetAsUint32())
	C.SegcoreSetKnowhereSearchThreadPoolNum(cKnowhereThreadPoolSize)

	cKnowhereFetchThreadPoolSize := C.uint32_t(paramtable.Get().QueryNodeCfg.KnowhereFetchThreadPoolSize.GetAsUint32())
	C.SegcoreSetKnowhereFetchThreadPoolNum(cKnowhereFetchThreadPoolSize)

	// override segcore SIMD type
	cSimdType := C.CString(paramtable.Get().CommonCfg.SimdType.GetValue())
	C.SegcoreSetSimdType(cSimdType)
	C.free(unsafe.Pointer(cSimdType))

	enableKnowhereScoreConsistency := paramtable.Get().QueryNodeCfg.KnowhereScoreConsistency.GetAsBool()
	if enableKnowhereScoreConsistency {
		C.SegcoreEnableKnowhereScoreConsistency()
	}

	// override segcore index slice size
	cIndexSliceSize := C.int64_t(paramtable.Get().CommonCfg.IndexSliceSize.GetAsInt64())
	C.SetIndexSliceSize(cIndexSliceSize)
	cLoadTransientBudgetBytes := C.int64_t(paramtable.Get().CommonCfg.LoadTransientBudgetBytes.GetAsInt64())
	C.SetLoadTransientBudgetBytes(cLoadTransientBudgetBytes)

	// set up thread pool for different priorities
	cHighPriorityThreadCoreCoefficient := C.float(paramtable.Get().CommonCfg.HighPriorityThreadCoreCoefficient.GetAsFloat())
	C.SetHighPriorityThreadCoreCoefficient(cHighPriorityThreadCoreCoefficient)
	cMiddlePriorityThreadCoreCoefficient := C.float(paramtable.Get().CommonCfg.MiddlePriorityThreadCoreCoefficient.GetAsFloat())
	C.SetMiddlePriorityThreadCoreCoefficient(cMiddlePriorityThreadCoreCoefficient)
	cLowPriorityThreadCoreCoefficient := C.float(paramtable.Get().CommonCfg.LowPriorityThreadCoreCoefficient.GetAsFloat())
	C.SetLowPriorityThreadCoreCoefficient(cLowPriorityThreadCoreCoefficient)
	cThreadPoolMaxThreadsSize := C.int(paramtable.Get().CommonCfg.ThreadPoolMaxThreadsSize.GetAsInt())
	C.SetThreadPoolMaxThreadsSize(cThreadPoolMaxThreadsSize)

	cCPUNum := C.int(hardware.GetCPUNum())
	C.InitCpuNum(cCPUNum)

	knowhereBuildPoolSize := uint32(float32(paramtable.Get().QueryNodeCfg.InterimIndexBuildParallelRate.GetAsFloat()) * float32(hardware.GetCPUNum()))
	if knowhereBuildPoolSize < uint32(1) {
		knowhereBuildPoolSize = uint32(1)
	}
	mlog.Info(ctx, "set up knowhere build pool size", mlog.Uint32("pool_size", knowhereBuildPoolSize))
	cKnowhereBuildPoolSize := C.uint32_t(knowhereBuildPoolSize)
	C.SegcoreSetKnowhereBuildThreadPoolNum(cKnowhereBuildPoolSize)

	cExprBatchSize := C.int64_t(paramtable.Get().QueryNodeCfg.ExprEvalBatchSize.GetAsInt64())
	C.SetDefaultExprEvalBatchSize(cExprBatchSize)

	cDeleteDumpBatchSize := C.int64_t(paramtable.Get().QueryNodeCfg.DeleteDumpBatchSize.GetAsInt64())
	C.SetDefaultDeleteDumpBatchSize(cDeleteDumpBatchSize)

	cEnableLatestDeleteSnapshotOptimization := C.bool(paramtable.Get().QueryNodeCfg.EnableLatestDeleteSnapshotOptimization.GetAsBool())
	C.SetEnableLatestDeleteSnapshotOptimization(cEnableLatestDeleteSnapshotOptimization)

	cOptimizeExprEnabled := C.bool(paramtable.Get().CommonCfg.EnabledOptimizeExpr.GetAsBool())
	C.SetDefaultOptimizeExprEnable(cOptimizeExprEnabled)

	cDriverPrefetchEnabled := C.bool(paramtable.Get().CommonCfg.EnableDriverPrefetch.GetAsBool())
	C.SetDefaultDriverPrefetchEnable(cDriverPrefetchEnabled)

	cJSONKeyStatsEnabled := C.bool(paramtable.Get().CommonCfg.EnabledJSONKeyStats.GetAsBool())
	C.SetDefaultJSONKeyStatsEnable(cJSONKeyStatsEnabled)

	cGrowingJSONKeyStatsEnabled := C.bool(paramtable.Get().CommonCfg.EnabledGrowingSegmentJSONKeyStats.GetAsBool())
	C.SetDefaultGrowingJSONKeyStatsEnable(cGrowingJSONKeyStatsEnabled)

	if paramtable.GetRole() != typeutil.StreamingNodeRole {
		cGpuMemoryPoolInitSize := C.uint32_t(paramtable.Get().GpuConfig.InitSize.GetAsUint32())
		cGpuMemoryPoolMaxSize := C.uint32_t(paramtable.Get().GpuConfig.MaxSize.GetAsUint32())
		C.SegcoreSetKnowhereGpuMemoryPoolSize(cGpuMemoryPoolInitSize, cGpuMemoryPoolMaxSize)
	}

	cEnableConfigParamTypeCheck := C.bool(paramtable.Get().CommonCfg.EnableConfigParamTypeCheck.GetAsBool())
	C.SetDefaultConfigParamTypeCheck(cEnableConfigParamTypeCheck)

	cExprResCacheEnabled := C.bool(paramtable.Get().QueryNodeCfg.ExprResCacheEnabled.GetAsBool())
	C.SetExprResCacheEnable(cExprResCacheEnabled)

	if paramtable.Get().QueryNodeCfg.ExprResCacheEnabled.GetAsBool() {
		UpdateExprResCacheConfig()
	}

	C.SetArrowIOThreadPoolCapacity(C.int(ResolveArrowIOThreadPoolCapacity()))

	cStorageV2CellTargetSizeBytes := C.int64_t(paramtable.Get().QueryNodeCfg.StorageV2CellTargetSizeBytes.GetAsInt64())
	C.SetStorageV2CellTargetSizeBytes(cStorageV2CellTargetSizeBytes)
	enableParquetStatsSkipIndex := paramtable.Get().CommonCfg.ParquetStatsSkipIndex.GetAsBool()
	C.SetDefaultEnableParquetStatsSkipIndex(C.bool(enableParquetStatsSkipIndex))

	err := InitArrowReaderConfig(paramtable.Get())
	if err != nil {
		return err
	}

	localDataRootPath := pathutil.GetPath(pathutil.LocalChunkPath, nodeID)

	if err := InitLocalChunkManager(localDataRootPath); err != nil {
		return err
	}

	err = InitRemoteChunkManager(paramtable.Get())
	if err != nil {
		return err
	}

	err = InitDiskFileWriterConfig(paramtable.Get())
	if err != nil {
		return err
	}

	err = InitStorageV2FileSystem(paramtable.Get())
	if err != nil {
		return err
	}

	err = InitMmapManager(paramtable.Get(), nodeID)
	if err != nil {
		return err
	}

	err = InitTieredStorage(paramtable.Get())
	if err != nil {
		return err
	}

	err = InitInterminIndexConfig(paramtable.Get())
	if err != nil {
		return err
	}

	err = InitGeometryCache(paramtable.Get())
	if err != nil {
		return err
	}

	InitTraceConfig(paramtable.Get())
	C.InitExecExpressionFunctionFactory()

	// init paramtable change callback for core related config
	SetupCoreConfigChangelCallback()
	return InitPluginLoader()
}

// SyncPreferFieldDataWhenIndexHasRawData pushes the current paramtable value
// of queryNode.preferFieldDataWhenIndexHasRawData into the segcore C++
// singleton. Safe to call repeatedly; tests invoke it after mutating the
// paramtable so the Go and C++ views of the flag stay in sync.
func SyncPreferFieldDataWhenIndexHasRawData(ctx context.Context, params *paramtable.ComponentParam) {
	v := params.QueryNodeCfg.PreferFieldDataWhenIndexHasRawData.GetAsBool()
	C.SegcoreSetPreferFieldDataWhenIndexHasRawData(C.bool(v))
	if v {
		mlog.Info(ctx, "preferFieldDataWhenIndexHasRawData=true: sealed retrieve will read field data instead of index raw data; "+
			"both will stay resident in memory, increasing the memory footprint for fields whose index also holds raw data")
	}
}

// SyncEnableGrowingSourceFlush pushes the effective growing-source flush switch
// into segcore so growing segments only retain raw chunks when the Go flush path
// may later persist them through StorageV3 FlushGrowingSegmentData.
func SyncEnableGrowingSourceFlush(ctx context.Context, params *paramtable.ComponentParam) {
	storageV3Enabled := params.CommonCfg.UseLoonFFI.GetAsBool()
	v := storageV3Enabled && params.CommonCfg.EnableGrowingSourceFlush.GetAsBool()
	C.SegcoreSetStorageV3Enabled(C.bool(storageV3Enabled))
	C.SegcoreSetEnableGrowingSourceFlush(C.bool(v))
	if v {
		mlog.Info(ctx, "enableGrowingSourceFlush=true: growing segments retain raw field chunks for StorageV3 growing-source flush")
	}
}