fix(python): allow selection of _rowid in Permutation (#3133)

Closes #3132

python/python/lancedb/permutation.py

@@ -885,7 +885,7 @@ class Permutation:
         This method refines the current selection, potentially removing columns.  It
         will not add back columns that were previously removed.
 
-        If any of the columns do not exist then an error will be raised
+        If any of the columns do not exist then an error will be raised.
 
         This does not introduce a post-processing step.  It simply reduces the amount
         of data we read.
@@ -898,9 +898,14 @@ class Permutation:
         for name in columns:
             value = self.selection.get(name, None)
             if value is None:
-                raise ValueError(
+                if name == "_rowid":
-                    f"Cannot select column `{name}` because it does not exist"
+                    # _rowid is a system column not in the default schema
-                )
+                    # but can be explicitly selected
+                    value = "_rowid"
+                else:
+                    raise ValueError(
+                        f"Cannot select column `{name}` because it does not exist"
+                    )
             new_selection[name] = value
         return self._with_selection(new_selection)
 

python/python/tests/test_permutation.py

@@ -1136,3 +1136,61 @@ def test_take_offsets_empty_permutation(some_permutation: Permutation):
     result = some_permutation.take_offsets([])
 
     assert result == []
+
+
+def test_select_rowid(some_permutation: Permutation):
+    """Test that _rowid can be selected alongside regular columns."""
+    perm_with_rowid = some_permutation.select_columns(["_rowid", "id"])
+    assert "_rowid" in perm_with_rowid.column_names
+    batches = list(perm_with_rowid.iter(100, skip_last_batch=False))
+    for batch in batches:
+        assert "_rowid" in batch[0]
+
+
+def test_select_rowid_only(some_permutation: Permutation):
+    """Test that _rowid can be selected as the sole column."""
+    perm_rowid_only = some_permutation.select_columns(["_rowid"])
+    assert perm_rowid_only.column_names == ["_rowid"]
+    batches = list(perm_rowid_only.iter(100, skip_last_batch=False))
+    assert len(batches) > 0
+    for batch in batches:
+        assert list(batch[0].keys()) == ["_rowid"]
+
+
+def test_select_rowid_not_in_default(some_permutation: Permutation):
+    """Test that _rowid is NOT in the default column_names or schema."""
+    assert "_rowid" not in some_permutation.column_names
+    assert "_rowid" not in some_permutation.schema.names
+
+
+def test_select_rowid_identity_permutation(mem_db):
+    """Test that _rowid works with an identity permutation."""
+    tbl = mem_db.create_table(
+        "test_rowid_identity", pa.table({"id": range(10), "value": range(10)})
+    )
+    perm = Permutation.identity(tbl)
+    perm_with_rowid = perm.select_columns(["_rowid", "id"])
+    batches = list(perm_with_rowid.iter(10, skip_last_batch=False))
+    assert len(batches) == 1
+    assert "_rowid" in batches[0][0]
+
+
+def test_rename_rowid(some_permutation: Permutation):
+    """Test that _rowid can be selected and then renamed."""
+    perm_with_rowid = some_permutation.select_columns(["_rowid", "id"])
+    renamed = perm_with_rowid.rename_column("_rowid", "my_row_id")
+    assert "my_row_id" in renamed.column_names
+    assert "_rowid" not in renamed.column_names
+    batches = list(renamed.iter(100, skip_last_batch=False))
+    for batch in batches:
+        assert "my_row_id" in batch[0]
+        assert "_rowid" not in batch[0]
+
+
+def test_remove_rowid_after_select(some_permutation: Permutation):
+    """Test that _rowid can be selected and then removed."""
+    perm_with_rowid = some_permutation.select_columns(["_rowid", "id"])
+    assert "_rowid" in perm_with_rowid.column_names
+    perm_without_rowid = perm_with_rowid.remove_columns(["_rowid"])
+    assert "_rowid" not in perm_without_rowid.column_names
+    assert perm_without_rowid.column_names == ["id"]