diff --git a/frontend/src/metadata/components/cell-editors/tags-editor/index.js b/frontend/src/metadata/components/cell-editors/tags-editor/index.js
index b389122a74..85b3c57f7f 100644
--- a/frontend/src/metadata/components/cell-editors/tags-editor/index.js
+++ b/frontend/src/metadata/components/cell-editors/tags-editor/index.js
@@ -41,6 +41,7 @@ const TagsEditor = forwardRef(({
   const [maxItemNum, setMaxItemNum] = useState(0);
   const [nodes, setNodes] = useState([]);
   const [keyNodeFoldedMap, setKeyNodeFoldedMap] = useState({});
+  const [searchedKeyNodeFoldedMap, setSearchedKeyNodeFoldedMap] = useState({});
   const [recentlyUsed, setRecentlyUsed] = useState([]);
   const itemHeight = 30;
   const editorContainerRef = useRef(null);
@@ -243,7 +244,8 @@ const TagsEditor = forwardRef(({
     }
   }, []);
 
-  const getShownNodes = useCallback((tree, keyNodeFoldedMap) => {
+  const getShownNodes = useCallback(() => {
+    const tree = tagsData?.rows_tree;
     if (!Array.isArray(tree)) {
       return [];
     }
@@ -253,17 +255,6 @@ const TagsEditor = forwardRef(({
       const nodeId = getTreeNodeId(node);
       const row = getRowById(tagsData, nodeId);
       if (!row) return;
-      if (searchValue) {
-        const value = searchValue.toLowerCase();
-        const tagName = getTagName(row).toLowerCase();
-        const tagColor = getTagColor(row).toLowerCase();
-        if (!tagName.includes(value) && !tagColor.includes(value)) return;
-        if (showTagsAsTree) {
-          const nodesWithAncestors = getNodesWithAncestors(node, tree).filter(node => checkIsTreeNodeShown(getTreeNodeKey(node), keyNodeFoldedMap));
-          shownNodes = [...shownNodes, ...nodesWithAncestors];
-          return;
-        }
-      }
       const nodeKey = getTreeNodeKey(node);
       if (row && checkIsTreeNodeShown(nodeKey, keyNodeFoldedMap)) {
         shownNodes.push({
@@ -273,19 +264,47 @@ const TagsEditor = forwardRef(({
       }
     });
     return shownNodes;
-  }, [tagsData, searchValue, showTagsAsTree]);
+  }, [tagsData, keyNodeFoldedMap]);
+
+  const getSearchedNodes = useCallback(() => {
+    const tree = tagsData?.rows_tree;
+    if (!Array.isArray(tree)) {
+      return [];
+    }
+    let searchedNodes = [];
+    tree.forEach((node, index) => {
+      const nodeId = getTreeNodeId(node);
+      const row = getRowById(tagsData, nodeId);
+      if (!row) return;
+      const value = searchValue.toLowerCase();
+      const tagName = getTagName(row).toLowerCase();
+      const tagColor = getTagColor(row).toLowerCase();
+      if (!tagName.includes(value) && !tagColor.includes(value)) return;
+      const nodesWithAncestors = getNodesWithAncestors(node, tree).filter(node => checkIsTreeNodeShown(getTreeNodeKey(node), searchedKeyNodeFoldedMap));
+      searchedNodes = [...searchedNodes, ...nodesWithAncestors];
+    });
+    return searchedNodes;
+  }, [tagsData, searchValue, searchedKeyNodeFoldedMap]);
 
   const toggleExpandTreeNode = useCallback((nodeKey) => {
-    const updatedKeyNodeFoldedMap = { ...keyNodeFoldedMap };
-    if (updatedKeyNodeFoldedMap[nodeKey]) {
-      delete updatedKeyNodeFoldedMap[nodeKey];
+    if (!searchValue) {
+      const updatedKeyNodeFoldedMap = { ...keyNodeFoldedMap };
+      if (updatedKeyNodeFoldedMap[nodeKey]) {
+        delete updatedKeyNodeFoldedMap[nodeKey];
+      } else {
+        updatedKeyNodeFoldedMap[nodeKey] = true;
+      }
+      setKeyNodeFoldedMap(updatedKeyNodeFoldedMap);
     } else {
-      updatedKeyNodeFoldedMap[nodeKey] = true;
+      const updatedSearchedKeyNodeFoldedMap = { ...searchedKeyNodeFoldedMap };
+      if (updatedSearchedKeyNodeFoldedMap[nodeKey]) {
+        delete updatedSearchedKeyNodeFoldedMap[nodeKey];
+      } else {
+        updatedSearchedKeyNodeFoldedMap[nodeKey] = true;
+      }
+      setSearchedKeyNodeFoldedMap(updatedSearchedKeyNodeFoldedMap);
     }
-    const updatedNodes = getShownNodes(tagsData.rows_tree, updatedKeyNodeFoldedMap);
-    setNodes(updatedNodes);
-    setKeyNodeFoldedMap(updatedKeyNodeFoldedMap);
-  }, [tagsData, keyNodeFoldedMap, getShownNodes]);
+  }, [keyNodeFoldedMap, searchedKeyNodeFoldedMap, searchValue]);
 
   useEffect(() => {
     if (editorRef.current) {
@@ -312,12 +331,26 @@ const TagsEditor = forwardRef(({
     setRecentlyUsed(tags);
   }, [tagsData, localStorage]);
 
+  useEffect(() => {
+    if (tagsData?.rows_tree && showTagsAsTree) {
+      const updatedKeyNodeFoldedMap = tagsData.rows_tree.reduce((acc, node) => {
+        const nodeKey = getTreeNodeKey(node);
+        if (checkTreeNodeHasChildNodes(node)) {
+          acc[nodeKey] = true;
+        }
+        return acc;
+      }, {});
+      setKeyNodeFoldedMap(updatedKeyNodeFoldedMap);
+    }
+  // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, []);
+
   useEffect(() => {
     if (tagsData?.rows_tree) {
-      const shownNodes = getShownNodes(tagsData.rows_tree, keyNodeFoldedMap);
+      const shownNodes = searchValue ? getSearchedNodes() : getShownNodes();
       setNodes(shownNodes);
     }
-  }, [tagsData, keyNodeFoldedMap, getShownNodes]);
+  }, [tagsData, searchValue, getSearchedNodes, getShownNodes]);
 
   const renderOptions = useCallback(() => {
     if (displayTags.length === 0) {
@@ -394,14 +427,14 @@ const TagsEditor = forwardRef(({
               onMouseLeave={onTreeMenuMouseLeave}
               depth={getTreeNodeDepth(node)}
               hasChildren={checkTreeNodeHasChildNodes(node)}
-              isFolded={keyNodeFoldedMap[nodeKey]}
+              isFolded={!searchValue ? keyNodeFoldedMap[nodeKey] : searchedKeyNodeFoldedMap[nodeKey]}
               onToggleExpand={() => toggleExpandTreeNode(nodeKey)}
             />
           );
         })}
       </>
     );
-  }, [nodes, tagsData, value, highlightNodeIndex, searchValue, recentlyUsedTags, renderRecentlyUsed, toggleExpandTreeNode, keyNodeFoldedMap, onSelectTag, onTreeMenuMouseEnter, onTreeMenuMouseLeave]);
+  }, [nodes, tagsData, value, highlightNodeIndex, searchValue, recentlyUsedTags, renderRecentlyUsed, toggleExpandTreeNode, keyNodeFoldedMap, searchedKeyNodeFoldedMap, onSelectTag, onTreeMenuMouseEnter, onTreeMenuMouseLeave]);
 
   return (
     <div className={classnames('sf-metadata-tags-editor', { 'tags-tree-container': showTagsAsTree })} style={style} ref={editorRef}>
diff --git a/frontend/src/metadata/components/cell-editors/tags-editor/tag-item/index.js b/frontend/src/metadata/components/cell-editors/tags-editor/tag-item/index.js
index 686aa9483a..6d7903993f 100644
--- a/frontend/src/metadata/components/cell-editors/tags-editor/tag-item/index.js
+++ b/frontend/src/metadata/components/cell-editors/tags-editor/tag-item/index.js
@@ -49,7 +49,7 @@ const TagItem = ({
           <div className="sf-metadata-tag-color" style={{ backgroundColor: tagColor }} />
           <div className="sf-metadata-tag-name">{tagName}</div>
         </div>
-        <div className="sf-metadata-tags-editor-tag-check-icon">
+        <div className="sf-metadata-tags-editor-tag-check-icon mr-1">
           {isSelected && <Icon className="sf-metadata-icon" symbol="check-mark" />}
         </div>
       </div>