优化从临时表到空间表的INSERT操作及索引构建

背景 #

通过Airflow编排ETL流程，将约8000万条数据从Snowflake经S3迁移至PostgreSQL，数据先加载至临时表，再转换插入生产表。

问题 #

从临时表插入生产表时，通过ST_MakePoint(longitude, latitude)生成地理点的操作，加上后续的GiST索引构建，占用了ETL流程的大部分时长。

示例SQL代码 #

-- 从临时表插入生产表
INSERT INTO public.target_table (geo_lookup, proximity_data, updated)
    SELECT
        ST_MakePoint(longitude, latitude) as geo_lookup,
        proximity_data,
        CURRENT_TIMESTAMP as updated
FROM public.source_table;

-- 为geo_lookup列创建GiST索引
CREATE INDEX IF NOT EXISTS target_table_geo_idx
ON public.target_table
USING GIST (geo_lookup);

环境详情 #

PostgreSQL: 16.1 on aarch64-unknown-linux-gnu
PostGIS: 3.4
Airflow: 2.9.2（AWS MWAA托管）
Airflow环境规格: mw1.small
调度器数量: 2
Worker数量: 最小1，最大10
Web服务器数量: 最小2，最大2

已尝试方案 #

• 改用自动提交的批量插入（仍串行循环处理批次），耗时大幅增加，效果更差；
• 提升AWS RDS Serverless实例硬件规格，对任务时长几乎无影响。

待尝试方向与诉求 #

尚未尝试通过Python异步/多进程实现并发插入，不确定如何结合Airflow实现（拆分为并行任务还是在单Operator内处理）；同时考虑从Snowflake端追踪变更数据，仅插入增量子集，寻求上述方案的实现方向及示例。

现有Python代码 #

def _insert_and_index_geo_data(self):
    insert_sql = """
        INSERT INTO public.target_table (geo_lookup, proximity_data, updated)
        SELECT
            ST_MakePoint(longitude, latitude) AS geo_lookup,
            proximity_data,
            CURRENT_TIMESTAMP AS updated
        FROM public.source_table;
    """
    self.pg_conn.execute_sql(query=insert_sql, id='insert_geo_data')

    index_sql = """
        CREATE INDEX IF NOT EXISTS target_table_geo_idx
        ON public.target_table
        USING GIST (geo_lookup);
    """
    self.pg_conn.execute_sql(query=index_sql, id='create_geo_index')

优化方案实现示例 #

1. Airflow并行任务拆分（基于Dynamic Task Mapping） #

利用Airflow的动态任务映射功能，将临时表数据分片，分配给多个Worker并行处理插入，最后统一执行索引构建。

步骤1：分片临时表

先给临时表添加分片键，示例SQL：

-- 给临时表添加分片标识
ALTER TABLE public.source_table ADD COLUMN shard_id INT;
UPDATE public.source_table SET shard_id = MOD(id, 8); -- 分成8个分片，匹配Worker最大数量

步骤2：Airflow DAG实现

from airflow import DAG
from airflow.providers.postgres.operators.postgres import PostgresOperator
from airflow.operators.python import PythonOperator
from datetime import datetime

with DAG(
    'geo_data_etl_parallel',
    start_date=datetime(2024, 1, 1),
    schedule_interval='@daily',
    catchup=False
) as dag:

    # 预分片任务（临时表每次加载后执行）
    add_shard_column = PostgresOperator(
        task_id='add_shard_column',
        postgres_conn_id='postgres_default',
        sql="""
            ALTER TABLE IF NOT EXISTS public.source_table ADD COLUMN shard_id INT;
            UPDATE public.source_table SET shard_id = MOD(id, 8);
        """
    )

    # 并行插入任务：动态生成8个分片处理任务
    insert_shard_task = PostgresOperator.partial(
        task_id='insert_shard',
        postgres_conn_id='postgres_default',
        sql="""
            INSERT INTO public.target_table (geo_lookup, proximity_data, updated)
            SELECT
                ST_MakePoint(longitude, latitude) AS geo_lookup,
                proximity_data,
                CURRENT_TIMESTAMP AS updated
            FROM public.source_table
            WHERE shard_id = {{ params.shard_id }};
        """
    ).expand(params=[{"shard_id": i} for i in range(8)])

    # 统一构建索引任务
    create_index = PostgresOperator(
        task_id='create_geo_index',
        postgres_conn_id='postgres_default',
        sql="""
            SET max_parallel_maintenance_workers = 4;
            CREATE INDEX IF NOT EXISTS target_table_geo_idx
            ON public.target_table
            USING GIST (geo_lookup);
        """
    )

    # 任务依赖
    add_shard_column >> insert_shard_task >> create_index

