Retain Custom Attributes & Methods of Pandas Series SubClass when assigning to DataFrame column

Viewed 463

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?

Custom function that does ETL, returns a Series with an extended class

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]})

Assign to a new object (new attributes and methods work)

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

Assign to a Pandas DataFrame (new attributes and methods break)

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'
2 Answers

Pandas allows you to extend its classes (Series, DataFrame). In your case, the solution is quite verbose, but I think it's the only way you have to reach your goal.

I try to get to the point without analyse complex cases, so the complete implementation of the interface is up to you, but I can give you an idea of what you can use.

I don't understand the utility of new_attribute, so for the moment is not taken into account. Basically, Pandas extensions let you extend only a 1-D array for what I know. Since you have multiple arrays (both normalized, original) you have to create another data type to get around the problem.

class NormX(object):
    def __init__(self, normalized, original):
        self.normalized = normalized
        self.original = original

    def __repr__(self,):
        if self.normalized is None:
            return 'Nan'
        return self.normalized

This allows you to create a simple base object like the following:

norm_obj = NormX('normalized_dog', 'dog')

This object will be the basic block of your custom array. To be able to exploit this this kind of class, you have to register a new type in Pandas:

from pandas.api.extensions import ExtensionDtype, register_extension_dtype
import numpy as np

@pd.api.extensions.register_extension_dtype
class NormXType(ExtensionDtype):
    name = 'normX'
    type = NormX
    kind = 'O'
    na_value = np.nan

Now you have all the elements to build a custom array based on Pandas framework. To do that, you have to extend its interface class named ExtensionArray. Here you can find the abstract methods that must be implemented by subclasses. I gave you a very basic implementation, but it should be declared in a proper way:

from pandas.api.extensions import ExtensionArray

class NormalizeX(ExtensionArray):
    def __init__(self, values):
        self.data = values
        
    def __repr__(self,):
        return "NormalizeX({!r})".format([(t.normalized, t.original) for t in self.data])
    
    def _from_sequence(self,):
        pass
    
    def _from_factorized(self,):
        pass
    
    def __getitem__(self, key):
        return self.data[key]
    
    # def __setitem__(self, key, value):
    #     self.normalized[key] = value
    #     return self
    
    def __len__(self,):
        return len(self.data)
    
    def __eq__(self, other):
        return False
    
    def dtype(self,):
        # return self._dtype
        return object
    
    def nbytes(self,):
        return sys.getsizeof(self.data)
    
    def isna(self,):
        return False
    
    def take(self,):
        pass
    
    def copy(self,):
        return type(self)(self.data)
    
    def _concat_same_type(self,):
        pass

More over, to define a custom method on that class, you have to define a custom Series accessor, as follows:

@pd.api.extensions.register_series_accessor("normx_accsr")
class NormalizeXAccessor:
    def __init__(self, obj):
        self.normalized = [o.normalized for o in obj]
        self.original = [o.original for o in obj]
        
    @property
    def conversion_errors(self):
        return [o != n for o, n in zip(pd.isnull(self.original), pd.isnull(self.normalized))]

In this way, the NormalizeX custom array implements all the requested methods to be successfully integrated into both Series and DataFrame. So, your example simply reduces to:

df = pd.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6], "C": ['dog', 'cat', 4]})
norm_list = [['normalized_'+ i, i] if i != 4 else (None, i) for i in df.C]
normx_objs = [NormX(*t) for t in norm_list]
normalizedX = NormalizeX(normx_objs)

df['new'] = normalizedX

# Use the previously defined attribute_accessor normx_accsr
df['new'].normx_accsr.conversion_errors # [False, False, True]

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 cause of the issue:

The reason why you got those attribute errors is sanitization processes were carried out on the Series which you passed to the DataFrame.

A brief check of 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.

Conclusion:

You have to override the DataFrame class to set your NormalizeX in a specified column keeping with its attributes.

Related