Python导出海量数据至文件时内存溢出的解决方法

Absolutely! Python has solid solutions for working with datasets too big to fit into memory—perfect for your use case of pulling data from a database and writing it to disk using a fixed-size buffer (like 1MB). Let’s walk through the core approaches, step by step.

Core Idea: Stream Data + Buffered Writes #

The key here is to process data in chunks (instead of loading everything into memory at once) and use a buffer to hold chunks until it hits your target size, then flush it to disk. You can either implement this buffer manually or use Python’s built-in tools to handle the heavy lifting.

1. Manual Fixed-Size Buffer Implementation #

If you want full control over the buffer, you can manage it yourself with a bytearray. Here’s a practical example using PostgreSQL’s psycopg2 (the logic translates easily to other DB clients like pymysql):

import psycopg2
from typing import Iterator

# Define your buffer size (1MB = 1024*1024 bytes)
BUFFER_SIZE = 1024 * 1024

def fetch_data_stream(conn) -> Iterator[bytes]:
    """Stream data from the database without loading everything into memory"""
    # Use a server-side cursor to avoid pulling all results to the client at once
    with conn.cursor(name="large_dataset_cursor") as cursor:
        cursor.execute("SELECT large_column FROM your_big_table")
        for row in cursor:
            # Adjust this conversion based on your data type (e.g., JSON to bytes, text to UTF-8)
            yield row[0].tobytes()  # Example for binary data; use .encode('utf-8') for text

def buffered_file_writer(data_stream: Iterator[bytes], output_path: str):
    """Write data to disk using a fixed-size buffer"""
    buffer = bytearray()
    with open(output_path, "wb") as output_file:
        for chunk in data_stream:
            buffer.extend(chunk)
            # Flush buffer to disk whenever it reaches the target size
            while len(buffer) >= BUFFER_SIZE:
                output_file.write(buffer[:BUFFER_SIZE])
                # Keep any leftover data in the buffer
                buffer = buffer[BUFFER_SIZE:]
        # Don't forget to write the final remaining buffer contents
        if buffer:
            output_file.write(buffer)

# Main workflow
if __name__ == "__main__":
    db_conn = psycopg2.connect("dbname=your_db user=your_username password=your_pw")
    try:
        data_stream = fetch_data_stream(db_conn)
        buffered_file_writer(data_stream, "output_large_file.bin")
    finally:
        db_conn.close()

Key Notes for This Approach:

• Server-Side Cursors: Critical for databases—they tell the database to send results in batches instead of all at once, preventing your Python process from being swamped with memory.
• Buffer Management: The bytearray acts as a temporary holding spot; once it hits 1MB, we write the full buffer to disk and keep any leftover data for the next chunk.

2. Use Python's Built-in io.BufferedWriter #

For a more concise solution, Python’s io module includes BufferedWriter, which lets you specify a custom buffer size. It handles flushing to disk automatically when the buffer is full:

import psycopg2
import io

BUFFER_SIZE = 1024 * 1024

def fetch_data_stream(conn):
    with conn.cursor(name="large_dataset_cursor") as cursor:
        cursor.execute("SELECT large_column FROM your_big_table")
        for row in cursor:
            yield row[0].tobytes()

if __name__ == "__main__":
    db_conn = psycopg2.connect("dbname=your_db user=your_username")
    try:
        # Wrap the raw file in a BufferedWriter with our custom buffer size
        with open("output_large_file.bin", "wb") as raw_file:
            with io.BufferedWriter(raw_file, buffer_size=BUFFER_SIZE) as buffered_file:
                for chunk in fetch_data_stream(db_conn):
                    buffered_file.write(chunk)
                # Optional: Explicit flush (the `with` block does this automatically on exit)
                buffered_file.flush()
    finally:
        db_conn.close()

This is great if you don’t need fine-grained control over the buffer—Python handles all the flushing logic for you.

Critical Best Practices #

• Avoid fetchall(): Never use cursor.fetchall() for large datasets—it loads every row into memory instantly. Always iterate over the cursor or use fetchmany(size=N) to pull batches.
• Data Type Handling: If you’re working with text (like CSV or JSON), convert strings to bytes before writing to the buffer (e.g., json.dumps(row).encode('utf-8')). For CSV specifically, combine csv.writer with a buffered text wrapper:

  import csv
  # ... (database setup same as above)
  with open("output.csv", "wb") as raw_file:
      with io.BufferedWriter(raw_file, buffer_size=BUFFER_SIZE) as buffered_file:
          text_wrapper = io.TextIOWrapper(buffered_file, encoding="utf-8", line_buffering=False)
          writer = csv.writer(text_wrapper)
          writer.writerow(["col1", "col2"])  # Write header
          for row in fetch_data_stream(db_conn):
              writer.writerow(row)
  

• Buffer Size Tuning: 1MB is a safe default, but adjust based on your system. Too large and you waste memory; too small and you’ll have more frequent disk I/O (slowing things down).
• Error Handling: Add try/except blocks around writes to catch issues mid-stream, and consider database transactions if you need to roll back partial writes.

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