Feature/pandas api idxmax

docs/user-guide/advanced/Pandas_API.ipynb

@@ -2373,6 +2373,44 @@
     "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 represent the column name / row number and the values are the result of calling `idxmax` on that column / row. |"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "da7cbf8f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tab.idxmax()"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "301ab2c2",

src/pykx/pandas_api/pandas_meta.py

@@ -228,6 +228,21 @@ def max(self, axis=0, skipna=True, numeric_only=False):
             res
         ), cols)
 
+    @convert_result
+    def idxmax(self, axis=0, skipna=True, numeric_only=False):
+        res, cols = preparse_computations(self, axis, skipna, numeric_only)
+        maximums = self.max(axis, skipna, numeric_only)
+        num_str = '9h$' if numeric_only else ''
+        func_str = '(::)' if axis==0 else 'column_names'
+        if(numeric_only):
+            (_, column_names) = _get_numeric_only_subtable_with_bools(self)
+        else:
+            column_names = q('cols', self)
+        op= q("{[row;col;maximums;column_names]"
+              f"{func_str}[({num_str}row) ?' maximums[col]]"
+              "}")
+        return (op(res, cols, maximums, column_names), 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,21 @@ 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()
+
+
 def test_pandas_all(q):
     tab = q(
         '([] sym: 100?`foo`bar`baz`qux; price: 250.0f - 100?500.0f; ints: 100 - 100?200;'