Lecturas de datos JDBC concurrentes de Spark

¿Alguna vez ha pasado por el proceso de implementación de Spark en su proyecto, determinando la cantidad óptima de particiones aleatorias, asignación de memoria para instancias de controlador y ejecutor, cantidad de núcleos ejecutores y todas esas cosas divertidas solo para leer datos de una fuente JDBC como esta? ?

jdbcDF = spark.read \
    .format("jdbc") \
    .option("driver", "org.postgresql.Driver") \
    .option("url", "jdbc:postgresql:dbserver") \
    .option("user", os.environ['user']) \
    .option("password", os.environ['pass']) \
    .option("query", query) \
    .load()

¿Cuál es el problema con el código anterior? Bueno... es muy lento :) Si ejecuta una consulta SQL en Spark como en el ejemplo anterior, usará solo un subproceso (puede ver en la interfaz de usuario de Spark que solo se ejecuta una tarea mientras se ejecuta el código) y el marco de datos creado se tener una partición.

La mayoría de nosotros usamos Spark en canalizaciones ETL/ELT para el procesamiento paralelo, por lo que leer datos de una fuente JDBC usando solo un subproceso probablemente no sea lo que queremos. Sin mencionar que tendremos que volver a dividir el marco de datos después de ejecutar la consulta porque los datos ahora están almacenados en 1 partición. Entonces, ¿cómo leemos datos al mismo tiempo en Spark?

# pass the original query to dbtable as a subquery 
final_query = f'({query}) as q'

# depends on your usecase, how many executor cores you have in the cluster
# what do you plan to do with the dataframe downstream, data size etc
partitionsCount = 100

# I'm hardcoding the date values here, but I'll show below how to identify
# these values at run time if you don't know them upfront
min_date = '2020-01-01'
max_date = '2022-12-31'

jdbcDF = spark.read \
    .format("jdbc") \
    .option("driver", "org.postgresql.Driver") \
    .option("url", "jdbc:postgresql:dbserver") \
    .option("user", os.environ['user']) \
    .option("password", os.environ['pass']) \
    .option("dbtable", final_query) \
    .option("numPartitions", partitionsCount) \
    .option("partitionColumn", "date") \
    .option("lowerBound", f"{min_date}") \
    .option("upperBound", f"{max_date}") \
    .load()

1. dbtable:
   según la documentación oficial de Spark, si queremos leer datos simultáneamente y particionarlos según nuestras necesidades, ya no podemos usar la opción de consulta y debemos pasar nuestra consulta SQL a través de la opción dbtable . La única diferencia entre estas dos opciones (aparte del nombre) es que la consulta ahora debe estar entre paréntesis y tener un alias; básicamente, pasar la consulta original como una subconsulta.
2. numPartitions:
   el número máximo de particiones que se pueden usar para el paralelismo en la lectura y escritura de tablas. Esto también determina el número máximo de conexiones JDBC simultáneas. Si el número de particiones para escribir excede este límite, lo reducimos a este límite llamando a coalesce(numPartitions) antes de escribir. ( fuente ) Tengo el mío configurado en 100 en el clúster EMR, pero una buena práctica sería configurar el número de partición entre 1 y 4 veces el número de núcleos. Sin embargo, no es un número escrito en piedra, así que siempre pruebe con diferentes valores para ver qué funciona para usted.
3. PartitionColumn:
   este es probablemente el parámetro más importante y el que hará la mayor diferencia en el tiempo de ejecución de sus consultas. Algunas cosas a tener en cuenta:
   - la columna en cuestión debe ser numérica, de fecha o de marca de tiempo
   - para obtener el mejor rendimiento posible, los valores de la columna deben estar distribuidos lo más uniformemente posible y tener una alta cardinalidad. Queremos evitar las columnas sesgadas de datos. Ampliaré esto a continuación.
   - si tiene varias columnas con las propiedades anteriores, seleccione la columna indexada.
4. lowerbound, upperbound :
   observe que lowerBound y upperBound solo se usan para decidir el paso de la partición, no para filtrar las filas en la tabla. Entonces, todas las filas de la tabla se dividirán y devolverán. ( fuente ) A continuación, verá cómo Spark genera cada partición. Esta es la razón por la que no desea una columna de cardinalidad baja como columna de partición, pero le daré un ejemplo a continuación en caso de que no tenga sentido ahora.

Imagine que la columna que desea usar para la partición tiene solo valores 0 y 1; un ejemplo extremo, pero un escenario común en ingeniería de datos donde los valores booleanos True y False se convierten en valores numéricos.

Debido a que solo tiene 0 y 1, usa 0 como límite inferior y 1 como límite superior. Según cómo se generan las particiones, vea la imagen de arriba, todas las filas con el valor 0 se enviarán a la primera partición (la cláusula WHERE de la primera partición es la única cláusula donde 0 no se filtra) y todas las filas con el valor 1 se enviará a la última partición (la cláusula WHERE de la última partición es la única cláusula en la que 1 no se filtra) .

Por lo tanto, independientemente de la cantidad de particiones que cree, solo se llenarán dos particiones y solo dos núcleos ejecutores harán todo el trabajo, los demás estarán en un estado inactivo si no hay otro trabajo programado para ellos. Debido a eso, el rendimiento de su consulta se verá afectado.

