Skip to content

Add DataFrame.insert_columns #231

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 6 commits into from
Closed
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 17 additions & 3 deletions spec/API_specification/dataframe_api/dataframe_object.py
Original file line number Diff line number Diff line change
Expand Up @@ -180,9 +180,9 @@ def get_rows_by_mask(self, mask: Column[Bool]) -> DataFrame:
"""
...

def insert_column(self, column: Column[Any]) -> DataFrame:
def insert_columns(self, columns: Column[Any] | Sequence[Column[Any]]) -> DataFrame:
"""
Insert column into DataFrame at rightmost location.
Insert column(s) into DataFrame at rightmost location.

The column's name will be used as the label in the resulting dataframe.
To insert the column with a different name, combine with `Column.rename`,
Expand All @@ -203,9 +203,23 @@ def insert_column(self, column: Column[Any]) -> DataFrame:
df = df.insert_column(new_column.rename('a_plus_1'))
df = df.get_columns_by_name(new_column_names)

If inserting multiple columns, then the order in which they are inserted
is not guaranteed and may vary across implementations. For example, the
following

.. code-block:: python

new_column_1 = df.get_column_by_name('a').rename('b')
new_column_2 = (new_column_1 + 2).rename('c')
df.insert_columns([new_column_1, new_column_2])

is not supported, as `new_column_2` is derived from `new_column_1`, which may
not be part of the dataframe if `new_column_2` is inserted first.
Comment on lines +206 to +217
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As of now our Column objects are immutable and have no concept of being derived from. Could you explain why this wouldn't be supported?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Say you start with

a b
1 4
1 5
2 6

and you want to add two new columns:

  • c: which is b + 1
  • d: which is a + c

These two new columns can't be added in any order, they need to be added first 'c' and then 'd'

If you try inserting 'd' before 'c', then you could get "KeyError: column not found, 'c' not part of dataframe"

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

e.g.

In [21]: df = pl.DataFrame({'a':[1,1,3], 'b':[4,5,6]})

In [22]: df.with_columns(c=pl.col('a')+1, d=pl.col('c')+pl.col('a'))
---------------------------------------------------------------------------
ColumnNotFoundError: c

Error originated just after this operation:
DF ["a", "b"]; PROJECT */2 COLUMNS; SELECTION: "None"

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Okay that makes sense as to why in a lazy implementation you can't use derived columns from other columns being inserted. Assuming that isn't allowed, then can inserting the columns in the order they are passed possible?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

eager too, the above example was eager

Should be possible, yes (we could always call .get_columns_by_name internally and reoder) I was just thinking about how to get the concept across without saying "independent"

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The need for a namespace.col seems natural for lazy implementations like Polars, but not so much for eager implementations as far as I can tell. So I do think this is a lazy vs eager thing.

I'm getting the sense that conflating lazy and eager into a single API is going to be the source of some tension (and likely a bad UX).

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also, as another point of reference, Ibis does support this:

import ibis
import polars as pl

con = ibis.polars.connect()

df = pl.DataFrame({'a':[1,1,3], 'b':[4,5,6]})
idf = con.create_table("df", df)

c = (idf['a'] + 1).name('c')
d = (c + idf['a']).name('d')

idf.select(['a', 'b', c, d])

Where it resolves it to:

Selection[r0]
  selections:
    a: r0.a
    b: r0.b
    c: r0.a + 1
    d: r0.a + 1 + r0.a

I'm guessing this is a design difference between Ibis and Polars-lazy since Ibis allows selecting columns, i.e. idf['a'] which returns a expression Column type of object which is bound to its owning table, whereas Polars expressions aren't bound to a table?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does Ibis insert the columns in parallel in select? If not, I think that's probably the underlying reason

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ibis generates SQL (or DataFrame code) and hands that generated code to the connected backend, so it's an implementation detail of the backend. But in general, Ibis binds its expressions to Tables which allows writing code like I did above which isn't dependent on a column named "c". On the other hand, the Polars code using expressions you posted isn't bound to a table so it is dependent on having a column named "c".

I think this is again a discussion more about expressions and Columns and what we do in the standard regarding them than this specific function at this point. Should we move discussion back to #229?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'll work on updating the Expressions proposal this week and next, but just wanted to note that that will make the independence condition very easy to state:

all expressions passed to insert_columns must have root names already present in the dataframe

I'll define as part of the proposal exactly what I mean by "root names"


Parameters
----------
column : Column
columns : Column | Sequence[Column]
Column(s) to insert. Insertion order is not guaranteed.
"""
...

Expand Down