Fix sort cell behavior for query results

querybook/webapp/components/DataDocStatementExecution/StatementResultTable.tsx

@@ -1,16 +1,10 @@
-import { produce } from 'immer';
 import classNames from 'classnames';
-import React, {
-    useRef,
-    useMemo,
-    useState,
-    useCallback,
-    useEffect,
-} from 'react';
+import React, { useMemo, useState, useCallback, useEffect } from 'react';
 import { useSelector } from 'react-redux';
 import styled from 'styled-components';
 
 import { UserSettingsFontSizeToCSSFontSize } from 'const/font';
+import { useImmer } from 'hooks/useImmer';
 import { IColumnTransformer } from 'lib/query-result/types';
 import { findColumnType } from 'lib/query-result/detector';
 import { isNumeric } from 'lib/utils/number';
@@ -76,10 +70,9 @@ export const StatementResultTable: React.FunctionComponent<{
     paginate: boolean;
     maxNumberOfRowsToShow?: number;
 }> = ({ data, paginate, maxNumberOfRowsToShow = 20, isPreview = false }) => {
-    const [expandedColumn, setExpandedColumn] = React.useState<
-        Record<string, boolean>
-    >({});
-    const [columnTransformerByIndex, setColumnTransformer] = useState<
+    const [expandedColumn, toggleExpandedColumn] = useExpandedColumn();
+
+    const [columnTransformerByIndex, setColumnTransformer] = useImmer<
         Record<string, IColumnTransformer>
     >({});
 
@@ -107,10 +100,9 @@ export const StatementResultTable: React.FunctionComponent<{
     );
     const setTransformerForColumn = useCallback(
         (colIndex: number, transformer: IColumnTransformer) => {
-            setColumnTransformer((old) => ({
-                ...old,
-                [colIndex]: transformer,
-            }));
+            setColumnTransformer((old) => {
+                old[colIndex] = transformer;
+            });
         },
         []
     );
@@ -120,7 +112,7 @@ export const StatementResultTable: React.FunctionComponent<{
             <StatementResultTableColumn
                 column={column}
                 expandedColumn={expandedColumn}
-                setExpandedColumn={setExpandedColumn}
+                toggleExpandedColumn={toggleExpandedColumn}
                 rows={rows}
                 colIndex={index}
                 colType={columnTypes[index]}
@@ -173,36 +165,79 @@ export const StatementResultTable: React.FunctionComponent<{
                         ? transformer.transform(row[index])
                         : row[index];
                 }}
-                sortCell={(a, b) => {
-                    if (a == null || a === 'null') {
-                        return -1;
-                    } else if (b == null || b === 'null') {
-                        return 1;
-                    }
-                    const aValue = isNumeric(a) ? Number(a) : String(a);
-                    const bValue = isNumeric(a) ? Number(b) : String(b);
-                    return aValue < bValue ? -1 : aValue > bValue ? 1 : 0;
-                }}
+                sortCell={useSortCell(rows)}
             />
         </StyledTableWrapper>
     );
 };
 
