diff --git a/x-pack/plugins/canvas/canvas_plugin_src/uis/arguments/datacolumn/index.js b/x-pack/plugins/canvas/canvas_plugin_src/uis/arguments/datacolumn/index.js
index b6343f5ca7ca..bc77f347530b 100644
--- a/x-pack/plugins/canvas/canvas_plugin_src/uis/arguments/datacolumn/index.js
+++ b/x-pack/plugins/canvas/canvas_plugin_src/uis/arguments/datacolumn/index.js
@@ -5,142 +5,128 @@
  * 2.0.
  */
 
-import React, { Component } from 'react';
-import { compose, withPropsOnChange, withHandlers } from 'recompose';
+import React, { useState, useCallback, useEffect } from 'react';
 import PropTypes from 'prop-types';
 import { EuiSelect, EuiFlexItem, EuiFlexGroup } from '@elastic/eui';
 import { sortBy } from 'lodash';
 import { getType } from '@kbn/interpreter/common';
-import { createStatefulPropHoc } from '../../../../public/components/enhance/stateful_prop';
 import { templateFromReactComponent } from '../../../../public/lib/template_from_react_component';
 import { ArgumentStrings } from '../../../../i18n';
 import { SimpleMathFunction } from './simple_math_function';
 import { getFormObject } from './get_form_object';
 
 const { DataColumn: strings } = ArgumentStrings;
+
 const maybeQuoteValue = (val) => (val.match(/\s/) ? `'${val}'` : val);
+const valueNotSet = (val) => !val || val.length === 0;
+
+const getMathValue = (argValue, columns) => {
+  if (getType(argValue) !== 'string') {
+    return { error: 'argValue is not a string type' };
+  }
+  try {
+    const matchedCol = columns.find(({ name }) => argValue === name);
+    const val = matchedCol ? maybeQuoteValue(matchedCol.name) : argValue;
+    const mathValue = getFormObject(val);
+    return { ...mathValue, column: mathValue.column || '' };
+  } catch (e) {
+    return { error: e.message };
+  }
+};
 
 // TODO: Garbage, we could make a much nicer math form that can handle way more.
