From c92c5667207780c45d9cf677c913d28287d4edf2 Mon Sep 17 00:00:00 2001 From: Samuel Maurer Date: Wed, 16 Sep 2026 15:28:25 -0700 Subject: [PATCH 1/2] Return to_frame columns in the requested order; accept multi-column broadcast keys DataFrameWrapper.to_frame(columns) built its column selections from sets, so registered columns came back in an order that varied from call to call (#57). Keep the requested order instead, dropping duplicates, and reorder the result so local and registered columns appear as requested. column_map likewise resolves columns in the order given, so merge_tables output is predictable too. broadcast documented cast_on and onto_on as equivalents of pandas left_on and right_on, but merge_tables assumed a single column name and failed with lists (#37). Normalize the keys to lists where merge_tables collects join columns and excludes them from the overlap check, and document that lists are accepted. Bump the version to 1.9.dev3 and add changelog entries. --- HISTORY.rst | 6 +++ orca/__init__.py | 2 +- orca/orca.py | 71 ++++++++++++++++++++++------------ orca/tests/test_mergetables.py | 32 +++++++++++++++ orca/tests/test_orca.py | 54 +++++++++++++++++++++++++- 5 files changed, 137 insertions(+), 28 deletions(-) diff --git a/HISTORY.rst b/HISTORY.rst index ed8c536..eccbd4a 100644 --- a/HISTORY.rst +++ b/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 ==== diff --git a/orca/__init__.py b/orca/__init__.py index c2fac98..09b880c 100644 --- a/orca/__init__.py +++ b/orca/__init__.py @@ -1,4 +1,4 @@ -__version__ = "1.9.dev2" +__version__ = "1.9.dev3" from .orca import * diff --git a/orca/orca.py b/orca/orca.py index 171e818..c1ce4b8 100644 --- a/orca/orca.py +++ b/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,11 +386,13 @@ 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} else: df = self.local.copy() @@ -404,6 +407,10 @@ def to_frame(self, columns=None): logger): df[name] = col() + if columns is not None: + # put local and registered columns in the requested order + df = df[[c for c in columns if c in df.columns]] + return df def update_col(self, column_name, series): @@ -1401,14 +1408,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 +1749,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 +1763,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 +1817,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 +1972,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 +2002,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) diff --git a/orca/tests/test_mergetables.py b/orca/tests/test_mergetables.py index 2d22964..3ab680b 100644 --- a/orca/tests/test_mergetables.py +++ b/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']]) diff --git a/orca/tests/test_orca.py b/orca/tests/test_orca.py index cf120de..8ebdf58 100644 --- a/orca/tests/test_orca.py +++ b/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(): From 758364a24fbdd3c6a042e96ce35816e2f2bc13ab Mon Sep 17 00:00:00 2001 From: Samuel Maurer Date: Wed, 16 Sep 2026 15:44:03 -0700 Subject: [PATCH 2/2] Insert registered columns at their final position instead of reordering the frame, avoiding a copy on pandas 2 --- orca/orca.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/orca/orca.py b/orca/orca.py index c1ce4b8..27fbc14 100644 --- a/orca/orca.py +++ b/orca/orca.py @@ -393,8 +393,12 @@ def to_frame(self, columns=None): 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( @@ -405,11 +409,10 @@ def to_frame(self, columns=None): 'computing column {!r} for table {!r}'.format( name, self.name), logger): - df[name] = col() - - if columns is not None: - # put local and registered columns in the requested order - df = df[[c for c in columns if c in df.columns]] + if order is None: + df[name] = col() + else: + df.insert(order.index(name), name, col()) return df