Skip to content
Tutoriais
PySpark
How to Convert PySpark DataFrame Column to Python List?

Como converter uma coluna de um DataFrame PySpark em uma lista em Python?

Se você está trabalhando com dataframes PySpark, pode ser necessário extrair uma coluna do dataframe e usá-la como uma lista em Python para análises adicionais. Neste artigo, exploraremos várias maneiras de extrair uma coluna de um dataframe PySpark para uma lista em Python. Começaremos com o básico do dataframe PySpark e depois mergulharemos em diversos métodos para extrair colunas.

Quer criar visualizações de dados rapidamente a partir de um dataframe Pandas em Python, sem escrever código?

PyGWalker é uma biblioteca em Python para Análise Exploratória de Dados com Visualizações. PyGWalker (opens in a new tab) pode simplificar sua análise de dados e fluxo de trabalho de visualização de dados no Jupyter Notebook, transformando seu dataframe pandas (e dataframe polares) em uma interface de usuário no estilo Tableau para exploração visual.

PyGWalker para visualização de dados (opens in a new tab)

Introdução ao DataFrame PySpark

Um DataFrame é uma coleção distribuída de dados organizados em colunas nomeadas. É conceitualmente equivalente a uma tabela em um banco de dados relacional ou a um dataframe em R ou Python. O DataFrame PySpark é um conjunto de dados distribuído e resiliente (RDD) de linhas que possui um esquema. O PySpark SQL fornece uma interface de programação para trabalhar com dados estruturados usando o Spark. O PySpark suporta a maioria das fontes de dados comuns, como CSV, JSON, Avro, Parquet, etc.

Para entender melhor, vamos criar um simples dataframe PySpark e verificar seu esquema usando o PySpark SQL.

from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, IntegerType, StringType
 
# Crie a SparkSession
spark = SparkSession.builder.appName("Exemplos_PySpark").getOrCreate()
 
# Defina o esquema
esquema = StructType([
    StructField("id", IntegerType(), True),
    StructField("nome", StringType(), True)])
 
# Crie o dataframe
dados = [(1,"João"),(2,"Maria"),(3,"Smith"),(4,"James")]
df = spark.createDataFrame(dados, schema=esquema)
 
# Mostre o dataframe
df.show()

A saída será como:

+---+-----+
| id| nome|
+---+-----+
|  1| João|
|  2|Maria|
|  3|Smith|
|  4|James|
+---+-----+

Extraindo uma única coluna como uma lista

Existem várias maneiras de extrair uma coluna do dataframe PySpark. Vamos explorar algumas delas nesta seção.

Método 1: Usando a função collect()

A função collect() no PySpark é usada para retornar todos os elementos do RDD (Resilient Distributed Datasets) para o programa driver como um array. Podemos usar collect() para converter uma coluna do dataframe PySpark em uma lista em Python. Veja como:

# Extraia a coluna nome usando o collect()
lista_nomes = df.select('nome').rdd.flatMap(lambda x: x).collect()
 
# Imprima a lista
print(lista_nomes)

A saída será como:

[u'João', u'Maria', u'Smith', u'James']

Aqui, usamos a função select() para selecionar a coluna "nome" do dataframe. Em seguida, usamos rdd.flatMap(lambda x: x) para converter a coluna em um RDD e, em seguida, usamos a função collect() do RDD para obter os dados na forma de uma lista em Python.

Método 2: Usando Compreensão de Lista

Outra maneira de extrair uma coluna de um dataframe PySpark como uma lista em Python é usar compreensão de lista. Veja como:

# Extraia a coluna nome usando compreensão de lista
lista_nomes = [linha.nome for linha in df.select('nome').collect()]
 
# Imprima a lista
print(lista_nomes)

A saída será como:

[u'João', u'Maria', u'Smith', u'James']

Aqui, usamos a compreensão de lista para extrair a coluna "nome" do dataframe como uma lista em Python. Primeiro, usamos a função select() para extrair a coluna e, em seguida, usamos a função collect() para recuperar os dados de volta para o programa driver.

Método 3: Usando a função toPandas()

Também podemos extrair uma coluna do dataframe PySpark como uma lista em Python usando a função toPandas(). No entanto, este método não é recomendado para dataframes grandes, pois pode causar erros de falta de memória. Veja como:

# Extraia a coluna nome usando toPandas()
lista_nomes = df.select('nome').toPandas()['nome'].tolist()
 
# Imprima a lista
print(lista_nomes)

A saída será como:

['João', 'Maria', 'Smith', 'James']

Aqui, usamos a função select() para selecionar a coluna "nome" do dataframe e depois convertemos o dataframe em um dataframe do Pandas usando a função toPandas(). Por fim, usamos a função tolist() para converter a série do Pandas em uma lista em Python.

Conclusão

Neste artigo, exploramos vários métodos para extrair uma coluna do dataframe PySpark para uma lista em Python, como usar a função collect(), compreensão de lista e a função toPandas(). O PySpark fornece uma maneira eficiente de processar grandes conjuntos de dados em paralelo, distribuindo o trabalho em vários nós em um cluster. Entender o dataframe PySpark e como extrair dados dele é uma habilidade valiosa para qualquer cientista de dados ou engenheiro que trabalhe com conjuntos de dados grandes.

Esperamos que este artigo seja útil e que você agora saiba como extrair uma coluna de um dataframe PySpark para uma lista em Python. Se você deseja aprender mais sobre PySpark e Pandas, confira nossos outros tutoriais.


Links: