with_columns for Pandas

[jernejfrank]

Creating the option to use with_columns on Pandas data frames.

Partially addressing #1158.

Changes

• added new decorator

How I tested this

• unit tests
• e2e
• example

Notes

Known issues:

• config.when does not work for individual nodes within the with_columns subtag.

The other thing is, at the moment, this extracts relevant columns from the df and then appends them. This would then be an eager execution.

Polars and PySpark can do with_columns lazy and use their optimizers under the hood. Pandas to my knowledge doesn't have that, but there are two ways how we could mimic laziness:

1. using eval

df.eval('new_column = f(previous_columns)'))

2. usingassign

df.assign(new_column = lambda x: x(previous_columns))

My reasoning so far is, that anybody who is concerned with latency will not be using Pandas in the first place so this seems like an overkill.

Checklist

• PR has an informative and human-readable title (this will be pulled into the release notes)
• Changes are limited to a single goal (no scope creep)
• Code passed the pre-commit check & code is left cleaner/nicer than when first encountered.
• Any change in functionality is tested
• New functions are documented (with a description, list of inputs, and expected output)
• Placeholder code is flagged / future TODOs are captured in comments
• Project documentation has been updated if adding/changing functionality.

[ellipsis-dev]

❌ Changes requested. Reviewed everything up to 0c41502 in 50 seconds

More details

• Looked at 755 lines of code in 7 files
• Skipped 2 files when reviewing.
• Skipped posting 4 drafted comments based on config settings.

1. tests/function_modifiers/test_recursive.py:778

• Draft comment:
  Typo in function name. Consider renaming substract_1_from_2 to subtract_1_from_2 for clarity and correctness. This applies to other occurrences as well.
• Reason this comment was not posted:
  Marked as duplicate.

2. hamilton/function_modifiers/recursive.py:778

• Draft comment:
  Typo in function name. Consider renaming substract_1_from_2 to subtract_1_from_2 for clarity and correctness. This applies to other occurrences as well.
• Reason this comment was not posted:
  Marked as duplicate.

3. tests/resources/with_columns.py:10

• Draft comment:
  Typo in function name substract_1_from_2. It should be subtract_1_from_2. This typo is present in multiple places, including the function definition and its usage in decorators.
• Reason this comment was not posted:
  Marked as duplicate.

4. tests/function_modifiers/test_recursive.py:778

• Draft comment:
  Typo in function name substract_1_from_2. It should be subtract_1_from_2. This typo is present in multiple places, including the function definition and its usage in decorators.
• Reason this comment was not posted:
  Marked as duplicate.

Workflow ID: wflow_tdIW9EERoU9TNVSK

---

Want Ellipsis to fix these issues? Tag @ellipsis-dev in a comment. You can customize Ellipsis with 👍 / 👎 feedback, review rules, user-specific overrides, quiet mode, and more.

[ellipsis-dev]

Typo in function name. Consider renaming substract_1_from_2 to subtract_1_from_2 for clarity and correctness.

Suggested change

	def substract_1_from_2(col_1: pd.Series, col_2: pd.Series) -> pd.Series:
	def subtract_1_from_2(col_1: pd.Series, col_2: pd.Series) -> pd.Series:

[jernejfrank]

Learned something new today:
Substract was formerly used in analogy with abstract. But in modern usage, it is written according to the Latin, subtract.

[zilto]

Could use

from hamilton import ..., node as hamilton_node

because the name node is overriden elsewhere (e.g., line 187)

[elijahbenizzy]

Looking good! Main comments about how to do the assign-style operations. Definitely on the right track!

[elijahbenizzy]

This should probably live in a pandas extension -- while pandas has, historically, been the dependency for Hamilton, we want to move it out (and this is very pandas-specific logic)

[elijahbenizzy]

Yep, would make this live here and have pandas/spark refer to it

[elijahbenizzy]

Might be worth expanding in the docstrings a bit more about what this means -- I'm getting a little lost int he details (E.G. is this a user error? What's the remediation?) More for internal docs.

[elijahbenizzy]

