Leituras de dados JDBC simultâneas do Spark

Você já passou pelo processo de implementação do Spark em seu projeto, determinando o número ideal de partições aleatórias, alocação de memória para instâncias de driver e executor, número de núcleos executores e todas essas coisas divertidas apenas para ler dados de uma fonte 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()

Qual é o problema com o código acima? Bem… é muito lento :) Se você executar uma consulta SQL no Spark como no exemplo acima, você usará apenas um thread (você pode ver na interface do usuário do Spark que apenas uma tarefa é executada durante a execução do código) e o dataframe criado será tem uma partição.

A maioria de nós usa Spark em pipelines ETL/ELT para processamento paralelo, portanto, ler dados de uma fonte JDBC usando apenas um thread provavelmente não é o que queremos. Sem mencionar que teremos que reparticionar o dataframe após executar a consulta porque os dados agora estão armazenados em 1 partição. Então, como realmente lemos os dados simultaneamente no 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:
   Com base na documentação oficial do Spark, se quisermos ler dados simultaneamente e particioná-los com base em nossas necessidades, não podemos mais usar a opção de consulta e devemos passar nossa consulta SQL pela opção dbtable . A única diferença entre essas duas opções (além do nome) é que a consulta agora deve estar entre parênteses e ter um alias — basicamente passar a consulta original como uma subconsulta.
2. numPartitions:
   O número máximo de partições que podem ser usadas para paralelismo na leitura e gravação de tabelas. Isso também determina o número máximo de conexões JDBC simultâneas. Se o número de partições a serem gravadas exceder esse limite, reduzimos para esse limite chamando coalesce(numPartitions) antes de gravar. ( source ) Tenho o meu definido como 100 no cluster EMR, mas uma boa prática seria definir o número da partição entre 1 e 4 vezes o número de núcleos. Não é um número definido em pedra, portanto, sempre teste com valores diferentes para ver o que funciona para você.
3. partitionColumn:
   Este é provavelmente o parâmetro mais importante e aquele que fará a maior diferença no tempo de execução de suas consultas. Algumas coisas a ter em mente:
   - a coluna em questão deve ser numérica, data ou timestamp
   - para o melhor desempenho possível, os valores da coluna devem ser tão uniformemente distribuídos quanto possível e ter alta cardinalidade. Queremos evitar colunas de dados distorcidas — vou expandir isso abaixo.
   - se você tiver várias colunas com as propriedades acima, selecione a coluna indexada.
4. lowerbound, upperbound :
   Observe que lowerBound e upperBound são usados ​​apenas para decidir o passo da partição, não para filtrar as linhas na tabela. Portanto, todas as linhas da tabela serão particionadas e retornadas. ( fonte ) Abaixo você verá como o Spark gera cada partição. É por isso que você não deseja uma coluna de baixa cardinalidade como sua coluna de partição, mas darei um exemplo abaixo caso não faça sentido agora.

Imagine que a coluna que você deseja usar para particionar tenha apenas os valores 0 e 1 — um exemplo extremo, de fato, mas um cenário comum em Engenharia de Dados onde valores booleanos True e False são convertidos em valores numéricos.

Como você tem apenas 0s e 1s, use 0 como limite inferior e 1 como limite superior. Com base em como as partições são geradas — veja a figura acima — todas as linhas com o valor 0 serão enviadas para a primeira partição (a cláusula WHERE da primeira partição é a única cláusula em que 0 não é filtrado) e todas as linhas com o valor 1 será enviado para a última partição (a cláusula WHERE da última partição é a única cláusula em que 1 não é filtrado) .

Portanto, independentemente de quantas partições você criar, apenas duas partições serão preenchidas e apenas dois núcleos executores farão todo o trabalho, os outros ficarão em estado ocioso se não houver outro trabalho agendado para eles. Por causa disso, o desempenho da sua consulta será afetado.

Distorção de dados

