Feature/pandas api idxmax

docs/user-guide/advanced/Pandas_API.ipynb

@@ -2373,6 +2373,72 @@
     "tab.max()"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "d98b298c",
+   "metadata": {},
+   "source": [
+    "### Table.idxmax()\n",
+    "\n",
+    "```\n",
+    "Table.idxmax(axis=0, skipna=True, numeric_only=False)\n",
+    "```\n",
+    "\n",
+    "Return index of first occurrence of maximum over requested axis.\n",
+    "\n",
+    "**Parameters:**\n",
+    "\n",
+    "| Name         | Type | Description                                                                      | Default |\n",
+    "| :----------: | :--: | :------------------------------------------------------------------------------- | :-----: |\n",
+    "| axis         | int  | The axis to calculate the idxmax across 0 is columns, 1 is rows.                | 0       |\n",
+    "| skipna       | bool | Ignore any null values along the axis.                                           | True    |\n",
+    "| numeric_only | bool | Only use columns of the table that are of a numeric data type.                   | False   |\n",
+    "\n",
+    "**Returns:**\n",
+    "\n",
+    "| Type               | Description                                                          |\n",
+    "| :----------------: | :------------------------------------------------------------------- |\n",
+    "| Dictionary         | A dictionary where the key represents the column name / row number and the values are the result of calling `idxmax` on that column / row. |"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "143f5483",
+   "metadata": {},
+   "source": [
+    "**Examples:**\n",
+    "\n",
+    "Calculate the idxmax across the columns of a table"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "da7cbf8f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tab.idxmax()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fb531e00",
+   "metadata": {},
+   "source": [
+    "Calculate the idxmax across the rows of a table using only columns thar are of a numeric data type"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9907226a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tab.idxmax(axis=1, numeric_only=True)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "301ab2c2",

src/pykx/pandas_api/pandas_meta.py

@@ -67,7 +67,7 @@ def preparse_computations(tab, axis=0, skipna=True, numeric_only=False, bool_onl
         skipna,
         axis
     )
-    return (res, cols if axis == 0 else q.til(len(res)))
+    return (res, cols if axis == 0 else q.til(len(res)), cols)
 
 
 # The simple computation functions all return a tuple of the results and the col names the results
@@ -212,33 +212,45 @@ def abs(self, numeric_only=False):
 
     @convert_result
     def all(self, axis=0, bool_only=False, skipna=True):
-        res, cols = preparse_computations(self, axis, skipna, bool_only=bool_only)
+        res, cols, _ = preparse_computations(self, axis, skipna, bool_only=bool_only)
         return (q('{"b"$x}', [all(x) for x in res]), cols)
 
     @convert_result
     def any(self, axis=0, bool_only=False, skipna=True):
-        res, cols = preparse_computations(self, axis, skipna, bool_only=bool_only)
+        res, cols, _ = preparse_computations(self, axis, skipna, bool_only=bool_only)
         return (q('{"b"$x}', [any(x) for x in res]), cols)
 
     @convert_result
     def max(self, axis=0, skipna=True, numeric_only=False):
-        res, cols = preparse_computations(self, axis, skipna, numeric_only)
+        res, cols, _ = preparse_computations(self, axis, skipna, numeric_only)
         return (q(
             '{[row] {$[11h=type x; {[x1; y1] $[x1 > y1; x1; y1]} over x; max x]} each row}',
             res
         ), cols)
 
     @convert_result
     def min(self, axis=0, skipna=True, numeric_only=False):
-        res, cols = preparse_computations(self, axis, skipna, numeric_only)
+        res, cols, _ = preparse_computations(self, axis, skipna, numeric_only)
         return (q(
             '{[row] {$[11h=type x; {[x1; y1] $[x1 < y1; x1; y1]} over x; min x]} each row}',
             res
         ), cols)
 
+    @convert_result
+    def idxmax(self, axis=0, skipna=True, numeric_only=False):
+        tab = self
+        axis = q('{$[11h~type x; `index`columns?x; x]}', axis)
+        res, cols, ix = preparse_computations(tab, axis, skipna, numeric_only)
+        return (q(
+            '''{[row;tab;axis]
+                row:{$[11h~type x; {[x1; y1] $[x1 > y1; x1; y1]} over x; max x]} each row;
+                m:$[0~axis; (::); flip] value flip tab;
+                $[0~axis; (::); cols tab] m {$[abs type y;x]?y}' row}
+            ''', res, tab[ix], axis), cols)
+
     @convert_result
     def prod(self, axis=0, skipna=True, numeric_only=False, min_count=0):
-        res, cols = preparse_computations(self, axis, skipna, numeric_only)
+        res, cols, _ = preparse_computations(self, axis, skipna, numeric_only)
         return (q(
             '{[row; minc] {$[y > 0; $[y>count[x]; 0N; prd x]; prd x]}[;minc] each row}',
             res,
@@ -247,7 +259,7 @@ def prod(self, axis=0, skipna=True, numeric_only=False, min_count=0):
 
     @convert_result
     def sum(self, axis=0, skipna=True, numeric_only=False, min_count=0):
-        res, cols = preparse_computations(self, axis, skipna, numeric_only)
+        res, cols, _ = preparse_computations(self, axis, skipna, numeric_only)
         return (q(
             '{[row; minc]'
             '{$[y > 0;'

tests/test_pandas_api.py

@@ -1811,6 +1811,29 @@ def test_pandas_max(q):
         assert float(qmax[i]) == float(pmax[i])
 
 
+def test_pandas_idxmax(q):
+    tab = q('([] sym: 100?`foo`bar`baz`qux; price: 250.0f - 100?500.0f; ints: 100 - 100?200)')
+    df = tab.pd()
+
+    p_m = df.idxmax()
+    q_m = tab.idxmax()
+    for c in q.key(q_m).py():
+        assert p_m[c] == q_m[c].py()
+
+    q_m = tab.idxmax(axis=1, numeric_only=True, skipna=True)
+    p_m = df.idxmax(axis=1, numeric_only=True, skipna=True)
+    for c in q.key(q_m).py():
+        assert p_m[c] == q_m[c].py()
+
+    tab = q('([]price: 250.0f - 100?500.0f; ints: 100 - 100?200)')
+    df = tab.pd()
+
+    q_m = tab.idxmax(axis=1)
+    p_m = df.idxmax(axis=1)
+    for c in q.key(q_m).py():
+        assert p_m[c] == q_m[c].py()
+
+
 def test_pandas_all(q):
     tab = q(
         '([] sym: 100?`foo`bar`baz`qux; price: 250.0f - 100?500.0f; ints: 100 - 100?200;'