Return to_frame columns in the requested order; accept multi-column broadcast keys

HISTORY.rst

@@ -13,6 +13,12 @@ Next release
   installs the package and builds the documentation.
 * Automate publishing to PyPI from GitHub releases using Trusted
   Publishing, and document the contribution and release process.
+* ``to_frame(columns)`` now returns columns in the order requested,
+  whether they are local or registered columns (previously the order
+  of registered columns varied from call to call), and
+  ``merge_tables(columns=...)`` resolves columns in the order given.
+* ``broadcast`` accepts lists of column names for ``cast_on`` and
+  ``onto_on``, for merges on multiple columns, as ``pandas.merge`` does.
 
 v1.8
 ====

orca/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.9.dev2"
+__version__ = "1.9.dev3"
 
 from .orca import *
 

orca/orca.py

@@ -372,9 +372,10 @@ def to_frame(self, columns=None):
         Parameters
         ----------
         columns : sequence or string, optional
-            Sequence of the column names desired in the DataFrame. A string
-            can also be passed if only one column is desired.
-            If None all columns are returned, including registered columns.
+            Sequence of the column names desired in the DataFrame, which
+            are returned in the order given. A string can also be passed
+            if only one column is desired. If None all columns are
+            returned, including registered columns.
 
         Returns
         -------
@@ -385,13 +386,19 @@ def to_frame(self, columns=None):
 
         if columns is not None:
             columns = [columns] if isinstance(columns, str) else columns
-            columns = set(columns)
-            set_extra_cols = set(extra_cols)
-            local_cols = set(self.local.columns) & columns - set_extra_cols
-            df = self.local[list(local_cols)].copy()
-            extra_cols = {k: extra_cols[k] for k in (columns & set_extra_cols)}
+            # keep the requested order (dropping any duplicates) so that
+            # the columns of the result are in a predictable order
+            columns = list(dict.fromkeys(columns))
+            local_cols = [c for c in columns
+                          if c in self.local.columns and c not in extra_cols]
+            df = self.local[local_cols].copy()
+            extra_cols = {c: extra_cols[c] for c in columns if c in extra_cols}
+            # final column order; inserting each registered column at its
+            # position avoids copying the frame again to reorder it
+            order = [c for c in columns if c in local_cols or c in extra_cols]
         else:
             df = self.local.copy()
+            order = None
 
         with log_start_finish(
                 'computing {!r} columns for table {!r}'.format(
@@ -402,7 +409,10 @@ def to_frame(self, columns=None):
                         'computing column {!r} for table {!r}'.format(
                             name, self.name),
                         logger):
-                    df[name] = col()
+                    if order is None:
+                        df[name] = col()
+                    else:
+                        df.insert(order.index(name), name, col())
 
         return df
 
@@ -1401,14 +1411,17 @@ def column_map(tables, columns):
     if not columns:
         return {t.name: None for t in tables}
 
-    columns = set(columns)
-    colmap = {
-        t.name: list(set(t.columns).intersection(columns)) for t in tables}
+    columns = list(dict.fromkeys(columns))
+    colmap = {}
+    for t in tables:
+        table_cols = set(t.columns)
+        colmap[t.name] = [c for c in columns if c in table_cols]
     foundcols = tz.reduce(
         lambda x, y: x.union(y), (set(v) for v in colmap.values()))
-    if foundcols != columns:
+    if foundcols != set(columns):
         raise RuntimeError('Not all required columns were found. '
-                           'Missing: {}'.format(list(columns - foundcols)))
+                           'Missing: {}'.format(
+                               [c for c in columns if c not in foundcols]))
     return colmap
 
 
