DataFrame apply results in a DataFrame when lambda returns a list of length equal to the number of DataFrame features · Issue #16353 · pandas-dev/pandas (original) (raw)
Hello guys, I'm using the apply method with axis=1 argument and returning a list for each row. The weird thing is when the length of this list is equal to the number of features in the original dataframe, this apply returns another dataframe, and a series if anything else. I don't know if this is intended, but it was very unexpected.
import numpy as np
import pandas as pd
df = pd.DataFrame(np.random.randn(6, 3), columns=['A', 'B', 'C'])
ans1 = df.apply(lambda x: [1, 2, 3], axis=1)
ans2 = df.apply(lambda x: [1, 2], axis=1)
assert type(ans1) is pd.DataFrame
assert type(ans2) is pd.Series