Repository: cogent3/cogent3
Pull request #1478 (2023-06-14 23:17:39)

ENH: more control over column naming in joined tables


Patch:
diff --git a/pyproject.toml b/pyproject.toml
index 366145343..08f07b1d5 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -158,7 +158,7 @@ src_paths = ["src", "tests","doc","noxfile.py"]
 
 [tool.scriv]
 format="md"
-categories=["Contributers","ENH","BUG","DOC","Deprecations","Discontinued"]
+categories=["Contributors", "ENH", "BUG", "DOC", "Deprecations", "Discontinued"]
 output_file="changelog.md"
 version="literal: src/cogent3/__init__.py:__version__"
 skip_fragments="README.*"
diff --git a/src/cogent3/util/table.py b/src/cogent3/util/table.py
index fe3d02de8..f7186e578 100644
--- a/src/cogent3/util/table.py
+++ b/src/cogent3/util/table.py
@@ -911,31 +911,38 @@ class Table:
     def array(self):
         return self.columns.array
 
-    def cross_join(self, other, **kwargs):
+    def cross_join(self, other, col_prefix="right_", **kwargs):
         """cross join, or full outer join, of self with other
 
-        Notes
-        -----
-        The column headers of the output are made unique by prepending
-        other column headers with _, e.g. 'Name' becomes _Name'.
+        Parameters
+        ----------
+        other
+            A table object which will be joined with this
+            table. other must have a title.
+        col_prefix
+            ensure <other> columns are unique by prepending col_prefix
         """
         self_range = range(self.shape[0])
         other_range = range(other.shape[0])
         self_selected, other_selected = list(zip(*product(self_range, other_range)))
         joined_data = {c: self.columns[c].take(self_selected) for c in self.columns}
-        col_prefix = "right" if not other.title else other.title
         other_data = {
-            f"{col_prefix}_{c}": other.columns[c].take(other_selected)
+            f"{col_prefix}{c}": other.columns[c].take(other_selected)
             for c in other.columns
         }
 
         joined_data.update(other_data)
         new_header = list(self.columns.order) + [
-            f"{col_prefix}_{c}" for c in other.columns
+            f"{col_prefix}{c}" for c in other.columns
         ]
-        attrs = self._get_persistent_attrs()
-        attrs.pop("title", None)
-        attrs |= kwargs
+        attrs = {**self._get_persistent_attrs()}
+        # we do not propagate current title as it is likely incorrect
+        # we also drop self.index_name since a cross-join generates duplicated
+        # values of that column
+        for key in ("title", "index_name"):
+            attrs.pop(key, None)
+
+        attrs.update(kwargs)
         joined = self.__class__(**attrs)
         for c in new_header:
             joined.columns[c] = joined_data[c]
@@ -947,6 +954,7 @@ class Table:
         columns_self=None,
         columns_other=None,
         use_index=True,
+        col_prefix="right_",
         **kwargs,
     ):
         """inner join of self with other
@@ -965,14 +973,9 @@ class Table:
         use_index
             if no columns specified and both self and other have a nominated
             index_name, this will be used.
-
-        Notes
-        -----
-        The column headers of the output are made unique by prepending
-        other column headers with _, e.g. 'Name' becomes _Name'.
+        col_prefix
+            ensure <other> columns are unique by prepending col_prefix
         """
-        col_prefix = "right" if not other.title else other.title
-
         if columns_self:
             columns_self = self.columns._get_keys_(columns_self)
 
@@ -1032,12 +1035,12 @@ class Table:
 
         joined_data = {c: self.columns[c][self_selected] for c in self.columns}
         other_data = {
-            f"{col_prefix}_{c}": other.columns[c][other_selected] for c in output_mask
+            f"{col_prefix}{c}": other.columns[c][other_selected] for c in output_mask
         }
 
         joined_data.update(other_data)
         new_header = list(self.columns.order) + [
-            f"{col_prefix}_{c}" for c in output_mask
+            f"{col_prefix}{c}" for c in output_mask
         ]
         attr = self._get_persistent_attrs()
         attr.pop("title", None)
@@ -1053,6 +1056,7 @@ class Table:
         columns_self=None,
         columns_other=None,
         inner_join=True,
+        col_prefix="right_",
         **kwargs,
     ):
         """returns a new table containing the join of this table and
@@ -1069,6 +1073,7 @@ class Table:
             columns_self=columns_self,
             columns_other=columns_other,
             use_index=False,
+            col_prefix=col_prefix,
             **kwargs,
         )
 


Test patch:
diff --git a/tests/test_util/test_table.py b/tests/test_util/test_table.py
index 0d71d9f98..eade243f8 100644
--- a/tests/test_util/test_table.py
+++ b/tests/test_util/test_table.py
@@ -1,5 +1,3 @@
-#!/usr/bin/env python
-
 """Unit tests for table.
 """
 import json
@@ -9,7 +7,7 @@ import pickle
 
 from collections import defaultdict
 from tempfile import TemporaryDirectory
-from unittest import TestCase, main, skipIf
+from unittest import TestCase, skipIf
 
 import numpy
 import pytest
@@ -1974,5 +1972,51 @@ def test_mixed_row_lengths(data):
         _ = Table(header=["a", "b", "c"], data=data)
 
 
-if __name__ == "__main__":
-    main()
+@pytest.fixture
+def t2():
+    header = ["id", "foo", "bar"]
+    rows = [
+        [1, "abc", 11],
+        [2, "bca", 22],
+        [3, "cab", 33],
+        [4, "abc", 44],
+        [5, "bca", 55],
+    ]
+    return Table(header=header, data=rows, index_name="id")
+
+
+@pytest.fixture
+def t3():
+    header = ["id", "foo2", "bar2"]
+    rows = [[6, "abc", 66], [7, "bca", 77]]
+    return Table(header=header, data=rows, index_name="id")
+
+
+@pytest.fixture
+def t4():
+    header = ["id", "foo2", "bar2"]
+    rows = [[6, "abc", 66], [7, "bca", 77]]
+    return Table(header=header, data=rows, index_name="id")
+
+
+def test_inner_join_col_naming(t2, t3):
+    """test the table joined method"""
+    # inner join with defaults
+    got = t2.inner_join(t3)
+    expect = list(t2.header) + [f"right_{c}" for c in t3.header[1:]]
+    assert list(got.header) == expect
+
+    # inner join with no col_prefix
+    got = t2.inner_join(t3, col_prefix="")
+    expect = list(t2.header) + list(t3.header[1:])
+    assert list(got.header) == expect
+
+
+def test_outer_join_col_naming(t2, t3):
+    """test the table joined method"""
+    # cross join with a new col prefix
+    col_prefix = "cp_"
+    got = t2.cross_join(t3, col_prefix=col_prefix)
+    expect = list(t2.header) + [f"{col_prefix}{c}" for c in t3.header]
+    assert list(got.header) == expect
+    assert got.shape[0] == t2.shape[0] * t3.shape[0]