Add include_key parameter to to_arrays() for primary key retrieval

- to_arrays('a', 'b', include_key=True) now returns (keys, a, b) where
  keys is a list of dicts containing primary key columns
- Keys are in the same format as table.keys(), usable for restrictions
- Added comprehensive tests for the new functionality

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <[email protected]>

Author: dimitri-yatsenko

src/datajoint/expression.py

@@ -678,13 +678,26 @@ def to_arrays(self, *attrs, include_key=False, order_by=None, limit=None, offset
         If attrs specified, returns a tuple of numpy arrays (one per attribute).
 
         :param attrs: attribute names to fetch (if empty, fetch all)
-        :param include_key: if True and attrs specified, include primary key columns
+        :param include_key: if True and attrs specified, prepend primary keys as list of dicts
         :param order_by: attribute(s) to order by, or "KEY"/"KEY DESC"
         :param limit: maximum number of rows to return
         :param offset: number of rows to skip
         :param squeeze: if True, remove extra dimensions from arrays
         :param download_path: path for downloading external data
-        :return: numpy recarray (no attrs) or tuple of arrays (with attrs)
+        :return: numpy recarray (no attrs) or tuple of arrays (with attrs).
+            With include_key=True: (keys, *arrays) where keys is list[dict]
+
+        Examples::
+
+            # Fetch as structured array
+            data = table.to_arrays()
+
+            # Fetch specific columns as separate arrays
+            a, b = table.to_arrays('a', 'b')
+
+            # Fetch with primary keys for later restrictions
+            keys, a, b = table.to_arrays('a', 'b', include_key=True)
+            # keys = [{'id': 1}, {'id': 2}, ...]  # same format as table.keys()
         """
         from functools import partial
 
@@ -702,6 +715,10 @@ def to_arrays(self, *attrs, include_key=False, order_by=None, limit=None, offset
             projected = expr.proj(*fetch_attrs)
             dicts = projected.to_dicts(squeeze=squeeze, download_path=download_path)
 
+            # Extract keys if requested
+            if include_key:
+                keys = [{k: d[k] for k in expr.primary_key} for d in dicts]
+
             # Extract arrays for requested attributes
             result_arrays = []
             for attr in attrs:
@@ -714,6 +731,8 @@ def to_arrays(self, *attrs, include_key=False, order_by=None, limit=None, offset
                     arr = np.array(values, dtype=object)
                 result_arrays.append(arr)
 
+            if include_key:
+                return (keys, *result_arrays)
             return result_arrays[0] if len(attrs) == 1 else tuple(result_arrays)
         else:
             # Fetch all columns as structured array

src/datajoint/version.py

@@ -1,4 +1,4 @@
 # version bump auto managed by Github Actions:
 # label_prs.yaml(prep), release.yaml(bump), post_release.yaml(edit)
 # manually set this version will be eventually overwritten by the above actions
-__version__ = "2.0.0a12"
+__version__ = "2.0.0a13"

tests/integration/test_fetch.py

@@ -344,3 +344,56 @@ def test_lazy_iteration(lang, languages):
     first = next(iter_obj)
     assert isinstance(first, dict)
     assert "name" in first and "language" in first
+
+
+def test_to_arrays_include_key(lang, languages):
+    """Test to_arrays with include_key=True returns keys as list of dicts"""
+    # Fetch with include_key=True
+    keys, names, langs = lang.to_arrays("name", "language", include_key=True, order_by="KEY")
+
+    # keys should be a list of dicts with primary key columns
+    assert isinstance(keys, list)
+    assert all(isinstance(k, dict) for k in keys)
+    assert all(set(k.keys()) == {"name", "language"} for k in keys)
+
+    # names and langs should be numpy arrays
+    assert isinstance(names, np.ndarray)
+    assert isinstance(langs, np.ndarray)
+
+    # Length should match
+    assert len(keys) == len(names) == len(langs) == len(languages)
+
+    # Keys should match the data
+    for key, name, language in zip(keys, names, langs):
+        assert key["name"] == name
+        assert key["language"] == language
+
+    # Keys should be usable for restrictions
+    first_key = keys[0]
+    restricted = lang & first_key
+    assert len(restricted) == 1
+    assert restricted.fetch1("name") == first_key["name"]
+
+
+def test_to_arrays_include_key_single_attr(subject):
+    """Test to_arrays include_key with single attribute"""
+    keys, species = subject.to_arrays("species", include_key=True)
+
+    assert isinstance(keys, list)
+    assert isinstance(species, np.ndarray)
+    assert len(keys) == len(species)
+
+    # Verify keys have only primary key columns
+    assert all("subject_id" in k for k in keys)
+
+
+def test_to_arrays_without_include_key(lang):
+    """Test that to_arrays without include_key doesn't return keys"""
+    result = lang.to_arrays("name", "language")
+
+    # Should return tuple of arrays, not (keys, ...)
+    assert isinstance(result, tuple)
+    assert len(result) == 2
+    names, langs = result
+    assert isinstance(names, np.ndarray)
+    assert isinstance(langs, np.ndarray)