The in operator in ResultSetCollection.append causes problems with numpy arrays

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## 0.11.0dev
 
 * [API Change] Disabled `%sql` and `%%sql` on Databricks ([#1047](https://github.com/ploomber/jupysql/issues/1047))
+* [Fix] Add support for numpy arrays in result sets
 
 ## 0.10.17 (2025-01-08)
 

src/sql/connection/connection.py

@@ -5,6 +5,7 @@
 from difflib import get_close_matches
 import atexit
 from functools import partial
+from itertools import starmap
 
 import sqlalchemy
 from sqlalchemy.engine import Engine
@@ -36,6 +37,7 @@
 from sql.warnings import JupySQLQuotedNamedParametersWarning, JupySQLRollbackPerformed
 from sql import _current
 from sql.connection import error_handling
+from sql.run.resultset import ResultSet
 
 BASE_DOC_URL = "https://jupysql.ploomber.io/en/latest"
 
@@ -107,13 +109,35 @@ def extract_module_name_from_NoSuchModuleError(e):
     return str(e).split(":")[-1].split(".")[-1]
 
 
+def _bool(x):
+    if x is True or x is False:
+        return x
+    return all(x)
+
+
+def _eq(a, b):
+    return a is b or _bool(a == b)
+
+
+def _results(x):
+    if isinstance(x, ResultSet):
+        return x._results
+    return x
+
+
 class ResultSetCollection:
     def __init__(self) -> None:
         self._result_sets = []
 
     def append(self, result):
-        if result in self._result_sets:
-            self._result_sets.remove(result)
+        for idx in reversed(
+            [
+                i
+                for i, item in enumerate(self._result_sets)
+                if all(starmap(_eq, zip(_results(result), _results(item))))
+            ]
+        ):
+            self._result_sets.pop(idx)
 
         self._result_sets.append(result)
 

src/tests/test_connection.py

@@ -1067,23 +1067,42 @@ def test_transpile_query_doesnt_transpile_if_it_doesnt_need_to(monkeypatch):
 
 def test_result_set_collection_append():
     collection = ResultSetCollection()
-    collection.append(1)
-    collection.append(2)
+    collection.append((1,))
+    collection.append((2,))
 
-    assert collection._result_sets == [1, 2]
+    assert collection._result_sets == [(1,), (2,)]
+
+
+def test_result_set_collection_append_numpy():
+    try:
+        import numpy as np
+
+        a1 = (np.array([1, 2]),)
+        a2 = (np.array([3, 4]),)
+
+        collection = ResultSetCollection()
+        collection.append(a1)
+        collection.append(a2)
+
+        assert len(collection._result_sets) == 2
+        assert collection._result_sets[0] is a1
+        assert collection._result_sets[1] is a2
+
+    except ImportError:
+        pass
 
 
 def test_result_set_collection_iterate():
     collection = ResultSetCollection()
-    collection.append(1)
-    collection.append(2)
+    collection.append((1,))
+    collection.append((2,))
 
-    assert list(collection) == [1, 2]
+    assert list(collection) == [(1,), (2,)]
 
 
 def test_result_set_collection_is_last():
     collection = ResultSetCollection()
-    first, second = object(), object()
+    first, second = (object(),), (object(),)
     collection.append(first)
 
     assert len(collection) == 1