I want to make a function that does an ETL to a list/Series and returns a Series with additional attributes and methods specific to that function. I can achieve this by creating a class to extend the Series and it works but then when I try to reassign the output from the function with the new class the updated class attributes and methods are stripped away. How can I extend the Series to have custom attributes and methods that are not stripped away when reassigning back to the dataframe?
import pandas as pd
def normalize_x(x: list, new_attribute: None):
normalized = pd.Series(['normalized_'+ i if i != 4 else None for i in x])
return NormalizeX(normalized = normalized, original = x, new_attribute = new_attribute)
class NormalizeX(pd.Series):
def __init__(self, normalized, original, new_attribute, *args, **kwargs,):
super().__init__(data = normalized, *args, **kwargs)
self.original = original
self.normalized = normalized
self.new_attribute = new_attribute
def conversion_errors(self):
return [o != n for o, n in zip(pd.isnull(self.original), pd.isnull(self.normalized))]
df = pd.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6], "C": ['dog', 'cat', 4]})
out = normalize_x(df.C, new_attribute = 'CoolAttribute')
out
## 0 normalized_dog
## 1 normalized_cat
## 2 None
## dtype: object
## Can still use Series methods
out.to_list()
## ['normalized_dog', 'normalized_cat', None]
## Can use the new methods and access attributes
out.conversion_errors()
## [False, False, True]
out.original
##0 dog
##1 cat
##2 4
##Name: C, dtype: object
df['new'] = normalize_x(df.C, new_attribute = 'CoolAttribute')
df['new']
## 0 normalized_dog
## 1 normalized_cat
## 2 None
## dtype: object
## Can't use the new methods or access attributes
df['new'].conversion_errors()
## AttributeError: 'Series' object has no attribute 'conversion_errors'
df['new'].original
## AttributeError: 'Series' object has no attribute 'original'
It's too difficult for me to implement your desired functionality, so I only share what I found in my investigation expecting it might be useful for other answerers.
The reason why you got those attribute errors is sanitization processes were carried out on the Series which you passed to the DataFrame.
ids:You can quickly confirm the difference between what out and df['new'] refer to by the following code:
out = normalize_x(df.C, new_attribute = 'CoolAttribute')
df['new'] = out
print(id(out))
print(id(df['new']))
1861777917792
1861770685504
you can see out and df['new'] are different from each other because of this id difference.
Let's dive into the pandas source code to see what goes on here.
DataFrame._set_item method:In the definition of the DataFrame class, _set_item method works when you try to add Series to DataFrame in a specified column.
def _set_item(self, key, value) -> None:
"""
Add series to DataFrame in specified column.
If series is a numpy-array (not a Series/TimeSeries), it must be the
same length as the DataFrames index or an error will be thrown.
Series/TimeSeries will be conformed to the DataFrames index to
ensure homogeneity.
"""
value = self._sanitize_column(value)
In this method, value = self._sanitize_column(value) in the first line except the docstring. This _sanitize_column method actually destroyed your original Series functionality. If you dig this method deeper, you'll finally reach the following lines:
def _reindex_for_setitem(value: FrameOrSeriesUnion, index: Index) -> ArrayLike:
# reindex if necessary
if value.index.equals(index) or not len(index):
return value._values.copy()
value._values.copy() is the direct cause of the disappearance of the NormalizeX attributes. It just copies the values from the given Series.
Therefore, the _set_item method should be modified in order to protect the NormalizeX attributes.
You have to override the DataFrame class to set your NormalizeX in a specified column keeping with its attributes.