You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "frreiss (via GitHub)" <gi...@apache.org> on 2024/03/11 21:08:06 UTC

Re: [I] [Python] Explode array column [arrow]

frreiss commented on issue #27923:
URL: https://github.com/apache/arrow/issues/27923#issuecomment-1989447329

   +1 to implementing this feature in Arrow.
   
   In the meantime, I appreciate the workaround. Here's a version that optionally retains information about offsets of values in the original lists:
   ```python
   def explode_table(table, column, list_ix_column: str=None):
       """
       :param table: Input table to explode
       :param column: Name of column (must be unique and list-valued) to explode
       :param list_ix_column:  name of new column in which to store offsets
                into the original lists; or ``None`` to skip generating that 
                column.  
       """
       null_filled = pc.fill_null(table[column], [None])
       flattened = pc.list_flatten(null_filled)
       other_columns = list(table.schema.names)
       other_columns.remove(column)
       if len(other_columns) == 0:
           result = pa.table({column: flattened})
       else:
           indices = pc.list_parent_indices(null_filled)
           result = table.select(other_columns).take(indices)
           result = result.append_column(
               pa.field(column, table.schema.field(column).type.value_type),
               flattened,
           )
   
       if list_ix_column is not None:
           # Generate indices into original lists, aligned with the unrolled
           # lists generated above.
           # We do this by mapping the list length to a categorical type, where
           # value k of the categorical type is a list from 0 to k.
           list_lengths = pc.list_value_length(table.column(column))
           max_list_len = pc.max(list_lengths).as_py()
           dictionary = [
               pa.scalar(list(range(l)), type=pa.list_(pa.int32()))
               for l in range(max_list_len + 1)
           ]
           id_lists = (
               pa.DictionaryArray.from_arrays(list_lengths, dictionary)
               .dictionary_decode())
           list_ix = pc.list_flatten(id_lists)
   
           # Add column to the left of the unrolled list elements.
           result = result.add_column(
               result.num_columns - 1,
               pa.field(list_ix_column, pa.int32()), list_ix
           )
           
       return result 
   
   ```


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org