pandas es una herramienta magnífica para analizar pequeños datasets en una única máquina. Cuando necesitamos trabajar con datasets más grandes, los Data Scientists normalmente escogemos hacerlo usando PySpark. Sin embargo, convertir código de pandas a PySpark no es sencillo, ya que las APIs de ambas herramientas son considerablemente distintas.
En este post presentamos Koalas, una librería que hace la curva de aprendizaje para pasar de pandas a PySpark mucho más fácil. Con Koalas, los Data Scientists pueden aprovecharse de las ventajas de PySpark con un esfuerzo mínimo, y de esta forma crear valor más rápido.
pandas
pandas es considerada una de las cuatro librerías más importantes para trabajar en el ecosistema Data Science de Python, junto a NumPy, matplotlib y scikit-learn. El nombre de pandas viene de “Python Data Analysis Library”.

pandas es una herramienta ideal para hacer Data Wrangling. Tiene una enorme cantidad de funciones para manipular y realizar análisis estadístico sobre datos utilizando los DataFrames, tales como:
- Indexar, renombrar, ordenar y unir DataFrames.
- Modificar la definición de un DataFrame actualizando, añadiendo o borrando columnas.
- Limpiar y preparar datos imputando datos ausentes.
- Calcular métricas como máximo, mínimo, media, suma, etc.
Inconvenientes de pandas
Desafortunadamente, pandas no es totalmente escalable, ya que todas las operaciones realizadas se ejecutan en un único nodo. En caso de estar trabajando con volúmenes muy grandes de datos, la ejecución de pandas es muy lenta y puede incluso llegar a ser imposible si la máquina no cuenta con suficiente memoria.
Para solventar este problema, es posible utilizar Apache Spark como alternativa.
Apache Spark – PySpark
Apache Spark es un framework open-source para procesamiento distribuido en clúster. Por otro lado, la API de Spark implementada en Python se denomina PySpark.

En la versión 1.3 de Spark se introdujeron los DataFrames, los cuales están inspirados en los DataFrames de pandas y ofrecen funcionalidades similares. En Spark es posible ejecutar operaciones en paralelo utilizando diferentes nodos de un cluster, lo cual hace que esta herramienta sea mucho más adecuada para procesar rápidamente grandes datasets.
Desventajas de PySpark
Las sintaxis de pandas y PySpark son muy diferentes. PySpark está muy influenciado por la sintaxis de SQL, mientras que pandas sigue un esquema más similar a la sintaxis clásica de Python. Para la mayoría de Data Scientists acostumbrados a pandas, escribir código en Spark resulta más engorroso y menos intuitivo, y la curva de aprendizaje resulta complicada.
A continuación, mostramos dos fragmentos de código para realizar exactamente las mismas operaciones, uno escrito con pandas y el otro con PySpark, para comparar ambas opciones:


Como podemos observar, el DataFrame que obtenemos como resultado es idéntico con pandas y PySpark.
Koalas
Koalas es una API de pandas programada sobre Apache Spark.

Tiene todas las ventajas de la implementación de los DataFrames de Spark para trabajar en clúster, pero utilizando la sintaxis de pandas.
Beneficios de Koalas
Koalas facilita que los Data Scientists familiarizados con pandas puedan trabajar casi de inmediato en entornos Big Data con Spark, aprendiendo a hacerlo de forma mucho más rápida y sencilla.
Además, permite emplear una única librería base para trabajar con conjuntos de datos de cualquier tamaño, en vez de tener que utilizar pandas para datasets pequeños y PySpark para grandes.
Instalación de Koalas
Para usar Koalas necesitamos tener instalado en nuestra máquina Python 3.6 o superior y PySpark. Luego podremos instalar Koalas tanto con conda como con pip. En este enlace se incluyen las instrucciones exactas.
También podemos utilizar Koalas con un clúster de Databricks, ya que está incluido en la versión DataBricks Runtime 7.3 y superiores, y se puede instalar como una librería de PyPI en las versiones 7.0 e inferiores.
Ejemplos de uso de Koalas
En esta sección vamos a mostrar varios ejemplos de código con Koalas para implementar tareas de Data Wrangling muy habituales en los proyectos de Data Science. Como podremos ver, la sintaxis es idéntica a la que usaríamos con pandas (cambiando el alias “pd” por “ks”, en nuestro caso):
Librerías necesarias
Importamos pandas, NumPy, matplotlib, seaborn, PySpark y Koalas:
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
%matplotlib inline
import findspark
findspark.init()
import pyspark
from pyspark import SparkContext, SparkConf
from pyspark.sql import SparkSession
sc = SparkContext.getOrCreate()
spark = SparkSession(sc)
import databricks.koalas as ks
Creación de objetos
Series
Un objeto Series
de Koalas se puede crear pasándole una lista de valores, de la misma forma que con las Series
de pandas. También es posible utilizar una Series
de pandas.
# Create a Koalas Series
kser = ks.Series([1, 3, 5, np.nan, 6, 8])
# Create a Koalas Series by passing a pandas Series
pser = pd.Series([1, 3, 5, np.nan, 6, 8])
kser = ks.Series(pser)
DataFrame
Podemos crear DataFrames
de Koalas de la misma manera que con pandas. De la misma manera, un DataFrame
de Koalas se puede componer a partir de un DataFrame
de pandas.
# Create a Koalas DataFrame
kdf = ks.DataFrame({'A': np.random.rand(5), 'B': np.random.rand(5)})
# Create a Koalas DataFrame by passing a pandas DataFrame
pdf = pd.DataFrame({'A': np.random.rand(5),
'B': np.random.rand(5)})
kdf = ks.DataFrame(pdf)
Entrada y salida de datos
Leer fichero CSV como DataFrame
Los ficheros CSV son muy claros y fáciles de emplear.
ks_df = ks.read_csv("yellow-small.csv", sep=',')
ks_df.head()