asimetría de datos

Imaginemos que tenemos un conjunto de datos con 20 millones de filas y 30 particiones, siendo los límites inferior y superior 2020–01–01 y 2022–12–31. A los efectos de este artículo, usaré una sesión de Spark con 30 núcleos ejecutores (lo que significa que podemos ejecutar hasta 30 tareas al mismo tiempo). Tenga en cuenta que Spark asigna una tarea por partición, por lo que cada partición será procesada por un núcleo ejecutor.

Cada año de datos se dividirá en 10 particiones y considerando que 1 tarea se asignará a 1 partición, cada tarea asignada a los años 2020 y 2021 será responsable de procesar 100k filas, pero cada tarea asignada al año 2022 tendrá que procesar 1,8 millones de filas. Eso es un aumento del 1700 % en los datos que debe procesar una tarea.

El resultado es que las primeras 20 tareas terminarán de procesar su partición asignada en poco tiempo y luego los 20 núcleos ejecutores asignados a las tareas completadas permanecerán inactivos (si no hay otro trabajo donde se puedan usar los núcleos ejecutores inactivos) hasta el último 10 tareas terminan de procesar 1700% más registros. Esto es una pérdida de recursos y un aumento en los costos operativos, especialmente si usa Spark en un clúster EMR, Glue Job o Databricks.

¿Cómo arreglamos eso? La solución más simple: use una columna diferente para la partición. Pero si es fundamental usar esta columna, podemos usar dos consultas y dos marcos de datos en lugar de uno. Todavía tendrá problemas con los datos sesgados en sentido descendente (un problema recurrente en los sistemas computacionales paralelos), pero abordamos un problema a la vez.

La primera consulta cargará los primeros dos años (los 30 núcleos ejecutores se dividirán entre 2020 y 2021), luego la segunda consulta cargará el año final (por lo que los 30 núcleos ejecutores ahora se usarán para 2022). Por eso, una determinada tarea asociada con 2022 ahora será responsable de procesar 600 000 filas en lugar de 1,8 millones de filas. Esa es una disminución del 66 % en las filas que debe procesar un solo núcleo ejecutor en comparación con la implementación original y, además, ahora usamos todos los núcleos ejecutores de nuestra sesión de Spark en lugar de dejarlos inactivos.

¿Qué pasa si no tengo ninguna columna numérica, de fecha o de marca de tiempo, o si la columna tiene datos sesgados?

No es un problema. Como ingenieros de datos, usar un CTE y una función de ventana para generar una columna numérica row_number()debería ser pan comido. Se introducirán algunos gastos generales a medida que se genere la nueva columna, pero incluso con estos gastos generales, la consulta final se ejecutará mucho más rápido usando la columna recién creada en lugar de intentar cargar todos los datos en una partición usando un núcleo ejecutor.

with cte as
(
    SELECT f1, f2, f3
    FROM table
)
select *, row_number() over (order by f1) as rn from cte

original_query = "SELECT f1, f2, f3 FROM table"

final_query = f'(with cte as ({original_query}) select *, row_number() over (order by f1) as rn from cte) as q'

Ejecutamos dos consultas en lugar de una:
- la primera consulta se ejecutará para encontrar los valores mínimo y máximo de la columna que usaremos como columna de partición
- en la segunda consulta, pasaremos los resultados de la primera consulta como argumentos de límite superior e inferior .

query_min_max = f"SELECT min(date), max(date) from ({query}) q"
final_query = f"({query}) as q"

df_min_max = spark.read \
    .format("jdbc") \
    .option("driver", "org.postgresql.Driver") \
    .option("url", "jdbc:postgresql:dbserver") \
    .option("user", os.environ['user']) \
    .option("password", os.environ['pass']) \
    .option("query", query_min_max) \
    .load()

min = df_min_max.first()["min"]
max = df_min_max.first()["max"]

jdbcDF = spark.read \
    .format("jdbc") \
    .option("driver", "org.postgresql.Driver") \
    .option("url", "jdbc:postgresql:dbserver") \
    .option("user", os.environ['user']) \
    .option("password", os.environ['pass']) \
    .option("dbtable", final_query) \
    .option("numPartitions", partitionsCount) \
    .option("partitionColumn", "date") \
    .option("lowerBound", f"{min}") \
    .option("upperBound", f"{max}") \
    .load()

• Si realiza el proceso de implementación de Spark en su aplicación, intente utilizarlo al máximo y leer los datos al mismo tiempo. Ejecutar una consulta SQL a spark.read() como se muestra en la documentación JDBC de Spark enviará todos los datos a una partición y usará solo un núcleo ejecutor, independientemente de cuántos núcleos haya configurado en su sesión de Spark.
• Si usa Spark hoy para leer datos de fuentes de datos JDBC utilizando el enfoque de subproceso único, intente refactorizar sus consultas y mejorar su rendimiento, incluso podría obtener una promoción mientras lo hace :)
• Al final del día, el tiempo es dinero. Si le preocupa el esfuerzo de volver a escribir sus consultas, puedo hablar por experiencia: algunas de las consultas de Spark que he refactorizado usando las técnicas descritas anteriormente ahora se ejecutan > 90 % más rápido que antes. Te puedes imaginar la acogida tan positiva que tuvo esta mejora en mi equipo.

gabriel