2. Snowflake增量同步实现 #

通过Snowflake的变更数据捕获（CDC）或时间戳追踪，仅同步新增/修改的数据，减少数据处理量。

方案A：基于时间戳的增量同步

假设Snowflake源表有last_modified字段，Airflow任务每次同步上次执行后的增量数据：

from airflow import DAG
from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
from airflow.providers.postgres.operators.postgres import PostgresOperator
from airflow.operators.python import PythonOperator
from airflow.models import Variable
from datetime import datetime

def get_last_sync_time():
    return Variable.get("last_geo_sync_time", default_var=str(datetime(2024, 1, 1)))

def update_last_sync_time(**context):
    Variable.set("last_geo_sync_time", str(datetime.now()))

with DAG(
    'geo_data_incremental_sync',
    start_date=datetime(2024, 1, 1),
    schedule_interval='@hourly',
    catchup=False
) as dag:

    # 从Snowflake导出增量数据到S3
    export_incremental = PythonOperator(
        task_id='export_incremental_to_s3',
        python_callable=lambda: SnowflakeHook(snowflake_conn_id='snowflake_default').run(
            sql=f"""
                COPY INTO @s3_stage/geo_incremental/{datetime.now().strftime('%Y%m%d%H')}
                SELECT longitude, latitude, proximity_data
                FROM snowflake_source.geo_data
                WHERE last_modified > '{get_last_sync_time()}';
            """
        )
    )

    # 省略：S3数据加载到PostgreSQL临时表的步骤（可使用AWS S3ToRedshiftOperator或自定义Operator）

    # 插入增量数据到生产表
    insert_incremental = PostgresOperator(
        task_id='insert_incremental_data',
        postgres_conn_id='postgres_default',
        sql="""
            INSERT INTO public.target_table (geo_lookup, proximity_data, updated)
            SELECT
                ST_MakePoint(longitude, latitude) AS geo_lookup,
                proximity_data,
                CURRENT_TIMESTAMP AS updated
            FROM public.staging_incremental_table;
        """
    )

    # 更新同步时间变量
    update_sync_time = PythonOperator(
        task_id='update_sync_time',
        python_callable=update_last_sync_time,
        provide_context=True
    )

    # 任务依赖
    export_incremental >> insert_incremental >> update_sync_time

方案B：Snowflake STREAM+TASK实现CDC

在Snowflake中创建流捕获源表变更，通过任务自动同步到S3：

-- 创建流捕获源表变更
CREATE OR REPLACE STREAM geo_data_stream ON TABLE snowflake_source.geo_data;

-- 创建定时任务将变更数据导出到S3
CREATE OR REPLACE TASK sync_geo_data_to_s3
WAREHOUSE = my_warehouse
SCHEDULE = '1 HOUR'
AS
COPY INTO @s3_stage/geo_cdc/
SELECT longitude, latitude, proximity_data
FROM geo_data_stream
WHERE METADATA$ACTION = 'INSERT' OR (METADATA$ACTION = 'DELETE' AND METADATA$ISUPDATE = TRUE);

3. SQL层面优化建议 #

• 预验证数据类型：确保longitude和latitude为FLOAT或NUMERIC类型，避免ST_MakePoint执行时的隐式类型转换；
• 并行索引构建：PostgreSQL 16支持并行GiST索引构建，通过SET max_parallel_maintenance_workers = 4;（根据实例CPU核数调整）加速索引创建；
• 禁用约束检查：插入前临时禁用生产表的非必要约束（如外键），插入完成后再启用，减少插入时的校验开销。

内容的提问来源于stack exchange，提问作者Olek