-class DatacolumnArgInput extends Component {
-  static propTypes = {
-    columns: PropTypes.array.isRequired,
-    onValueChange: PropTypes.func.isRequired,
-    mathValue: PropTypes.object.isRequired,
-    setMathFunction: PropTypes.func.isRequired,
-    typeInstance: PropTypes.object.isRequired,
-    renderError: PropTypes.func.isRequired,
-    argId: PropTypes.string.isRequired,
-  };
+const DatacolumnArgInput = ({
+  onValueChange,
+  columns,
+  argValue,
+  renderError,
+  argId,
+  typeInstance,
+}) => {
+  const [mathValue, setMathValue] = useState(getMathValue(argValue, columns));
 
-  inputRefs = {};
+  useEffect(() => {
+    setMathValue(getMathValue(argValue, columns));
+  }, [argValue, columns]);
 
-  render() {
-    const {
-      onValueChange,
-      columns,
-      mathValue,
-      setMathFunction,
-      renderError,
-      argId,
-      typeInstance,
-    } = this.props;
-
-    if (mathValue.error) {
-      renderError();
-      return null;
-    }
-
-    const allowedTypes = typeInstance.options.allowedTypes || false;
-    const onlyShowMathFunctions = typeInstance.options.onlyMath || false;
-    const valueNotSet = (val) => !val || val.length === 0;
-
-    const updateFunctionValue = () => {
-      const fn = this.inputRefs.fn.value;
-      const column = this.inputRefs.column.value;
+  const allowedTypes = typeInstance.options.allowedTypes || false;
+  const onlyShowMathFunctions = typeInstance.options.onlyMath || false;
 
+  const updateFunctionValue = useCallback(
+    (fn, column) => {
       // if setting size, auto-select the first column if no column is already set
       if (fn === 'size') {
-        const col = column || (columns[0] && columns[0].name);
+        const col = column || (columns[0] && columns[0].name) || '';
         if (col) {
           return onValueChange(`${fn}(${maybeQuoteValue(col)})`);
         }
       }
 
-      // this.inputRefs.column is the column selection, if there is no value, do nothing
+      // if there is no column value, do nothing
       if (valueNotSet(column)) {
-        return setMathFunction(fn);
+        return setMathValue({ ...mathValue, fn });
       }
 
-      // this.inputRefs.fn is the math function to use, if it's not set, just use the value input
+      // if fn is not set, just use the value input
       if (valueNotSet(fn)) {
         return onValueChange(column);
       }
 
-      // this.inputRefs.fn has a value, so use it as a math.js expression
+      // fn has a value, so use it as a math.js expression
       onValueChange(`${fn}(${maybeQuoteValue(column)})`);
-    };
+    },
+    [mathValue, onValueChange, columns]
+  );
 
-    const column =
-      columns.map((col) => col.name).find((colName) => colName === mathValue.column) || '';
+  const onChangeFn = useCallback(
+    ({ target: { value } }) => updateFunctionValue(value, mathValue.column),
+    [mathValue.column, updateFunctionValue]
+  );
 
-    const options = [{ value: '', text: 'select column', disabled: true }];
+  const onChangeColumn = useCallback(
+    ({ target: { value } }) => updateFunctionValue(mathValue.fn, value),
+    [mathValue.fn, updateFunctionValue]
+  );
 
-    sortBy(columns, 'name').forEach((column) => {
-      if (allowedTypes && !allowedTypes.includes(column.type)) {
-        return;
-      }
-      options.push({ value: column.name, text: column.name });
-    });
-
-    return (
-      <EuiFlexGroup gutterSize="s" id={argId} direction="row">
-        <EuiFlexItem grow={false}>
-          <SimpleMathFunction
-            id={argId}
-            value={mathValue.fn}
-            inputRef={(ref) => (this.inputRefs.fn = ref)}
-            onlymath={onlyShowMathFunctions}
-            onChange={updateFunctionValue}
-          />
-        </EuiFlexItem>
-        <EuiFlexItem>
-          <EuiSelect
-            compressed
-            options={options}
-            value={column}
-            inputRef={(ref) => (this.inputRefs.column = ref)}
-            onChange={updateFunctionValue}
-          />
-        </EuiFlexItem>
-      </EuiFlexGroup>
-    );
+  if (mathValue.error) {
+    renderError();
+    return null;
   }
-}
 
-const EnhancedDatacolumnArgInput = compose(
-  withPropsOnChange(['argValue', 'columns'], ({ argValue, columns }) => ({
-    mathValue: ((argValue) => {
-      if (getType(argValue) !== 'string') {
-        return { error: 'argValue is not a string type' };
-      }
-      try {
-        const matchedCol = columns.find(({ name }) => argValue === name);
-        const val = matchedCol ? maybeQuoteValue(matchedCol.name) : argValue;
-        return getFormObject(val);
-      } catch (e) {
-        return { error: e.message };
-      }
-    })(argValue),
-  })),
-  createStatefulPropHoc('mathValue', 'setMathValue'),
-  withHandlers({
-    setMathFunction: ({ mathValue, setMathValue }) => (fn) => setMathValue({ ...mathValue, fn }),
-  })
-)(DatacolumnArgInput);
+  const firstColumnOption = { value: '', text: 'select column', disabled: true };
+  const options = sortBy(columns, 'name')
+    .filter((column) => !allowedTypes || allowedTypes.includes(column.type))
+    .map(({ name }) => ({ value: name, text: name }));
 
-EnhancedDatacolumnArgInput.propTypes = {
-  argValue: PropTypes.oneOfType([PropTypes.string, PropTypes.object]).isRequired,
+  return (
+    <EuiFlexGroup gutterSize="s" id={argId} direction="row">
+      <EuiFlexItem grow={false}>
+        <SimpleMathFunction
+          id={argId}
+          value={mathValue.fn}
+          onlymath={onlyShowMathFunctions}
+          onChange={onChangeFn}
+        />
+      </EuiFlexItem>
+      <EuiFlexItem>
+        <EuiSelect
+          compressed
+          options={[firstColumnOption, ...options]}
+          value={mathValue.column}
+          onChange={onChangeColumn}
+        />
+      </EuiFlexItem>
+    </EuiFlexGroup>
+  );
+};
+
+DatacolumnArgInput.propTypes = {
   columns: PropTypes.array.isRequired,
+  onValueChange: PropTypes.func.isRequired,
+  typeInstance: PropTypes.object.isRequired,
+  renderError: PropTypes.func.isRequired,
+  argId: PropTypes.string.isRequired,
+  argValue: PropTypes.oneOfType([PropTypes.string, PropTypes.object]).isRequired,
 };
 
 export const datacolumn = () => ({
@@ -148,5 +134,5 @@ export const datacolumn = () => ({
   displayName: strings.getDisplayName(),
   help: strings.getHelp(),
   default: '""',
-  simpleTemplate: templateFromReactComponent(EnhancedDatacolumnArgInput),
+  simpleTemplate: templateFromReactComponent(DatacolumnArgInput),
 });
diff --git a/x-pack/plugins/canvas/canvas_plugin_src/uis/arguments/datacolumn/simple_math_function.js b/x-pack/plugins/canvas/canvas_plugin_src/uis/arguments/datacolumn/simple_math_function.js
index c54b44d3cc5b..5fd859a789d8 100644
--- a/x-pack/plugins/canvas/canvas_plugin_src/uis/arguments/datacolumn/simple_math_function.js
+++ b/x-pack/plugins/canvas/canvas_plugin_src/uis/arguments/datacolumn/simple_math_function.js
@@ -12,7 +12,7 @@ import { ArgumentStrings } from '../../../../i18n';
 
 const { DataColumn: strings } = ArgumentStrings;
 
-export const SimpleMathFunction = ({ onChange, value, inputRef, onlymath }) => {
+export const SimpleMathFunction = ({ onChange, value, onlymath }) => {
   const options = [
     { text: strings.getOptionAverage(), value: 'mean' },
     { text: strings.getOptionCount(), value: 'size' },
@@ -29,15 +29,12 @@ export const SimpleMathFunction = ({ onChange, value, inputRef, onlymath }) => {
     options.unshift({ text: strings.getOptionValue(), value: '' });
   }
 
-  return (
-    <EuiSelect compressed options={options} inputRef={inputRef} value={value} onChange={onChange} />
-  );
+  return <EuiSelect compressed options={options} value={value} onChange={onChange} />;
 };
 
 SimpleMathFunction.propTypes = {
   onChange: PropTypes.func,
   value: PropTypes.string,
-  inputRef: PropTypes.func,
   onlymath: PropTypes.bool,
 };