diff --git a/src/components/QueryEditor.tsx b/src/components/QueryEditor.tsx
index d6b8805..57afc55 100644
--- a/src/components/QueryEditor.tsx
+++ b/src/components/QueryEditor.tsx
@@ -1,8 +1,9 @@
 import React, { ChangeEvent, useRef, FormEventHandler } from 'react';
 import _ from 'lodash';
 import defaults from 'lodash/defaults';
-import { InlineField, InlineFieldRow, HorizontalGroup, InlineSwitch, Input, Select } from '@grafana/ui';
-import { QueryEditorProps } from '@grafana/data';
+import { InlineField, InlineSwitch, Input, Select, Button, Icon, useStyles2 } from '@grafana/ui';
+import { QueryEditorProps, GrafanaTheme2 } from '@grafana/data';
+import { css } from '@emotion/css';
 import { DataSource } from '../datasource';
 import { TIMESHIFT_FORMAT_REG, TIMESHIFT_VAR_FORMAT_REG } from '../config';
 import { CompareQueriesOptions, CompareQueriesQuery, defaultQuery } from '../types';
@@ -21,6 +22,15 @@ export function QueryEditor({ query, onChange, onRunQuery }: Props) {
   const aliasTypes = ['suffix', 'prefix', 'absolute'];
   let target: Record<string, any> = defaults(query, defaultQuery);
 
+  const getStyles = (theme: GrafanaTheme2) => ({
+    root: css({
+      display: 'flex',
+    }),
+    addButton: css({
+      padding: '0px 10px',
+    }),
+  });
+
   // Query refId change event
   const onQueryRefChange = (event: ChangeEvent<HTMLInputElement>) => {
     target.query = event.target.value;
@@ -154,117 +164,102 @@ export function QueryEditor({ query, onChange, onRunQuery }: Props) {
     return !TIMESHIFT_FORMAT_REG.test(amountValue) && !TIMESHIFT_VAR_FORMAT_REG.test(amountValue)
   };
 
+  const styles = useStyles2(getStyles);
+
   return (
     <div>
-      <div className="gf-form-inline">
-        <InlineFieldRow>
-          <div className="gf-form">
-            <InlineField
-              label='Query'
-              className='query-keyword'
-              labelWidth={8}
-            >
-              <Input
-                spellCheck={false}
-                placeholder='query'
-                value={target.query}
-                onChange={onQueryRefChange}
-                onBlur={handleBlur}
-              />
-            </InlineField>
-
-            <InlineField label="Process TimeShift" className='width-12'>
-              <InlineSwitch value={target.process} onChange={onProcessChange}/>
-            </InlineField>
-
-          </div>
-        </InlineFieldRow>
+      <div className={styles.root}>
+        <InlineField
+          label='Query'
+          className='query-keyword'
+          labelWidth={8}
+        >
+          <Input
+            spellCheck={false}
+            placeholder='query'
+            value={target.query}
+            onChange={onQueryRefChange}
+            onBlur={handleBlur}
+          />
+        </InlineField>
+
+        <InlineField label="Process TimeShift" className='width-12'>
+          <InlineSwitch value={target.process} onChange={onProcessChange}/>
+        </InlineField>
       </div>
 
-      <HorizontalGroup>
-        <InlineFieldRow>
-          {
-            target.timeShifts?.map((timeShift: any) => {
-              return (
-                <div className="gf-form" key={timeShift}>
-                  <span className="gf-form-label">
-                    <i className="fa fa-clock-o"></i>
-                  </span>
-                  <span className="gf-form-label width-6">Time shift</span>
-
-                  <InlineField
-                    labelWidth={8}
-                    label='Amount'
-                    invalid={timeShift.value && isInvalidAmount(timeShift.value)}
-                    error={timeShift.value && isInvalidAmount(timeShift.value) ? 'Amount format is invalid' : ''}
-                  >
-                    <Input
-                      className='width-8'
-                      placeholder='1h'
-                      value={timeShift.value}
-                      required
-                      onChange={onChangeHandler('value', { jsonData: timeShift })} onBlur={handleBlur}
-                    />
-                  </InlineField>
-
-                  <InlineField
-                    label="alias"
-                    labelWidth={8}
-                  >
-                    <Input
-                      className='width-8'
-                      placeholder='auto'
-                      value={timeShift.alias}
-                      onChange={onChangeHandler('alias', { jsonData: timeShift })}
-                      onBlur={handleBlur}
-                    />
-                  </InlineField>
-
-                  <InlineField
-                    label='alias type'
-                    labelWidth={12}
-                  >
-                    <Select
-                      className='width-7'
-                      defaultValue={'suffix'}
-                      value={timeShift.aliasType || 'suffix'}
-                      onChange={({ value }) => onAliasTypeChange(value, { jsonData: timeShift })}
-                      onBlur={handleBlur}
-                      options={aliasTypes.map(val => ({value: val, label: val}))}
-                    />
-                  </InlineField>
-
-                  <span className="gf-form-label width-4" title="only valid when alias type is suffix or prefix">delimiter</span>
+      <div>
+        {
+          target.timeShifts?.map((timeShift: any) => {
+            return (
+              <div className={styles.root} key={timeShift}>
+                <span className="gf-form-label width-6"><Icon name='history' size='sm'/>Time shift</span>
+
+                <InlineField
+                  labelWidth={8}
+                  label='Amount'
+                  invalid={timeShift.value && isInvalidAmount(timeShift.value)}
+                  error={timeShift.value && isInvalidAmount(timeShift.value) ? 'Amount format is invalid' : ''}
+                >
                   <Input
-                      disabled={timeShift.aliasType === 'absolute'}
-                      className='width-8'
-                      placeholder='default:_'
-                      value={timeShift.delimiter}
-                      onChange={onChangeHandler('delimiter', { jsonData: timeShift })}
-                      onBlur={handleBlur}
+                    className='width-8'
+                    placeholder='1h'
+                    value={timeShift.value}
+                    required
+                    onChange={onChangeHandler('value', { jsonData: timeShift })} onBlur={handleBlur}
                   />
+                </InlineField>
 
-                  {
-                    target.timeShifts.length > 1 && (<label className="gf-form-label">
-                      <a className="pointer" onClick={() => removeTimeShift(timeShift)}>
-                        <i className="fa fa-trash"></i>
-                      </a>
-                    </label>)
-                  }
-                </div>)
-            })
-          }
-        </InlineFieldRow>
-
-      </HorizontalGroup>
-
-      <button
-        className="btn btn-secondary gf-form-btn"
-        onClick={addTimeShifts}
-      >
-        Add time shift
-      </button>
+                <InlineField
+                  label="alias"
+                  labelWidth={8}
+                >
+                  <Input
+                    className='width-8'
+                    placeholder='auto'
+                    value={timeShift.alias}
+                    onChange={onChangeHandler('alias', { jsonData: timeShift })}
+                    onBlur={handleBlur}
+                  />
+                </InlineField>
+
+                <InlineField
+                  label='alias type'
+                  labelWidth={12}
+                >
+                  <Select
+                    className='width-7'
+                    defaultValue={'suffix'}
+                    value={timeShift.aliasType || 'suffix'}
+                    onChange={({ value }) => onAliasTypeChange(value, { jsonData: timeShift })}
+                    onBlur={handleBlur}
+                    options={aliasTypes.map(val => ({value: val, label: val}))}
+                  />
+                </InlineField>
+
+                <span className="gf-form-label width-4" title="only valid when alias type is suffix or prefix">delimiter</span>
+                <Input
+                    disabled={timeShift.aliasType === 'absolute'}
+                    className='width-8'
+                    placeholder='default:_'
+                    value={timeShift.delimiter}
+                    onChange={onChangeHandler('delimiter', { jsonData: timeShift })}
+                    onBlur={handleBlur}
+                />
+
+                {
+                  target.timeShifts.length > 1 && (<label className="gf-form-label">
+                    <a className="pointer" onClick={() => removeTimeShift(timeShift)}>
+                      <i className="fa fa-trash"></i>
+                    </a>
+                  </label>)
+                }
+              </div>)
+          })
+        }
+      </div>
 
+      <Button variant='secondary' icon='plus' className={styles.addButton} onClick={addTimeShifts}>Add time shift</Button>
     </div>
   );
 }