So I think we should keep this consistenct with the h_spark with_columns -- maybe we add a single string ability for that? Or we remove it for this?

[jernejfrank]

Will make it consistent with spark -- in case this becomes a pain point can always add it later on.

[elijahbenizzy]

This should solve the problem of config wiring through...

[elijahbenizzy]

Let's remote this for now? Should also be able to remove from h_spark -- it's an artifact and was left in erroneously.

[elijahbenizzy]

Ooh we should be able to share logic with extract_columns... Might be good to leave this as a TODO note

[elijahbenizzy]

SHould be able to validate the type for the parameter as well?

[elijahbenizzy]

Hmm, we're mutating this... We should be careful as that breaks assumptions.

The nice thing about pandas is that the dataframes shouldn't mess with memory if you create new ones, the series are what takes up the memory. So we should be able to continually modify it by creating new ones (I think...)

[elijahbenizzy]

Might be worth using assign or something like that... Looking here (you mentioned similar things in your comments on the code): https://stackoverflow.com/questions/72291290/how-to-create-new-column-dynamically-in-pandas-like-we-do-in-pyspark-withcolumn

[zilto]

Looking at the way Narwhals implements the Polars .with_columns() for other backends can also be informative.

[jernejfrank]

I looked a bit deeper:

1. assign does a copy of the data frame. (It should be shallow unless the flag is set to copy="deep". From this thread my impression is that assign was created for chaining operations together and leaving the original data frame intact.
2. Narwhals uses concat(dfs, axis=1, **extra_kwargs) under the hood and this does have a warning: "concat() makes a full copy of the data, and iteratively reusing concat() can create unnecessary copies. Collect all DataFrame or Series objects in a list before using concat()." However, for our use case I guess we can do this only once so should also be ok.

And tested it -- seems like concat is the way to go in terms of performance:

[EDIT] Concat is annoying because you would need to filter out individual columns that have changed / are provided as selected by the user and safeguard since concat would just duplicate them in the data frame. .assign() seems to do this automatically (probably where the discrepancy is coming from).

%timeit pd.concat([df:(10000, 5000), cols:2000], axis=1)
280 ms ± 22.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit df:(10000, 5000).assign(**cols:2000)
4.44 s ± 765 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%memit pd.concat([df:(10000, 5000), cols:2000], axis=1)
peak memory: 2107.36 MiB, increment: 0.20 MiB

%memit df:(10000, 5000).assign(**cols:2000)
peak memory: 2525.92 MiB, increment: 37.08 MiB

%load_ext memory_profiler
import numpy as np
import pandas as pd
import warnings  
warnings.simplefilter(action='ignore', category=pd.errors.PerformanceWarning)

pd_size = 10000
df = pd.DataFrame(np.random.randint(0,pd_size,size=(pd_size, pd_size//2)))#, columns=list('ABCD'))

cols = []
for i in range(pd_size//10):
    cols.append(np.random.rand(pd_size))

concat_cols = [pd.Series(col) for col in cols]
assign_cols = {f"{pd_size+i+1}":col for i,col in enumerate(cols)}


print(f'\n%timeit pd.concat([df:{df.shape}, cols:{len(concat_cols)}], axis=1)')
temp_df = [df.copy()]
temp_df.extend(concat_cols)
%timeit l=pd.concat(temp_df, axis=1)

print(f'\n%timeit df:{df.shape}.assign(**cols:{len(assign_cols)})')
temp_df = df.copy()
%timeit temp_df.assign(**assign_cols)


print(f'\n%memit pd.concat([df:{df.shape}, cols:{len(concat_cols)}], axis=1)')
temp_df = [df.copy()]
temp_df.extend(concat_cols)
%memit l=pd.concat(temp_df, axis=1)

print(f'\n%memit df:{df.shape}.assign(**cols:{len(assign_cols)})')
temp_df = df.copy()
%memit temp_df.assign(**assign_cols)

[elijahbenizzy]

Other tests you may want to consider

1. Inputs from outside the subdag
2. Inputs from external stuff
3. Config

[elijahbenizzy]

THis looks good -- let's rebase, I'll take one more look tomorrow morning, then we can merge and release!