diff --git a/CHANGELOG.md b/CHANGELOG.md
index 42d5b8978f4ee0432c56ef7e69076a31366b5c89..b6c1e323ecf55b37e9776d471273690abef1c614 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -13,6 +13,7 @@ Please mark all change in change log and use the ticket from JIRA.
 - \#353 - Rename config.h.in to version.h.in
 - \#374 - sdk_simple return empty result
 - \#397 - sdk_simple return incorrect result
+- \#440 - Query API in customization still uses old version
 
 ## Feature
 - \#12 - Pure CPU version for Milvus
diff --git a/core/src/scheduler/action/PushTaskToNeighbour.cpp b/core/src/scheduler/action/PushTaskToNeighbour.cpp
index f49f1d871f40c5efa91d2866a6977ece46656a93..562041cee1adc48413d3d4aed691604021a5c28c 100644
--- a/core/src/scheduler/action/PushTaskToNeighbour.cpp
+++ b/core/src/scheduler/action/PushTaskToNeighbour.cpp
@@ -106,41 +106,6 @@ Action::SpecifiedResourceLabelTaskScheduler(const ResourceMgrPtr& res_mgr, Resou
                                             std::shared_ptr<LoadCompletedEvent> event) {
     auto task_item = event->task_table_item_;
     auto task = event->task_table_item_->task;
-    //    if (resource->type() == ResourceType::DISK) {
-    //        // step 1: calculate shortest path per resource, from disk to compute resource
-    //        auto compute_resources = res_mgr->GetComputeResources();
-    //        std::vector<std::vector<std::string>> paths;
-    //        std::vector<uint64_t> transport_costs;
-    //        for (auto& res : compute_resources) {
-    //            std::vector<std::string> path;
-    //            uint64_t transport_cost = ShortestPath(resource, res, res_mgr, path);
-    //            transport_costs.push_back(transport_cost);
-    //            paths.emplace_back(path);
-    //        }
-    //        if (task->job_.lock()->type() == JobType::BUILD) {
-    //            // step2: Read device id in config
-    //            // get build index gpu resource
-    //            server::Config& config = server::Config::GetInstance();
-    //            int32_t build_index_gpu;
-    //            Status stat = config.GetResourceConfigIndexBuildDevice(build_index_gpu);
-    //
-    //            bool find_gpu_res = false;
-    //            if (res_mgr->GetResource(ResourceType::GPU, build_index_gpu) != nullptr) {
-    //                for (uint64_t i = 0; i < compute_resources.size(); ++i) {
-    //                    if (compute_resources[i]->name() ==
-    //                        res_mgr->GetResource(ResourceType::GPU, build_index_gpu)->name()) {
-    //                        find_gpu_res = true;
-    //                        Path task_path(paths[i], paths[i].size() - 1);
-    //                        task->path() = task_path;
-    //                        break;
-    //                    }
-    //                }
-    //            }
-    //            if (not find_gpu_res) {
-    //                task->path() = Path(paths[0], paths[0].size() - 1);
-    //            }
-    //        }
-    //    }
 
     if (resource->name() == task->path().Last()) {
         resource->WakeupExecutor();
diff --git a/core/unittest/db/test_db.cpp b/core/unittest/db/test_db.cpp
index d8614dd5d14b580de1751db78a3489cdc5a38d87..9e01b5ea1ac2602bdf81eff7ff8b25110c1881f3 100644
--- a/core/unittest/db/test_db.cpp
+++ b/core/unittest/db/test_db.cpp
@@ -305,24 +305,30 @@ TEST_F(DBTest, SEARCH_TEST) {
     // test FAISS_IVFSQ8H optimizer
     index.engine_type_ = (int)milvus::engine::EngineType::FAISS_IVFSQ8H;
     db_->CreateIndex(TABLE_NAME, index);  // wait until build index finish
+    std::vector<std::string> partition_tag;
+    milvus::engine::ResultIds result_ids;
+    milvus::engine::ResultDistances result_dists;
 
     {
-        milvus::engine::QueryResults results;
-        stat = db_->Query(TABLE_NAME, k, nq, 10, xq.data(), results);
+        result_ids.clear();
+        result_dists.clear();
+        stat = db_->Query(TABLE_NAME, partition_tag, k, nq, 10, xq.data(), result_ids, result_dists);
         ASSERT_TRUE(stat.ok());
     }
 
     {
-        milvus::engine::QueryResults large_nq_results;
-        stat = db_->Query(TABLE_NAME, k, 200, 10, xq.data(), large_nq_results);
+        result_ids.clear();
+        result_dists.clear();
+        stat = db_->Query(TABLE_NAME, partition_tag, k, 200, 10, xq.data(), result_ids, result_dists);
         ASSERT_TRUE(stat.ok());
     }
 
     {  // search by specify index file
         milvus::engine::meta::DatesT dates;
         std::vector<std::string> file_ids = {"1", "2", "3", "4", "5", "6"};
-        milvus::engine::QueryResults results;
-        stat = db_->Query(TABLE_NAME, file_ids, k, nq, 10, xq.data(), dates, results);
+        result_ids.clear();
+        result_dists.clear();
+        stat = db_->QueryByFileID(TABLE_NAME, file_ids, k, nq, 10, xq.data(), dates, result_ids, result_dists);
         ASSERT_TRUE(stat.ok());
     }