Escribir DataFrame
a fichero CSV
ks_df.to_csv("yellow-small-pandas.csv", index=False)
Leer fichero Parquet como DataFrame
Parquet es un formato de fichero eficiente y compacto diseñado para leer y escribir datos de forma más rápida.
ks_df = ks.read_parquet("yellow-small.parquet")
ks_df.head()

Escribir DataFrame
a fichero Parquet
ks_df.to_parquet("yellow-small-pandas.parquet")
Visualización de datos
Ver las primeras n filas de un DataFrame
head()
suele generar confusión al pasar de pandas a PySpark, debido a que se comporta de manera diferente en cada uno. En el caso de Koalas, funciona de la misma forma que en pandas, usando la función limit()
de PySpark internamente.
ks_df.head()

Devolver index de un DataFrame
ks_df.index

Devolver columnas de un DataFrame
ks_df.columns

Tipos de datos de las columnas de un DataFrame
ks_df.dtypes

Tamaño de un DataFrame
El atributo shape de un DataFrame
contiene el número de filas y columnas, de la misma forma que en pandas.
ks_df.shape

Transformación de datos
Selección de una columna de un DataFrame
De la misma forma que con los DataFrames
de pandas, podemos seleccionar una columna de un DataFrame
de Koalas, lo que nos devolverá una Series
.
ks_df['color'] # equivalente a ks_df.color

Selección de varias columnas de un DataFrame
Asimismo, podemos seleccionar múltiples columnas de un DataFrame
de Koalas (nos devuelve un DataFrame
de Koalas).
ks_df[['color', 'size']]

Acceder a filas y columnas específicas de un DataFrame
Las funciones de slicing loc
e iloc
están disponibles en Koalas:
ks_df.loc[1]

ks_df.iloc[:3, 1:2]

Renombrar columnas de un DataFrame
ks_df_renamed=ks_df.rename(columns={"color":"COLOR"})
ks_df_renamed.head()

Crear nuevas columnas de un DataFrame
ks_df['size_2'] = ks_df['size'] * 2
ks_df.head()

Eliminar columnas de un DataFrame
ks_df_drop = ks_df.drop(['color'], axis=1)
ks_df_drop.head()

Aplicar operaciones y sustituir valores sobre un DataFrame
de Koalas
La función apply()
se utiliza para aplicar una función sobre un eje de un DataFrame
o a los valores de una Series
:
ks_df_apply = ks_df.copy()
ks_df_apply['size_apply'] = ks_df['size'].apply(lambda x: x ** 2)
ks_df_apply.head(10)

applymap()
se usa para aplicar una función sobre todos los elementos de un DataFrame
:
ks_df_len = ks_df.applymap(lambda x: len(str(x)))
ks_df_len.head(10)

map()
sirve para sustituir cada valor de un objeto Series
por otro valor:
ks_df_map = ks_df.copy()
ks_df_map['color_minus'] = ks_df_map['color'].map({'YELLOW': 'yellow', 'PURPLE': 'purple'})
ks_df_map.sample(frac=1/5)

Exploración de datos
Ordenar datos de un DataFrame
ks_df.sort_values(by='size')

Filtrar datos de un DataFrame
ks_df[(ks_df['size'] == 1) & (ks_df['inflated'] == 'T')]

Group By y Aggregation sobre un DataFrame
ks_df.groupby('color').max()['size']

Contar número de valores de las categorías de una columna de un DataFrame
ks_df['color'].value_counts()

Calcular métricas de un DataFrame
ks_df['size'].mean(), ks_df['size'].min(), ks_df['size'].max(), ks_df['size'].median()

ks_df.describe()

Calcular variables dummy de un DataFrame
ks_df_dummies = ks.get_dummies(ks_df)
ks_df_dummies.head()

