diff --git a/docs/user-guide/advanced/Pandas_API.ipynb b/docs/user-guide/advanced/Pandas_API.ipynb index a240a90..006db60 100644 --- a/docs/user-guide/advanced/Pandas_API.ipynb +++ b/docs/user-guide/advanced/Pandas_API.ipynb @@ -3330,6 +3330,105 @@ "\n", "Example Table." ] + }, + { + "cell_type": "markdown", + "id": "7f08eb84", + "metadata": {}, + "source": [ + "## Comparison\n", + "\n", + "### Table.isin()\n", + "\n", + "```\n", + "Table.isin(\n", + " values\n", + ")\n", + "```\n", + "\n", + "Whether each element in the DataFrame is contained in values.\n", + "\n", + "**Parameters:**\n", + "\n", + "| Name | Type | Description | Default |\n", + "| :--------------: | :---------------------------------: | :-------------------------------------------------------------------------- | :------: |\n", + "| values | Union[List, dict, Table, KeyedTable] | The result will only be true at a location if all the labels match. If values is a dict, the keys must be the column names, which must match. If values is a Table or KeyedTable, then both the index and column labels must match. | None|\n", + "\n", + "\n", + "**Returns:**\n", + "\n", + "| Type | Description |\n", + "| :-----------------------: | :---------------------------------------------- |\n", + "| Table | Boolean type Table/KeyedTable showing whether each element in the DataFrame is contained in values.|\n", + "\n", + "**Examples:**\n", + "\n", + "Example Table." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f6e453c8", + "metadata": {}, + "outputs": [], + "source": [ + "tab = kx.Table(data={'x': list(range(3)), 'y': [\"A\", \"B\", \"C\"]})" + ] + }, + { + "cell_type": "markdown", + "id": "aadd23c1", + "metadata": {}, + "source": [ + "Find if element \"A\" or \"1\" is in the table:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d41d40e0", + "metadata": {}, + "outputs": [], + "source": [ + "tab.isin([\"A\", 1])" + ] + }, + { + "cell_type": "markdown", + "id": "cff856fe", + "metadata": {}, + "source": [ + "Find if element \"A\" is in colum \"y\":" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bccf59d9", + "metadata": {}, + "outputs": [], + "source": [ + "tab.isin({\"y\": [\"A\"]})" + ] + }, + { + "cell_type": "markdown", + "id": "ed704cce", + "metadata": {}, + "source": [ + "Find if element \"A\" is in the first position of \"y\" column:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "41840cc0", + "metadata": {}, + "outputs": [], + "source": [ + "tab.isin(kx.Table(data={\"y\":[\"A\"]}))" + ] } ], "metadata": { diff --git a/src/pykx/pandas_api/pandas_meta.py b/src/pykx/pandas_api/pandas_meta.py index fab85fb..217f410 100644 --- a/src/pykx/pandas_api/pandas_meta.py +++ b/src/pykx/pandas_api/pandas_meta.py @@ -290,6 +290,52 @@ def abs(self, numeric_only=False): return q.abs(tab) @api_return + def isin(self, values): + tab = self + key_table = 'KeyedTable' in str(type(tab)) + key_value = 'KeyedTable' in str(type(values)) + n_rows = 0 + false_dataframe_f = q("""{u:(cols x); + v:(count[u],count[x])#0b; + flip u!v}""") + if key_value and not key_table: + return false_dataframe_f(tab) + if key_table: + kcols = q.key(tab) + if key_value: + n_rows, tab = q("""{n_rows:max 0, count[x]- + count rows:(key y) inter key x; + (n_rows; x each rows)}""", tab, values) + values = q.value(values) + else: + tab = q.value(tab) + dic_value, is_tab = q("""{$[98h = type x; + (flip x; 1b); + (x; 0b)]}""", values) + if key_table and not key_value and is_tab: + ftable = false_dataframe_f(tab) + else: + ftable = q("""{ [table; values; is_tab; n_rows] + flip (cols table)! + {[col_name; tab; values; v_is_tab; n_rows] + col: tab col_name; + ltype: .Q.ty col; + values: $[99h~type values; values col_name; values]; + $[v_is_tab or ltype=" "; ; + values@:where (lower ltype) = .Q.t abs type each values]; + $[0 = count values; + (n_rows + count[col])#0b; + $[v_is_tab; + $[any ltype = (" ";"C"); ~'; =] + [mlen#col;mlen#values], + (n_rows + max 0,count[col]- + mlen: min count[values], + count[col])#0b; + any $[any ltype = (" ";"C"); ~/:\:; =\:][values;col] + ]]}[; table; values; is_tab; n_rows] + each cols table}""", tab, dic_value, is_tab, n_rows) + return ftable.set_index(kcols) if key_table else ftable + def round(self, decimals: Union[int, Dict[str, int]] = 0): tab = self if 'Keyed' in str(type(tab)): diff --git a/tests/test_pandas_api.py b/tests/test_pandas_api.py index 0a1dbf8..8e93545 100644 --- a/tests/test_pandas_api.py +++ b/tests/test_pandas_api.py @@ -2228,6 +2228,30 @@ def test_keyed_loc_fixes(q): mkt['k1'] +def test_pandas_isin(kx): + tab = kx.q("""([] k1: 0n 1. 0n 2. 0n; + k2: ("A";" ";"B";" ";"A"); + k3: (`a;1.;`c;5;`d))""") + keyed_tab = kx.q("""([`a`b`c`d`e] + k1: 0n 1. 0n 2. 0n; + k2: ("A";" ";"B";" ";"A"); + k3: (`a;1.;`c;5;`d))""") + + list_value = kx.q('(`a;1.;"A")') + dict_value = {"k1": [1., 2., 3.]} + tab_value = kx.q('([] k1: 1. 2. 3.; k2: ("A";"B";"C"))') + keyed_tab_value = kx.q('([`a`b] k1: 1. 2.; k2: ("A";"B"))') + + assert tab.isin(list_value).pd().equals(tab.pd().isin(list_value.py())) + assert tab.isin(dict_value).pd().equals(tab.pd().isin(dict_value)) + assert tab.isin(tab_value).pd().equals(tab.pd().isin(tab_value.pd())) + assert tab.isin(keyed_tab_value).pd().equals(tab.pd().isin(keyed_tab_value)) + assert keyed_tab.isin(list_value).pd().equals(keyed_tab.pd().isin(list_value.py())) + assert keyed_tab.isin(dict_value).pd().equals(keyed_tab.pd().isin(dict_value)) + assert keyed_tab.isin(keyed_tab_value).pd().equals(keyed_tab.pd().isin(keyed_tab_value.pd())) + assert keyed_tab.isin(tab_value).pd().equals(keyed_tab.pd().isin(tab_value)) + + def test_nunique(kx, q): tab = kx.q('([]a:4 0n 7 6;b:4 0n 0n 7;c:``foo`foo`)') df = tab.pd()