diff --git a/js/src/components/query/QueryForm.tsx b/js/src/components/query/QueryForm.tsx
index 1f13f5e0..1c51592f 100644
--- a/js/src/components/query/QueryForm.tsx
+++ b/js/src/components/query/QueryForm.tsx
@@ -1,4 +1,5 @@
 import { useLineageGraphContext } from "@/lib/hooks/LineageGraphContext";
+import { DropdownValuesInput } from "@/utils/DropdownValuesInput";
 import { InfoIcon } from "@chakra-ui/icons";
 import {
   Flex,
@@ -7,17 +8,7 @@ import {
   FormLabel,
   Tooltip,
 } from "@chakra-ui/react";
-import {
-  AutoComplete,
-  AutoCompleteCreatable,
-  AutoCompleteInput,
-  AutoCompleteItem,
-  AutoCompleteList,
-  AutoCompleteTag,
-  Item,
-  ItemTag,
-} from "@choc-ui/chakra-autocomplete";
-import { useMemo } from "react";
+import { useMemo, useRef, useState } from "react";
 
 interface QueryFormProps extends FlexProps {
   defaultPrimaryKeys: string[] | undefined;
@@ -30,13 +21,14 @@ export const QueryForm = ({
   ...prob
 }: QueryFormProps) => {
   const { lineageGraph } = useLineageGraphContext();
+  const labelInfo =
+    "Provide a primary key to perform query diff in data warehouse and only return changed rows.";
 
-  const columns = useMemo(() => {
+  const availableColumns = useMemo(() => {
     if (!lineageGraph) {
       return [];
     }
-
-    const columnSet = new Set();
+    const columnSet = new Set<string>();
     for (const modelName in lineageGraph.nodes) {
       const model = lineageGraph.nodes[modelName];
       const baseColumns = model.data.base?.columns;
@@ -53,57 +45,24 @@ export const QueryForm = ({
     return Array.from(columnSet).sort();
   }, [lineageGraph]);
 
-  const labelInfo =
-    "Provide a primary key to perform query diff in data warehouse and only return changed rows.";
-
   return (
     <Flex {...prob}>
       <FormControl m="4px 8px">
-        <FormLabel>
-          Primary key{" "}
+        <FormLabel fontSize={"8pt"} margin={"0"}>
+          Diff with Primary Key(s) (suggested){" "}
           <Tooltip label={labelInfo}>
             <InfoIcon color="gray.600" boxSize="3" />
           </Tooltip>
         </FormLabel>
-        <AutoComplete
-          restoreOnBlurIfEmpty={false}
-          multiple
-          creatable
-          filter={(query: string, optionValue: Item["value"]) => {
-            return optionValue.startsWith(query);
-          }}
-          onChange={(vals: string[]) => onPrimaryKeysChange(vals)}
-          defaultValues={
-            defaultPrimaryKeys !== undefined && defaultPrimaryKeys.length !== 0
-              ? defaultPrimaryKeys
-              : undefined
-          }
-        >
-          <AutoCompleteInput
-            placeholder="Select primary key..."
-            variant="outline"
-          >
-            {({ tags }: { tags: ItemTag[] }) =>
-              tags.map((tag, tid) => (
-                <AutoCompleteTag
-                  key={tid}
-                  label={tag.label}
-                  onRemove={tag.onRemove}
-                />
-              ))
-            }
-          </AutoCompleteInput>
-          <AutoCompleteList>
-            {columns.map((column, cid) => (
-              <AutoCompleteItem key={`option-${cid}`} value={column}>
-                {column}
-              </AutoCompleteItem>
-            ))}
-            <AutoCompleteCreatable>
-              {({ value }) => <Flex>Add &apos;{value}&apos; to List</Flex>}
-            </AutoCompleteCreatable>
-          </AutoCompleteList>
-        </AutoComplete>
+        <DropdownValuesInput
+          unitName="key"
+          defaultValues={defaultPrimaryKeys}
+          suggestionList={availableColumns}
+          onValuesChange={onPrimaryKeysChange}
+          size="xs"
+          width={"240px"}
+          placeholder="Start by typing key name..."
+        />
       </FormControl>
     </Flex>
   );
diff --git a/js/src/components/query/QueryPage.tsx b/js/src/components/query/QueryPage.tsx
index 5d8b31ac..5fc8a2b9 100644
--- a/js/src/components/query/QueryPage.tsx
+++ b/js/src/components/query/QueryPage.tsx
@@ -122,63 +122,61 @@ export const QueryPage = () => {
   });
 
   return (
-    <HSplit
-      sizes={[90, 10]}
-      minSize={200}
-      style={{ flex: "1", height: "100%" }}
-    >
-      <Flex direction="column" height="100%">
-        <Flex
-          justifyContent="right"
-          alignItems="center"
-          padding="4pt 8pt"
-          gap="5px"
-          height="54px"
-          borderBottom="1px solid lightgray"
-          flex="0 0 54px"
+    <Flex direction="column" height="100%">
+      <Flex
+        justifyContent="right"
+        alignItems="center"
+        padding="4pt 8pt"
+        gap="5px"
+        height="54px"
+        borderBottom="1px solid lightgray"
+        flex="0 0 54px"
+      >
+        <HistoryToggle />
+        <QueryModeToggle />
+        <Spacer />
+        <QueryForm
+          defaultPrimaryKeys={primaryKeys}
+          onPrimaryKeysChange={setPrimaryKeys}
+        />
+        <Button
+          colorScheme="blue"
+          onClick={() => runQuery("query_diff")}
+          isDisabled={isPending}
+          size="sm"
+          leftIcon={<Icon as={VscDiff} />}
         >
-          <HistoryToggle />
-          <QueryModeToggle />
-          <Spacer />
-          <Button
-            colorScheme="blue"
-            onClick={() => runQuery("query_diff")}
-            isDisabled={isPending}
-            size="sm"
-            leftIcon={<Icon as={VscDiff} />}
-          >
-            Run Diff
-          </Button>
-        </Flex>
-
-        <Box width="100%" flex="1">
-          {isCustomQueries ? (
-            <DualSqlEditor
-              value={sqlQuery}
-              baseValue={baseSqlQuery}
-              onChange={setSqlQuery}
-              onChangeBase={setBaseSqlQuery}
-              onRun={() => runQuery("query")}
-              onRunBase={() => runQuery("query_base")}
-              onRunDiff={() => runQuery("query_diff")}
-            />
-          ) : (
-            <SqlEditor
-              value={sqlQuery}
-              onChange={setSqlQuery}
-              onRun={() => runQuery("query")}
-              onRunDiff={() => runQuery("query_diff")}
-            />
-          )}
-        </Box>
+          Run Diff
+        </Button>
       </Flex>
-      <QueryForm
-        p="5px"
-        border="1px"
-        borderColor="gray.300"
-        defaultPrimaryKeys={primaryKeys}
-        onPrimaryKeysChange={setPrimaryKeys}
-      />
-    </HSplit>
+
+      <Box width="100%" flex="1">
+        {isCustomQueries ? (
+          <DualSqlEditor
+            value={sqlQuery}
+            baseValue={baseSqlQuery}
+            onChange={setSqlQuery}
+            onChangeBase={setBaseSqlQuery}
+            onRun={() => runQuery("query")}
+            onRunBase={() => runQuery("query_base")}
+            onRunDiff={() => runQuery("query_diff")}
+          />
+        ) : (
+          <SqlEditor
+            value={sqlQuery}
+            onChange={setSqlQuery}
+            onRun={() => runQuery("query")}
+            onRunDiff={() => runQuery("query_diff")}
+          />
+        )}
+      </Box>
+    </Flex>
+    /* <QueryForm
+      p="5px"
+      border="1px"
+      borderColor="gray.300"
+      defaultPrimaryKeys={primaryKeys}
+      onPrimaryKeysChange={setPrimaryKeys}
+    /> */
   );
 };
diff --git a/js/src/components/query/SqlEditor.tsx b/js/src/components/query/SqlEditor.tsx
index 257e14fd..7bd24e51 100644
--- a/js/src/components/query/SqlEditor.tsx
+++ b/js/src/components/query/SqlEditor.tsx
@@ -1,6 +1,19 @@
-import { Flex, Text, Stack, Badge, Spacer, IconButton } from "@chakra-ui/react";
+import {
+  Flex,
+  Text,
+  Stack,
+  Badge,
+  Spacer,
+  IconButton,
+  Button,
+  Icon,
+  Box,
+} from "@chakra-ui/react";
 import { EditorProps, DiffEditor, Editor } from "@monaco-editor/react";
 import { on } from "events";
+import { BiLogoXing } from "react-icons/bi";
+import { FaPlay } from "react-icons/fa6";
+import { RiPlayMiniFill } from "react-icons/ri";
 import { VscDebugStart } from "react-icons/vsc";
 
 interface SqlEditorProps {
@@ -39,26 +52,27 @@ const SqlEditor: React.FC<SqlEditorProps> = ({
     <>
       {(label || onRun || onRunBase) && (
         <Flex
-          backgroundColor="#E8EFF5"
-          height="18px"
-          fontSize={"12px"}
-          marginBottom={"6px"}
-          paddingX="2"
+          backgroundColor="#EDF2F880"
+          height="40px"
+          fontSize={"14px"}
+          align="center"
+          margin={"0"}
+          padding={"0px 16px"}
         >
           <Text as="b">{label ? label.toUpperCase() : ""}</Text>
           <Spacer />
           {(onRun || onRunBase) && (
-            <IconButton
-              borderRadius={"2px"}
-              aria-label="Run"
-              icon={<VscDebugStart />}
+            <Button
+              size="sm"
+              variant="outline"
               onClick={onRun || onRunBase}
-              color="green"
-              fontSize="18px"
-              size="18px"
-              marginX={2}
-              title={runButtonTitle}
-            />
+              backgroundColor={"white"}
+              // leftIcon={<Icon as={RiPlayMiniFill} />}
+              leftIcon={<Icon as={FaPlay} />}
+              padding={"6px 12px"}
+            >
+              Run Query
+            </Button>
           )}
         </Flex>
       )}
@@ -120,23 +134,29 @@ export const DualSqlEditor: React.FC<SqlEditorProps> = ({
 }: SqlEditorProps) => {
   return (
     <>
-      <Flex height={"100%"}>
-        <Stack height={"100%"} width={"50%"}>
+      <Flex height={"100%"} gap={0}>
+        <Stack
+          height={"100%"}
+          width={"50%"}
+          gap={0}
+          borderRight={"1px"}
+          borderColor={"#D4DBE4"}
+        >
           <SqlEditor
-            label="Current"
-            value={value}
-            onChange={onChange}
-            onRun={onRun}
+            label="Base"
+            value={baseValue || ""}
+            onChange={onChangeBase}
+            onRunBase={onRunBase}
             options={options}
             {...props}
           />
         </Stack>
-        <Stack height={"100%"} width={"50%"}>
+        <Stack height={"100%"} width={"50%"} gap={0}>
           <SqlEditor
-            label="Base"
-            value={baseValue || ""}
-            onChange={onChangeBase}
-            onRunBase={onRunBase}
+            label="Current"
+            value={value}
+            onChange={onChange}
+            onRun={onRun}
             options={options}
             {...props}
           />
diff --git a/js/src/utils/DropdownValuesInput.tsx b/js/src/utils/DropdownValuesInput.tsx
new file mode 100644
index 00000000..7a67fe2d
--- /dev/null
+++ b/js/src/utils/DropdownValuesInput.tsx
@@ -0,0 +1,184 @@
+import React, { useState, useRef } from "react";
+import {
+  Button,
+  Input,
+  InputGroup,
+  InputProps,
+  InputRightElement,
+  Menu,
+  MenuButton,
+  MenuDivider,
+  MenuGroup,
+  MenuItem,
+  MenuList,
+  Portal,
+  Tag,
+  TagCloseButton,
+  TagLabel,
+  Wrap,
+  WrapItem,
+} from "@chakra-ui/react";
+
+export interface DropdownValuesInputProps extends InputProps {
+  unitName: string;
+  suggestionList?: string[];
+  defaultValues?: string[];
+  onValuesChange: (values: string[]) => void;
+}
+
+export const DropdownValuesInput = (props: DropdownValuesInputProps) => {
+  const { defaultValues, suggestionList, onValuesChange } = props;
+  const [values, setValues] = useState<string[]>(defaultValues || []);
+  const [filter, setFilter] = useState<string>("");
+  const [isTyping, setIsTyping] = useState<boolean>(false);
+  const inputRef = useRef<HTMLInputElement>(null);
+
+  const showNumberOfValuesSelected = (tags: string[]) => {
+    if (tags.length > 1) {
+      return `${tags.length} ${props.unitName}s selected`;
+    } else if (tags.length === 1) {
+      return `${tags.length} ${props.unitName} selected`;
+    }
+    return "";
+  };
+
+  const handleSelect = (value: string) => {
+    if (!values.includes(value)) {
+      setFilter("");
+      setValues([...values, value]);
+      onValuesChange([...values, value]);
+    }
+  };
+
+  const handleClear = () => {
+    setFilter("");
+    setValues([]);
+    onValuesChange([]);
+  };
+
+  return (
+    <InputGroup size={props.size} width={props.width}>
+      <Menu isLazy closeOnSelect={false}>
+        <MenuButton width={"100%"}>
+          <Input
+            placeholder={props.placeholder}
+            size={props.size}
+            borderRadius={"4px"}
+            value={showNumberOfValuesSelected(values)}
+          />
+        </MenuButton>
+        <Portal>
+          <MenuList
+            zIndex={"dropdown"}
+            fontSize={props.size}
+            width={props.width}
+          >
+            {/* Input Filter & Show Tags */}
+            <MenuGroup>
+              <Wrap
+                border={"1px solid #e2e8f0"}
+                borderRadius={"4px"}
+                width={"calc(100% - 8px)"}
+                marginX={"4px"}
+                padding={"4px"}
+              >
+                {values.map((value, cid) => (
+                  <WrapItem key={`tag-${cid}`}>
+                    <Tag key={value} size={props.size}>
+                      <TagLabel paddingLeft={"8px"}>{value}</TagLabel>
+                      <TagCloseButton
+                        paddingRight={"8px"}
+                        onClick={() => {
+                          setValues(values.filter((v) => v !== value));
+                          onValuesChange(values.filter((v) => v !== value));
+                        }}
+                      />
+                    </Tag>
+                  </WrapItem>
+                ))}
+                <WrapItem width={"100%"}>
+                  <Input
+                    ref={inputRef}
+                    placeholder="Filter keys or add custom"
+                    variant={"unstyled"}
+                    size={props.size}
+                    value={filter}
+                    onChange={(e) => {
+                      setFilter(e.target.value);
+                      setIsTyping(true);
+                    }}
+                    onKeyDown={(e) => {
+                      const newText = e.currentTarget.value
+                        .trim()
+                        .replace(",", "");
+                      switch (e.key) {
+                        case ",":
+                        case "Enter":
+                          handleSelect(newText);
+                          setFilter("");
+                          break;
+                        case "Backspace":
+                          if (
+                            e.currentTarget.value === "" &&
+                            values.length > 0
+                          ) {
+                            setValues(values.slice(0, -1));
+                            onValuesChange(values.slice(0, -1));
+                          }
+                          break;
+                        default:
+                          break;
+                      }
+                    }}
+                    onBlur={() => {
+                      if (inputRef && isTyping) inputRef?.current?.focus();
+                    }}
+                  />
+                </WrapItem>
+              </Wrap>
+            </MenuGroup>
+            <MenuDivider />
+            {/* Suggestion List */}
+            <MenuGroup>
+              {filter !== "" && !suggestionList?.includes(filter) && (
+                <MenuItem
+                  key={"custom-value-by-filter"}
+                  onClick={() => {
+                    handleSelect(filter);
+                    setIsTyping(false);
+                  }}
+                >
+                  Add &apos;{filter}&apos; to the list
+                </MenuItem>
+              )}
+              {suggestionList
+                ?.filter((value) => filter === "" || value.includes(filter))
+                .filter((value) => !values.includes(value))
+                ?.map((value, cid) => (
+                  <MenuItem
+                    key={`option-${cid}`}
+                    onClick={() => handleSelect(value)}
+                  >
+                    {value}
+                  </MenuItem>
+                ))}
+            </MenuGroup>
+          </MenuList>
+        </Portal>
+      </Menu>
+      <InputRightElement>
+        <Button
+          variant={"link"}
+          color={"#3182CE"}
+          fontSize={props.size}
+          paddingTop="4px"
+          paddingRight={"24px"}
+          hidden={values.length === 0}
+          onClick={handleClear}
+        >
+          Clear
+        </Button>
+      </InputRightElement>
+    </InputGroup>
+  );
+};