大CSV文件拆分批量导入PostgreSQL的S3+Lambda方案求助

Hey there! Let's work through this problem together—you're trying to load a 1M-row CSV from S3 into PostgreSQL using Lambda, processing 10k rows at a time. I’ve tackled similar batch processing tasks before, so I’ll walk you through a solid approach, common pitfalls to watch for, and code examples to get you unblocked.

First, let’s align on the key constraints here: Lambda has a 15-minute timeout cap, and loading the entire 1M-row CSV into memory is inefficient (and risky if the file grows). So our plan will focus on:

• Streaming the CSV directly from S3 (no full download)
• Processing rows in 10k batches for efficient DB inserts
• Tracking progress to resume if Lambda times out or fails
• Splitting work across multiple Lambda invocations if needed

1. Stream the CSV from S3 (Skip Full File Download) #

Instead of downloading the entire CSV to Lambda’s temporary storage, we’ll use boto3 to get a streaming object. Wrapping it in a buffered reader lets us track byte offsets—critical for picking up where we left off if processing is interrupted.

import boto3
import io
import csv

s3 = boto3.client('s3')
BUCKET_NAME = 'your-bucket-name'
FILE_KEY = 'path/to/your/large-file.csv'

# Fetch the S3 object as a stream
s3_response = s3.get_object(Bucket=BUCKET_NAME, Key=FILE_KEY)
stream = io.BufferedReader(s3_response['Body'])

2. Track Progress with DynamoDB #

To avoid reprocessing rows if Lambda times out, we’ll store the last processed byte offset in a DynamoDB table. Create a table (e.g., csv_processing_progress) with file_key as the primary key.

dynamodb = boto3.resource('dynamodb')
progress_table = dynamodb.Table('csv_processing_progress')

def get_last_processed_offset(file_key):
    """Get the last byte offset we processed, or 0 if it's the first run"""
    try:
        response = progress_table.get_item(Key={'file_key': file_key})
        return response['Item']['last_offset']
    except KeyError:
        return 0

def update_processing_offset(file_key, new_offset):
    """Update the progress table with the latest byte offset"""
    progress_table.put_item(Item={'file_key': file_key, 'last_offset': new_offset})

3. Batch Processing & PostgreSQL Inserts #

For fast batch inserts, psycopg2’s copy_from method is way more efficient than executemany—it mimics PostgreSQL’s native COPY command, which is optimized for bulk loads.

First, set up your DB connection:

import psycopg2
from psycopg2 import sql

def get_db_connection():
    """Create a PostgreSQL connection (replace with your DB credentials)"""
    return psycopg2.connect(
        host='your-db-host',
        database='your-db-name',
        user='your-db-user',
        password='your-db-password',
        port='5432'
    )

Now, process the stream in 10k-row batches:

BATCH_SIZE = 10000

def load_batch_to_postgres(rows, cursor):
    """Load a batch of rows into PostgreSQL using copy_from"""
    # Use StringIO to create an in-memory "file" for copy_from
    buffer = io.StringIO()
    for row in rows:
        # Adjust delimiter if your CSV uses something other than tabs
        buffer.write('\t'.join(str(col) for col in row) + '\n')
    buffer.seek(0)
    
    # Replace with your target table and column names
    cursor.copy_from(
        buffer,
        'your_target_table',
        sep='\t',
        columns=('column_1', 'column_2', 'column_3')
    )

def process_csv():
    last_offset = get_last_processed_offset(FILE_KEY)
    stream.seek(last_offset)  # Jump to our last processed position
    
    csv_reader = csv.reader(stream, delimiter=',', quotechar='"')  # Match your CSV format
    next(csv_reader)  # Skip header row if your CSV has one
    
    batch = []
    current_offset = last_offset
    
    for row in csv_reader:
        batch.append(row)
        current_offset = stream.tell()  # Track where we are in the stream
        
        if len(batch) >= BATCH_SIZE:
            # Insert the batch into PostgreSQL
            conn = get_db_connection()
            cursor = conn.cursor()
            try:
                load_batch_to_postgres(batch, cursor)
                conn.commit()
                update_processing_offset(FILE_KEY, current_offset)
                print(f"Successfully loaded {BATCH_SIZE} rows. Updated offset to {current_offset}")
                batch = []
            except Exception as e:
                conn.rollback()
                print(f"Error loading batch: {str(e)}")
                raise e
            finally:
                cursor.close()
                conn.close()
    
    # Process any remaining rows after the final full batch
    if batch:
        conn = get_db_connection()
        cursor = conn.cursor()
        try:
            load_batch_to_postgres(batch, cursor)
            conn.commit()
            # Clean up progress record once we're done
            progress_table.delete_item(Key={'file_key': FILE_KEY})
            print(f"Loaded final {len(batch)} rows. Processing complete!")
        except Exception as e:
            conn.rollback()
            print(f"Error loading final batch: {str(e)}")
            raise e
        finally:
            cursor.close()
            conn.close()

4. Handle Lambda Timeouts #

If processing all 100 batches takes longer than 15 minutes (Lambda’s max timeout), modify the code to stop after, say, 50 batches, update the offset, and trigger another Lambda invocation to continue. You can use boto3.client('lambda').invoke() to call the same function again, passing the current offset as an event parameter.

• Memory Constraints: Even with streaming, make sure your Lambda has enough memory (at least 512MB) to handle batch processing and DB connections.
• DB Connection Limits: PostgreSQL has default connection limits—if you’re running concurrent Lambda invocations, use a connection pool like pgBouncer to avoid hitting limits.
• CSV Parsing Errors: Double-check your csv.reader parameters (delimiter, quotechar) to match your CSV’s format—mismatches will break row parsing.
• Idempotency: Add unique constraints to your PostgreSQL table and use ON CONFLICT DO NOTHING if needed, so retries don’t create duplicate rows.

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