diff --git a/CMakeLists.txt b/CMakeLists.txt
index d45b223693cc09c6c76f44a6a6d5faf178b4e5cc..1adfefbd7e323b5c5a83dfa8b5cb0963e3d7cbfb 100755
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -64,6 +64,12 @@ if (ENABLE_OPEN_SRC)
         ${CMAKE_CURRENT_LIST_DIR}/tf_adapter/optimizers/*.cc
         ${CMAKE_CURRENT_LIST_DIR}/tf_adapter/util/*.cc
     )
+    file(GLOB_RECURSE DISCARD_SOURCES
+        ${CMAKE_CURRENT_LIST_DIR}/tf_adapter/kernels/aicpu/map_dataset_op.cc
+        ${CMAKE_CURRENT_LIST_DIR}/tf_adapter/kernels/aicpu/map_and_batch_dataset_op.cc
+        ${CMAKE_CURRENT_LIST_DIR}/tf_adapter/kernels/aicpu/dataset_function.cc
+    )
+    list(REMOVE_ITEM SOURCES ${DISCARD_SOURCES})
     add_library(_tf_adapter SHARED ${SOURCES} ${CMAKE_BINARY_DIR}/dist/swig/ge_plugin_wrap.cxx)
 
     foreach (COMPILE_FLAG ${COMPILE_FLAGS})
diff --git a/tf_adapter/kernels/aicpu/host_queue_dataset_op.cc b/tf_adapter/kernels/aicpu/host_queue_dataset_op.cc
index 574abee93aae6813ffdf4524e971661487f8dc8a..d7f72079cce465b457f555c65f3158024f02b335 100644
--- a/tf_adapter/kernels/aicpu/host_queue_dataset_op.cc
+++ b/tf_adapter/kernels/aicpu/host_queue_dataset_op.cc
@@ -39,13 +39,13 @@
 #include "tf_adapter/util/util.h"
 #include "tf_adapter/util/npu_attrs.h"
 #include "tf_adapter/util/acl_channel.h"
+#include "tf_adapter/util/ge_plugin.h"
 #include "tf_adapter/util/host_queue.h"
 #include "tf_adapter/util/memory_pool.h"
 #include "tf_adapter/util/host_thread_pool.h"
 #include "tf_adapter/util/host_allocator.h"
 #include "tf_adapter/kernels/aicpu/data_item_deliver.h"
 #include "tf_adapter/kernels/aicpu/npu_tensor.h"
-#include "tf_adapter/kernels/aicpu/dataset_function.h"
 
 namespace tensorflow {
 namespace data {
@@ -1092,7 +1092,7 @@ class HostQueueDatasetOp : public DatasetOpKernel {
             TF_RETURN_IF_ERROR(data_deliver_->InitSocketServer());
           }
         }
-        TF_RETURN_IF_ERROR(DatasetFunction::RegisterNpuCancellation(
+        TF_RETURN_IF_ERROR(RegisterNpuCancellationCallback(
             [this]() {
               mutex_lock lck(mu_);
               finish_send_ = true;
diff --git a/tf_adapter/tests/st/CMakeLists.txt b/tf_adapter/tests/st/CMakeLists.txt
index 637cb6edc47753d03b7cabc082b4ec5be7fdd8c6..16830c9a1250e09f1a725569cdf9a61053902c2d 100644
--- a/tf_adapter/tests/st/CMakeLists.txt
+++ b/tf_adapter/tests/st/CMakeLists.txt
@@ -28,7 +28,12 @@ file(GLOB_RECURSE ST_SOURCES
     "kernels/testcase/*.cc"
     "util/testcase/*.cc"
 )
-
+file(GLOB_RECURSE DISCARD_ST_SOURCES
+    "kernels/testcase/dataset/map_dataset_st.cc"
+    "kernels/testcase/dataset/map_and_batch_dataset_st.cc"
+    "kernels/testcase/dataset/dataset_function_st.cc"
+)
+list(REMOVE_ITEM ST_SOURCES ${DISCARD_ST_SOURCES})
 add_executable(tfadapter_stest
                "main.cc"
                ${ST_SOURCES}
diff --git a/tf_adapter/tests/ut/CMakeLists.txt b/tf_adapter/tests/ut/CMakeLists.txt
index d40f64ae0ccb29241e9634c4b11e4dc4604bf2db..4caa091ee2fca917aa8e02505e700167091e6fd2 100644
--- a/tf_adapter/tests/ut/CMakeLists.txt
+++ b/tf_adapter/tests/ut/CMakeLists.txt
@@ -29,7 +29,12 @@ file(GLOB_RECURSE UT_SOURCES
     "kernels/testcase/*.cc"
     "common/testcase/*.cc"
 )
-
+file(GLOB_RECURSE DISCARD_UT_SOURCES
+    "kernels/testcase/dataset/map_dataset_ut.cc"
+    "kernels/testcase/dataset/map_and_batch_dataset_ut.cc"
+    "kernels/testcase/dataset/dataset_function_ut.cc"
+)
+list(REMOVE_ITEM UT_SOURCES ${DISCARD_UT_SOURCES})
 add_executable(tfadapter_utest
                "main.cc"
                ${UT_SOURCES}