Pyarrow drop a column in a nested structure

user_18152545 0
Viewed 474

I would like to drop columns in my pyarrow table that are null type. Basically NullType columns are columns where all the rows have null data. I would like to drop them since they are not used by me and they cause a conflict when I import them in Spark. Here is an exemple of how I do this right now:

import pyarrow as pa

def drop_na_columns(df):
    null_columns = []
    schema = df.schema
    for (name_, type_) in zip(schema.names, schema.types):
        if type_ == pa.null():
          null_columns.append(name_)
    return df.drop(null_columns)

def create_dataframe(list_dict: dict) -> pa.table:
    fields = set()
    for d in list_dict:
        fields = fields.union(d.keys())
    dataframe = pa.table({f: [row.get(f) for row in list_dict] for f in fields})
    return drop_na_columns(dataframe)

This works fine for most of my use cases but I also have nested structures in my tables and sometimes one column in a nested structure is a null type column. My question is to know if there is a way to identify the type of every sub-column struct, and then to drop a column in a nested structure. If any of you have an idea or has already tackled this problem I would be delighted to hear about it. Thanks

pyarrow python
1 Answers

It's reasonalby easy to identify which of your table columns are struct and which fields among these structs are null.

It's a bit harder to remove them though.

Here's an example:

struct_type = pa.struct([
    pa.field('ok', pa.string()),
    pa.field('bad', pa.null())
])

schema = pa.schema([
    pa.field('col1', pa.string()),
    pa.field('col2', struct_type),
])

table = pa.Table.from_arrays(
    [
        ["#0", "#1"],
        [
            ("OK1", None),
            ("OK2", None),
        ]
        
    ],
    schema=schema
)

You can tell that the second column is a struct by doing:

pa.types.is_struct(table.schema[1].type)
>>> True

And you can iterate through its fields:

[field.name for field in table.schema[1].type]
>>> ['ok', 'bad']

And tell which one are null:

[pa.types.is_null(field.type) for field in table.schema[1].type]
>>> [False, True]

Here's the code to recursively do what you are trying to achieve:

def get_clean_array(array):
    if pa.types.is_struct(array.type):
        if isinstance(array, pa.ChunkedArray):
            array = array.combine_chunks()
        arrays = [
            get_clean_array(array.field(index))
            for index, field in enumerate(array.type)
            if not pa.types.is_null(field.type)
        ]
        names = [
            field.name 
            for field in array.type
            if not pa.types.is_null(field.type)
        ]
        return pa.StructArray.from_arrays(arrays, names)
    else:
        return array

def get_clean_table(table):
    
    return pa.Table.from_pydict({
        field.name: get_clean_array(table[field.name])
        for field in table.schema if not pa.types.is_null(field.type)
    })
user_1095250
Related