From a4d69031f121561c880046571d8a3f6faad017b8 Mon Sep 17 00:00:00 2001
From: Lior Friedman <lior.friedman@il.kioxia.com>
Date: Thu, 6 Nov 2025 02:53:34 +0200
Subject: [PATCH] fix: Add AiSAQ index type RAM estimation implementation on
 the query node. (#45246)

Currently, the index type AiSAQ RAM usage estimation is not being
calculated correctly.
AiSAQ index type consumes less RAM usage while loading the index than
DISKANN does, and the query node module is missing the implementation of
the RAM usage estimation for that AiSAQ index type.
We suggest that the AiSAQ RAM usage estimation calculation should be as
follows:

UsedDiskMemoryRatioAisaq = 1024 (contrary to the UsedDiskMemoryRatio,
which is 4)
neededMemSize = indexInfo.IndexSize / UsedDiskMemoryRatioAisaq
neededDiskSize = indexInfo.IndexSize

Reported issue is #45247

---------

Signed-off-by: Lior Friedman <lior.friedman@il.kioxia.com>
Signed-off-by: friedl <lior.friedman@kioxia.com>
Co-authored-by: friedl <lior.friedman@kioxia.com>
---
 .../querynodev2/segments/index_attr_cache.go  |  5 +++++
 .../segments/index_attr_cache_test.go         | 19 +++++++++++++++++++
 .../querynodev2/segments/segment_loader.go    |  3 ++-
 internal/util/vecindexmgr/vector_index_mgr.go |  5 +++++
 4 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/internal/querynodev2/segments/index_attr_cache.go b/internal/querynodev2/segments/index_attr_cache.go
index 0e6ae07809..c2b0cdacdb 100644
--- a/internal/querynodev2/segments/index_attr_cache.go
+++ b/internal/querynodev2/segments/index_attr_cache.go
@@ -68,6 +68,11 @@ func (c *IndexAttrCache) GetIndexResourceUsage(indexInfo *querypb.FieldIndexInfo
 		neededDiskSize := indexInfo.IndexSize - neededMemSize
 		return uint64(neededMemSize), uint64(neededDiskSize), nil
 	}
+	if vecindexmgr.GetVecIndexMgrInstance().IsAISAQ(indexType) {
+		neededMemSize := indexInfo.IndexSize / UsedDiskMemoryRatioAisaq
+		neededDiskSize := indexInfo.IndexSize
+		return uint64(neededMemSize), uint64(neededDiskSize), nil
+	}
 	if indexType == indexparamcheck.IndexINVERTED {
 		neededMemSize := 0
 		// we will mmap the binlog if the index type is inverted index.
diff --git a/internal/querynodev2/segments/index_attr_cache_test.go b/internal/querynodev2/segments/index_attr_cache_test.go
index e8966da3c2..7f55ebbb2e 100644
--- a/internal/querynodev2/segments/index_attr_cache_test.go
+++ b/internal/querynodev2/segments/index_attr_cache_test.go
@@ -78,6 +78,25 @@ func (s *IndexAttrCacheSuite) TestDiskANN() {
 	s.EqualValues(75, disk)
 }
 
+func (s *IndexAttrCacheSuite) TestAISAQ() {
+	info := &querypb.FieldIndexInfo{
+		IndexParams: []*commonpb.KeyValuePair{
+			{Key: common.IndexTypeKey, Value: "AISAQ"},
+		},
+		CurrentIndexVersion: 0,
+		IndexSize:           1024,
+	}
+
+	memory, disk, err := s.c.GetIndexResourceUsage(info, paramtable.Get().QueryNodeCfg.MemoryIndexLoadPredictMemoryUsageFactor.GetAsFloat(), nil)
+	s.Require().NoError(err)
+
+	_, has := s.c.loadWithDisk.Get(typeutil.NewPair[string, int32]("AISAQ", 0))
+	s.False(has, "AISAQ shall never be checked load with disk")
+
+	s.EqualValues(16, memory)
+	s.EqualValues(1024, disk)
+}
+
 func (s *IndexAttrCacheSuite) TestInvertedIndex() {
 	info := &querypb.FieldIndexInfo{
 		IndexParams: []*commonpb.KeyValuePair{
diff --git a/internal/querynodev2/segments/segment_loader.go b/internal/querynodev2/segments/segment_loader.go
index b77b46b256..68bb5608ec 100644
--- a/internal/querynodev2/segments/segment_loader.go
+++ b/internal/querynodev2/segments/segment_loader.go
@@ -66,7 +66,8 @@ import (
 )
 
 const (
-	UsedDiskMemoryRatio = 4
+	UsedDiskMemoryRatio      = 4
+	UsedDiskMemoryRatioAisaq = 64
 )
 
 var errRetryTimerNotified = errors.New("retry timer notified")
diff --git a/internal/util/vecindexmgr/vector_index_mgr.go b/internal/util/vecindexmgr/vector_index_mgr.go
index e46324f7b9..36a17fcfec 100644
--- a/internal/util/vecindexmgr/vector_index_mgr.go
+++ b/internal/util/vecindexmgr/vector_index_mgr.go
@@ -77,6 +77,7 @@ type VecIndexMgr interface {
 	IsNoTrainIndex(indexType IndexType) bool
 	IsVecIndex(indexType IndexType) bool
 	IsDiskANN(indexType IndexType) bool
+	IsAISAQ(indexType IndexType) bool
 	IsGPUVecIndex(indexType IndexType) bool
 	IsDiskVecIndex(indexType IndexType) bool
 	IsMMapSupported(indexType IndexType) bool
@@ -108,6 +109,10 @@ func (mgr *vecIndexMgrImpl) IsDiskANN(indexType IndexType) bool {
 	return indexType == "DISKANN"
 }
 
+func (mgr *vecIndexMgrImpl) IsAISAQ(indexType IndexType) bool {
+	return indexType == "AISAQ"
+}
+
 func (mgr *vecIndexMgrImpl) init() {
 	size := int(C.GetIndexListSize())
 	if size == 0 {