@@ -1739,8 +1752,8 @@ def broadcast(cast, onto, cast_on=None, onto_on=None,
     ----------
     cast, onto : str
         Names of registered tables.
-    cast_on, onto_on : str, optional
-        Column names used for merge, equivalent of ``left_on``/``right_on``
+    cast_on, onto_on : str or list of str, optional
+        Column name(s) used for merge, equivalent of ``left_on``/``right_on``
         parameters of pandas.merge.
     cast_index, onto_index : bool, optional
         Whether to use table indexes for merge. Equivalent of
@@ -1753,6 +1766,19 @@ def broadcast(cast, onto, cast_on=None, onto_on=None,
         Broadcast(cast, onto, cast_on, onto_on, cast_index, onto_index)
 
 
+def _join_columns(key):
+    """
+    Return the column name(s) in a broadcast's ``cast_on`` or ``onto_on``
+    as a list, which is empty if the join uses the table index instead.
+
+    """
+    if key is None:
+        return []
+    if isinstance(key, str):
+        return [key]
+    return list(key)
+
+
 def _get_broadcasts(tables):
     """
     Get the broadcasts associated with a set of tables.
@@ -1794,10 +1820,10 @@ def get_broadcast(cast_name, onto_name):
 
         - cast: the name of the table being broadcast
         - onto: the name of the table onto which "cast" is broadcast
-        - cast_on: The optional name of a column on which to join.
-          None if the table index will be used instead.
-        - onto_on: The optional name of a column on which to join.
-          None if the table index will be used instead.
+        - cast_on: The optional name (or list of names) of a column on
+          which to join. None if the table index will be used instead.
+        - onto_on: The optional name (or list of names) of a column on
+          which to join. None if the table index will be used instead.
         - cast_index: True if the table index should be used for the join.
         - onto_index: True if the table index should be used for the join.
 
@@ -1949,10 +1975,8 @@ def merge_tables(target, tables, columns=None, drop_intersection=True):
     if columns:
         columns = list(columns)
         for c in casts.values():
-            if c.onto_on:
-                columns.append(c.onto_on)
-            if c.cast_on:
-                columns.append(c.cast_on)
+            columns.extend(_join_columns(c.onto_on))
+            columns.extend(_join_columns(c.cast_on))
 
     # get column map for which columns go with which table
     colmap = column_map(tables.values(), columns)
@@ -1981,8 +2005,8 @@ def merge_tables(target, tables, columns=None, drop_intersection=True):
                 intersection = set(onto_table.columns).\
                     intersection(cast_table.columns)
                 # intersection is ok if it's the join key
-                intersection.discard(bc.onto_on)
-                intersection.discard(bc.cast_on)
+                intersection.difference_update(_join_columns(bc.onto_on))
+                intersection.difference_update(_join_columns(bc.cast_on))
                 # otherwise drop so as not to create conflicts
                 if drop_intersection:
                     cast_table = cast_table.drop(intersection, axis=1)

orca/tests/test_mergetables.py

@@ -255,3 +255,35 @@ def test_merge_tables_dup_columns():
         'zone_id': [1, 1, 2]
     })
     assert_frames_equal(df, expected)
+
+
+def test_merge_tables_multi_column_join():
+    # cast_on / onto_on may be lists of column names, as in pandas.merge
+    df_hh = pd.DataFrame(
+        {'hh_id': [1, 1, 2],
+         'person_num': [1, 2, 1],
+         'age': [40, 38, 65]},
+        index=['pa', 'pb', 'pc'])
+    df_trips = pd.DataFrame(
+        {'hh_id': [1, 1, 2, 2, 1],
+         'person_num': [2, 1, 1, 1, 2],
+         'mode': ['walk', 'bike', 'car', 'walk', 'car']},
+        index=['ta', 'tb', 'tc', 'td', 'te'])
+
+    orca.add_table('persons', df_hh)
+    orca.add_table('trips', df_trips)
+    orca.broadcast(
+        'persons', 'trips',
+        cast_on=['hh_id', 'person_num'], onto_on=['hh_id', 'person_num'])
+
+    merged = orca.merge_tables('trips', ['persons', 'trips'])
+    expected = pd.merge(
+        df_trips, df_hh, on=['hh_id', 'person_num'], how='inner')
+    assert_frames_equal(merged, expected)
+
+    # the join columns are added automatically when a column list is given
+    merged = orca.merge_tables(
+        'trips', ['persons', 'trips'], columns=['mode', 'age'])
+    assert list(merged.columns) == ['mode', 'hh_id', 'person_num', 'age']
+    assert_frames_equal(
+        merged, expected[['mode', 'hh_id', 'person_num', 'age']])

orca/tests/test_orca.py

@@ -69,6 +69,49 @@ def test_func(test_frame):
     assert table.columns == ['a', 'b']
 
 
+def test_to_frame_column_order(df):
+    orca.add_table('test_frame', df)
+
+    @orca.column('test_frame')
+    def c(test_frame):
+        return test_frame.a + test_frame.b
+
+    @orca.column('test_frame')
+    def d(test_frame):
+        return test_frame.a * 2
+
+    table = orca.get_table('test_frame')
+
+    # requested columns come back in the requested order, whether they
+    # are local or registered columns, and no matter how many times
+    # the frame is built
+    for _ in range(5):
+        assert list(table.to_frame(['d', 'a', 'c', 'b']).columns) == \
+            ['d', 'a', 'c', 'b']
+        assert list(table.to_frame(['c', 'd']).columns) == ['c', 'd']
+        assert list(table.to_frame(['b', 'a']).columns) == ['b', 'a']
+
+    # duplicates are dropped, missing columns are ignored
+    assert list(table.to_frame(['b', 'b', 'a', 'nope']).columns) == \
+        ['b', 'a']
+
+    # with no columns given, local columns come first, then registered
+    # columns in the order they were registered
+    assert list(table.to_frame().columns) == ['a', 'b', 'c', 'd']
+
+    # function tables behave the same way
+    @orca.table()
+    def test_func(test_frame):
+        return test_frame.to_frame(['a', 'b'])
+
+    @orca.column('test_func')
+    def e(test_func):
+        return test_func.a
+
+    assert list(orca.get_table('test_func').to_frame(['e', 'b']).columns) \
+        == ['e', 'b']
+
+
 def test_table_func_cache(df):
     orca.add_injectable('x', 2)
 
@@ -512,11 +555,18 @@ def test_column_map_none(fta, ftb):
 def test_column_map(fta, ftb):
     result = orca.column_map([fta, ftb], ['aa', 'by', 'bz'])
     assert result['a'] == ['aa']
-    assert sorted(result['b']) == ['by', 'bz']
+    assert result['b'] == ['by', 'bz']
 
     result = orca.column_map([fta, ftb], ['by', 'bz'])
     assert result['a'] == []
-    assert sorted(result['b']) == ['by', 'bz']
+    assert result['b'] == ['by', 'bz']
+
+
+def test_column_map_order(fta, ftb):
+    # columns come back in the order requested, without duplicates
+    result = orca.column_map([fta, ftb], ['bz', 'aa', 'by', 'bz'])
+    assert result['a'] == ['aa']
+    assert result['b'] == ['bz', 'by']
 
 
 def test_is_step():