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 minimum 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

@@ -228,6 +228,14 @@ def max(self, axis=0, skipna=True, numeric_only=False):
             res
         ), cols)
 
+    @convert_result
+    def idxmax(self, axis=0, skipna=True, numeric_only=False):
+        tab = self
+        res, cols = preparse_computations(tab, axis, skipna, numeric_only)
+        col_names = _get_numeric_only_subtable_with_bools(tab)[1] if numeric_only else tab.columns
+        max_vals = [elems.index(max(elems)) for elems in res]
+        return (max_vals if axis == 0 else [col_names[i] for i in max_vals], cols)
+
     @convert_result
     def min(self, axis=0, skipna=True, numeric_only=False):
         res, cols = preparse_computations(self, axis, skipna, numeric_only)

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;'