Calcular matriz de correlación de un DataFrame
plt.figure(figsize=(6,6))
ks_correlation_matrix=ks_df_dummies.corr()
pd_correlation_matrix=ks_correlation_matrix.to_pandas()
sns.heatmap(pd_correlation_matrix,xticklabels=list(pd_correlation_matrix),yticklabels=list(pd_correlation_matrix), annot=True,fmt='.1f',linewidths=0.5,cmap="coolwarm")

Limpieza de datos
Identificar datos ausentes de un DataFrame
ks_df_nulls = ks.DataFrame({'a': [1, None, 3], 'b': [4, 5, 6]})
ks_df_nulls.head()

for col in ks_df_nulls:
print(col)
print(ks_df_nulls[ks_df_nulls[col].isnull()])
print()

Imputar datos ausentes de un DataFrame
ks_df_nulls['a'].fillna(0)

Eliminar datos ausentes de un DataFrame
ks_df_dropna=ks_df_nulls.dropna()
ks_df_dropna.head()

Conversión de DataFrames
entre librerías
En caso de encontrarnos con alguna función de pandas no disponible en Koalas, o si queremos convertir un DataFrame
de Koalas a PySpark para trabajar con él utilizando alguna librería de Spark (como Spark MLlib), podemos utilizar las funciones de conversión de DataFrames
. También es posible convertir un DataFrame
de pandas o PySpark a Koalas de manera muy sencilla.
Lo único que debemos tener en cuenta al convertir un DataFrame
de Koalas a pandas es que, si el tamaño del DataFrame
es muy grande, la conversión puede dar lugar a un error out-of-memory (lo mismo que ocurre con PySpark).
Transformar DataFrame
de pandas a Koalas
pd_df=pd.DataFrame({'a'[1,2,3],'b':[4,5,6]})
ks_df_from_pd=ks.from_pandas(pd_df)
ks_df_from_pd.head()

Transformar DataFrame
de Spark a Koalas
spark_df = spark.createDataFrame(pd_df)
ks_df_from_spark = spark_df.to_koalas()
ks_df_from_spark.head()

Transformar DataFrame
de Koalas a pandas
pd_df_from_koalas = ks_df.to_pandas()
pd_df_from_koalas.head()

Transformar DataFrame
de Koalas a Spark
spark_df_from_koalas = ks_df.to_spark()
spark_df_from_koalas.show(5)

Gráficos y visualizaciones de datos
En pandas podemos usar DataFrame.plot()
para implementar visualizaciones. Igualmente está disponible en Koalas.
El siguiente fragmento de código crea un diagrama de barras a partir de un DataFrame
de Koalas:
speed = [0.1, 17.5, 40, 48, 52, 69, 88]
lifespan = [2, 8, 70, 1.5, 25, 12, 28]
index = ['snail', 'pig', 'elephant', 'rabbit', 'giraffe', 'coyote', 'horse']
kdf = ks.DataFrame({'speed': speed, 'lifespan': lifespan}, index=index)
kdf.plot.bar()

Otro ejemplo sería un gráfico de línea:
kdf = ks.DataFrame({'pig': [20, 18, 489, 675, 1776], 'horse': [4, 25, 281, 600, 1900]}, index=[1990, 1997, 2003, 2009, 2014])
kdf.plot.line()

Problemas de Koalas
La release 1.0 de Koalas (lanzada en junio de 2020) cubría alrededor del 70-80% de la API de pandas, por lo que no todas las funcionalidades de pandas estaban disponibles en Koalas. En el momento de escribir este post, Koalas ya soporta Spark 3.1 e inferiores, y será incluido en la próxima release de Spark 3.2, por lo que se espera que las funcionalidades restantes se desarrollen rápidamente gracias a la gran comunidad de desarrolladores que existe detrás de este proyecto. El registro y documentación de las funciones que aún no están soportadas en Koalas se puede ver aquí.
Por otro lado, es importante resaltar que Koalas no ha sido creado como reemplazo de PySpark. Para dominar Koalas y aprovechar al máximo todas sus ventajas, es imprescindible conocer los básicos sobre cómo funciona Spark. Además, PySpark sigue siendo necesario para varios casos. Por ejemplo, si queremos implementar algoritmos de Machine Learning tendremos que usar el módulo de Spark MLlib, mientras que para crear aplicaciones streaming utilizaremos Spark Streaming.
Conclusión
Koalas resulta muy beneficioso para aquellos Data Scientists que ya están usando pandas y quieren comenzar a trabajar con datasets más grandes de manera sencilla, o para convertir código ya implementado en una solución escalable.
Es muy probable que, en un futuro cercano, Koalas se convierta en el estándar para proyectos Data Science que necesiten utilizar un entorno distribuido, mientras que PySpark continuará siendo la tecnología de referencia para las tareas que realizan los Data Engineers.