.
Случился стандартный рабочий кейс: надо затащить несколько внешних таблиц в клик.
Источник: Postgres
Сущность: N вьюх (у одной из под капотом было около десятка JOIN 🥶, дадим ей имя вьюха
G)История умалчивает подробности и даже почему выбран такой способ загрузки: целевая таблица полностью перегружается каждый день.
.
Работа шла стандартным флоу (практикую частичное решение задач):
- написал даг с коннектом к источнику
- таски для создания инфры
- несколько таск групп с небольшим селектом из источника (проверить, что всё ок и источник отдает данные)
- заварил 🫖 и стал налюдать за даграном
.
Всё шло хорошо до момента тестирования вьюхи G. В логах вижу инфра создается, как только дело доходит до забора данных, то таска ждет и падает по
SIGTERM - явно что-то идет не так..
В процессе обсуждения с разработчиком со стороны источника выяснилось:
- про десятки JOIN
- там в таблице 10+млн строк
- .....
.
Но тут я подумал: ведь я же батчами забираю, причем здесь млн строк🤔.
.
Код для загрузки батчами тривиален:
<code class="language-python"><code class="language-python">import psycopg2
con = psycopg2.connect(.....)
...
for chunk in pd.read_sql(stmt, con=con, chunk_size=100500):
yield chunk
.
Крутил, вертел никак, надо идти курить доку (тимлид подкинул про именованные курсоры). И от тут истина и открылась, если делать так как выше, то:
-
stmt должно выполниться полностью на стороне сервера- полные данные передаются на клиент
- и батчи получаются на стороне клиента - то есть у нас клиентский курсор
И естественно дождаться не получается пока собирутся 10+млн строк.
.
Из доки psycopg2 узнал про реализацию серверного курсора (тот который будет формировать батчи на стороне сервера и отдавать их). Для организиации серверного курсора меняем код, использовать красивую pandas-обертку уже не получится:
<code class="language-python">import psycopg2
from psycopg2.extras import RealDictCursor
with psycopg2.connect(...) as conn:
with conn.cursor(name="It's server cursor", cursor_factory=RealDictCursor()) as cursor:
cur.execute(stmt)
chunk = cur.fetchmany(chunk_size)
while chunk:
yield pd.DataFrame(chunk)
chunk = cur.fetchmany(chunk_size)
.
Бинго, теперь батчи отдаются по мере готовности и не нужно ждать сбора всех данных - пошло чтение данных. Далее я дописал загрузку в Clickhouse.
.
ps: на источнике окаалось 36+млн строк 🤨
#case #work #psycopg2 #def
