如何基于Python列表从Redshift(PostgreSQL)选列拉取并追加CSV

Hey there, let's work through your problem step by step! You're hitting a few key roadblocks here: getting your Python column lists into valid SQL syntax, fixing the table name formatting, correctly appending to your CSV, and a critical hidden issue with random sampling that would break the alignment between your numeric and categorical data.

First, let's break down the core issues in your current code: #

1. Column list formatting: Using {num_cols} directly in your SQL string doesn't convert the Python list into comma-separated column names SQL can understand.
2. Table name mistake: str(seg) turns your list into something like ['seg1','seg2'], which is invalid SQL syntax for a table name.
3. CSV append error: The 'rw' file mode doesn't exist in Python, and you're not using pandas' built-in tools to append correctly.
4. Random sampling misalignment: Running ORDER BY random() twice will pull two completely different sets of rows—your numeric and categorical columns won't match up at all.

Here's the fixed, robust solution: #

1. Setup & Basic Fixes

First, let's clean up your variables (I'll assume seg is a single table name; if you need to loop through multiple tables, we can adjust that later):

import pandas as pd
# If you're using psycopg2 for Redshift, we'll use this for safer SQL formatting later
from psycopg2 import sql

# Your column lists
cat_cols = ['stat','zip','turned_off','turned_on']
num_cols = ['acu_m1','acu_cnt_m1','acu_cnt_m2','acu_wifi_m2']
# Table name (use a string, not a list, for valid SQL)
target_table = 'seg1'
# Assume your database connection `cnxn` is already initialized correctly

2. Best Option: Pull All Data At Once (Avoid Sampling Issues)

The simplest and most reliable way to avoid misaligned data is to pull all columns in one query, then split them into your CSV:

# Combine columns into a single comma-separated string
all_columns = ','.join(num_cols + cat_cols)
# Build the SQL query
sql_query = f"""SELECT {all_columns} FROM public."{target_table}" ORDER BY random() LIMIT 50000;"""

# Pull the full dataset
full_df = pd.read_sql(sql_query, cnxn)

# First write numeric columns to CSV (no index)
full_df[num_cols].to_csv("df_sample.csv", index=False)
# Append categorical columns—use mode='a' and header=False to avoid duplicate headers
full_df[cat_cols].to_csv("df_sample.csv", mode='a', header=False, index=False)

3. If You Must Pull Data In Two Separate Queries

If you have a reason to split the queries (e.g., very large dataset), you need to lock in the same set of rows for both pulls. Here's how to do that using a unique ID column (assuming your table has one, like id):

# Step 1: Pull numeric columns + unique ID to lock the sample
num_cols_with_id = num_cols + ['id']
num_cols_str = ','.join(num_cols_with_id)
sql_num = f"""SELECT {num_cols_str} FROM public."{target_table}" ORDER BY random() LIMIT 50000;"""
df_num = pd.read_sql(sql_num, cnxn)

# Write numeric columns to CSV
df_num[num_cols].to_csv("df_sample.csv", index=False)

# Step 2: Get the IDs of the sampled rows to pull matching categorical data
sampled_ids = tuple(df_num['id'].tolist())
cat_cols_str = ','.join(cat_cols)
sql_cat = f"""SELECT {cat_cols_str} FROM public."{target_table}" WHERE id IN {sampled_ids} ORDER BY id;"""
df_cat = pd.read_sql(sql_cat, cnxn)

# Append categorical columns (sort by ID to match the numeric data order)
df_cat.to_csv("df_sample.csv", mode='a', header=False, index=False)

4. Safer SQL Formatting (Prevent SQL Injection)

If your column or table names might ever come from user input (or even just to follow best practices), use psycopg2.sql to format your query safely:

# Build numeric column query safely
sql_num_safe = sql.SQL("SELECT {} FROM public.{} ORDER BY random() LIMIT 50000").format(
    sql.SQL(',').join(map(sql.Identifier, num_cols)),
    sql.Identifier(target_table)
)
# Execute the query
df_num = pd.read_sql(sql_num_safe.as_string(cnxn), cnxn)

Key Things to Remember: #

• Table name quotes: If your table name uses uppercase letters or special characters, wrap it in double quotes (like public."Seg1")—Redshift converts unquoted names to lowercase automatically.
• No duplicate random sampling: Never run ORDER BY random() twice for related data sets—you'll end up with mismatched rows.
• CSV append rules: Always use header=False when appending to a CSV that already has headers, otherwise you'll duplicate column names.

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