Vamos imaginar que temos um conjunto de dados com 20 milhões de linhas e 30 partições, sendo os limites inferior e superior 2020–01–01 e 2022–12–31. Para o propósito deste artigo, usarei uma Spark Session com 30 núcleos executores (o que significa que podemos executar até 30 tarefas ao mesmo tempo). Lembre-se de que o Spark atribui uma tarefa por partição, portanto, cada partição será processada por um núcleo executor.

Cada ano de dados será dividido em 10 partições e considerando que 1 tarefa será atribuída a 1 partição, cada tarefa atribuída aos anos 2020 e 2021 será responsável por processar 100k linhas, mas cada tarefa atribuída ao ano 2022 terá que processar 1,8 milhões de linhas. Isso representa um aumento de 1700% nos dados que devem ser processados ​​por uma tarefa.

O resultado é que as primeiras 20 tarefas terminarão o processamento de suas partições atribuídas rapidamente e, depois disso, os 20 núcleos executores atribuídos às tarefas concluídas ficarão ociosos (se não houver outro trabalho em que os núcleos executores ociosos possam ser usados) até o último 10 tarefas terminam de processar 1700% mais registros. Isso é um desperdício de recursos e aumento nos custos operacionais, especialmente se você usar o Spark em um cluster EMR, Glue Job ou Databricks.

Como consertamos isso? A solução mais simples: use uma coluna diferente para particionar. Mas se for crítico usar esta coluna, podemos usar duas consultas e dois dataframes em vez de um. Você ainda terá problemas com downstream de dados distorcidos (um problema recorrente em sistemas computacionais paralelos), mas resolvemos um problema de cada vez.

A primeira consulta carregará os dois primeiros anos (os 30 núcleos executores serão divididos entre 2020 e 2021), depois a segunda consulta carregará o último ano (portanto, todos os 30 núcleos executores serão usados ​​agora para 2022). Por causa disso, uma determinada tarefa associada a 2022 agora será responsável pelo processamento de 600 mil linhas em vez de 1,8 milhão de linhas. Isso representa uma redução de 66% nas linhas que devem ser processadas por um único núcleo executor em relação à implementação original e, além disso, agora usamos todos os núcleos executores de nossa Spark Session em vez de deixá-los ociosos.

E se eu não tiver nenhuma coluna numérica, de data ou carimbo de data/hora ou se a coluna tiver dados distorcidos?

Não é um problema. Como engenheiros de dados, usar um CTE e uma função de janela para gerar uma coluna numérica row_number()deve ser muito fácil. Alguma sobrecarga será introduzida à medida que a nova coluna for gerada, mas mesmo com essa sobrecarga, a consulta final será executada muito mais rapidamente usando a coluna recém-criada em vez de tentar carregar todos os dados em uma partição usando um núcleo executor.

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'

Executamos duas consultas em vez de uma:
- a primeira consulta será executada para encontrar os valores mínimo e máximo da coluna que usaremos como partitionColumn
- na segunda consulta passaremos os resultados da primeira consulta como argumentos de limite inferior e superior .

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()

• Se você passar pelo processo de implementação do Spark em seu aplicativo, tente usá-lo ao máximo e ler os dados simultaneamente. A execução de uma consulta SQL para spark.read(), como mostrado na documentação JDBC do Spark, enviará todos os dados em uma partição e usará apenas um núcleo executor, independentemente de quantos núcleos você configurou em sua sessão do Spark.
• Se você usa o Spark hoje para ler dados de fontes de dados JDBC usando a abordagem de encadeamento único, tente refatorar suas consultas e melhorar seu desempenho, você pode até conseguir uma promoção enquanto estiver fazendo isso :)
• No final do dia, tempo é dinheiro. Se você está preocupado com o esforço de reescrever suas consultas, posso falar por experiência própria — algumas das consultas Spark que refatorei usando as técnicas descritas acima agora são executadas 90% mais rápido do que antes. Você pode imaginar a recepção positiva que essa melhoria teve na minha equipe.

Gabriel