-function useLastNotNull<T>(value: T): T {
-    const [s, set] = useState(null);
-    useEffect(() => {
-        if (value != null) {
-            set(value);
-        }
-    }, [value]);
+function useExpandedColumn() {
+    const [expandedColumn, setExpandedColumn] = useImmer<
+        Record<string, boolean>
+    >({});
+    const toggleExpandedColumn = useCallback((column: string) => {
+        setExpandedColumn((old) => {
+            if (column in old) {
+                delete old[column];
+            } else {
+                old[column] = true;
+            }
+        });
+    }, []);
+    return [expandedColumn, toggleExpandedColumn] as const;
+}
 
-    return s;
+type ColumnSortType = 'string' | 'number';
+function isCellNull(cell: any) {
+    return cell === 'null' || cell == null;
+}
+
+function useSortCell(rows: string[][]) {
+    const columnTypeCache: Record<number, ColumnSortType> = useMemo(
+        () => ({}),
+        [rows]
+    );
+
+    const sortCell = useCallback(
+        (colIdx: number, a: any, b: any) => {
+            if (isCellNull(a)) {
+                return -1;
+            } else if (isCellNull(b)) {
+                return 1;
+            }
+
+            if (!(colIdx in columnTypeCache)) {
+                columnTypeCache[colIdx] = getColumnType(rows, colIdx);
+            }
+
+            const colType = columnTypeCache[colIdx];
+            if (colType === 'number') {
+                a = Number(a);
+                b = Number(b);
+            }
+            return a < b ? -1 : a > b ? 1 : 0;
+        },
+        [rows, columnTypeCache]
+    );
+    return sortCell;
+}
+
+function getColumnType(rows: any[][], colIdx: number): ColumnSortType {
+    for (const row of rows) {
+        const cell = row[colIdx];
+        if (isCellNull(cell)) {
+            continue;
+        } else if (!isNumeric(cell)) {
+            return 'string';
+        }
+    }
+    return 'number';
 }
 
 const StatementResultTableColumn: React.FC<{
     column: string;
     expandedColumn: Record<string, boolean>;
-    setExpandedColumn: (c: Record<string, boolean>) => any;
+    toggleExpandedColumn: (c: string) => any;
 
     // These props are for the column info
     colIndex: number;
@@ -217,7 +252,7 @@ const StatementResultTableColumn: React.FC<{
 }> = ({
     column,
     expandedColumn,
-    setExpandedColumn,
+    toggleExpandedColumn,
 
     colIndex,
     colType,
@@ -259,15 +294,7 @@ const StatementResultTableColumn: React.FC<{
                         e.preventDefault();
                         e.stopPropagation();
 
-                        setExpandedColumn(
-                            produce(expandedColumn, (draft) => {
-                                if (isExpanded) {
-                                    delete draft[column];
-                                } else {
-                                    draft[column] = true;
-                                }
-                            })
-                        );
+                        toggleExpandedColumn(column);
                     }}
                 />
                 <Dropdown
@@ -333,3 +360,14 @@ const StatementResultTableColumn: React.FC<{
         </Level>
     );
 };
+
+function useLastNotNull<T>(value: T): T {
+    const [s, set] = useState(null);
+    useEffect(() => {
+        if (value != null) {
+            set(value);
+        }
+    }, [value]);
+
+    return s;
+}

querybook/webapp/hooks/useImmer.ts

@@ -0,0 +1,16 @@
+import produce from 'immer';
+import { useCallback, useState } from 'react';
+
+export type SetStateImmer<S> = (f: (draft: S) => void) => void;
+export function useImmer<S>(
+    initialValue: S | (() => S)
+): [S, SetStateImmer<S>] {
+    const [state, setState] = useState(initialValue);
+
+    return [
+        state,
+        useCallback((updater) => {
+            setState((prevState) => produce(prevState, updater));
+        }, []),
+    ];
+}

querybook/webapp/lib/query-result/detector.ts

@@ -6,11 +6,7 @@ const columnDetectors: IColumnDetector[] = [
     {
         type: 'string',
         priority: 0,
-        checker: (colName: string, values: any[]) =>
-            detectTypeForValues(values, (v) => {
-                const vType = typeof v;
-                return vType === 'string';
-            }),
+        checker: (colName: string, values: any[]) => true,
     },
     {
         type: 'number',

querybook/webapp/lib/utils/number.ts

@@ -75,14 +75,13 @@ export function roundNumberToDecimal(n: number, decimalPlaces: number) {
 }
 
 export function isNumeric(n: any): boolean {
-    if (typeof n === 'bigint') {
-        return true;
-    }
-    if (typeof n === 'number') {
-        return !isNaN(n);
-    } else if (typeof n === 'string') {
+    if (typeof n === 'string') {
         // typescript thinks isNaN can only take number
         return !isNaN(n as any) && !isNaN(parseFloat(n));
+    } else if (typeof n === 'number') {
+        return !isNaN(n);
+    } else if (typeof n === 'bigint') {
+        return true;
     }
     return false;
 }

querybook/webapp/ui/Table/Table.tsx

@@ -28,7 +28,7 @@ export interface ITableProps extends Partial<TableProps> {
         row: any
     ) => React.ReactNode;
 
-    sortCell?: (cellA: any, cellB: any) => -1 | 0 | 1;
+    sortCell?: (columnIdx: number, cellA: any, cellB: any) => -1 | 0 | 1;
 }
 
 export class Table extends React.Component<ITableProps> {
@@ -46,13 +46,22 @@ export class Table extends React.Component<ITableProps> {
         widthObj?: Record<string, number>,
         alignObj?: Record<string, 'left' | 'right' | 'center'>
     ): Column[] {
-        return columns.map((column: string, columnIndex) => {
+        return columns.map((column: Column | string, columnIndex: number) => {
             let formattedColumn: Column;
             if (typeof column === 'string') {
                 formattedColumn = {
                     Header: titleize(column, '_', ' '),
                     accessor: column,
                 };
+                if (widthObj && widthObj[column]) {
+                    formattedColumn.minWidth = widthObj[column];
+                }
+
+                if (alignObj && alignObj[column]) {
+                    formattedColumn.style = {
+                        textAlign: alignObj[column],
+                    };
+                }
             } else {
                 formattedColumn = column;
             }
@@ -65,18 +74,7 @@ export class Table extends React.Component<ITableProps> {
                 );
             }
             if (sortCell) {
-                formattedColumn.sortMethod = sortCell;
-            }
-
-            if (widthObj && widthObj[column]) {
-                formattedColumn.minWidth = widthObj[column];
-            }
-
-            if (alignObj && alignObj[column]) {
-                formattedColumn.style = {
-                    ...formattedColumn.style,
-                    textAlign: alignObj[column],
-                };
+                formattedColumn.sortMethod = sortCell.bind(null, columnIndex);
             }
 
             return formattedColumn;