如何使用Python在大型CSV文件中批量搜索关键词？

Hey there! Let's adapt your PostgreSQL multi-keyword logic to Python for large CSV files. Your original SQL looks for rows where the description contains any of hotel, travel, taxi, or food, then returns distinct rows by id. Here's how to replicate that efficiently:

Step-by-Step Solution #

We'll use Python's built-in csv module (ideal for large files since it reads rows one at a time, no need to load the entire CSV into memory). We'll also track seen IDs to mimic the DISTINCT ON (id) behavior from PostgreSQL.

Full Code Example #

import csv

def multi_keyword_search(csv_file_path, keywords, output_file=None):
    # Track seen IDs to ensure uniqueness (matches DISTINCT ON (id))
    seen_ids = set()
    # Columns we want to keep (matches your SQL SELECT clause)
    target_columns = ['id', 'year', 'cost', 'description']
    
    with open(csv_file_path, mode='r', newline='', encoding='utf-8') as infile:
        reader = csv.DictReader(infile)
        
        # Optional: Save results to a new CSV if an output path is provided
        if output_file:
            with open(output_file, mode='w', newline='', encoding='utf-8') as outfile:
                writer = csv.DictWriter(outfile, fieldnames=target_columns)
                writer.writeheader()
                
                for row in reader:
                    # Check if any keyword exists in the description
                    description = row.get('description', '').lower()  # Case-insensitive search
                    if any(keyword.lower() in description for keyword in keywords):
                        row_id = row['id']
                        if row_id not in seen_ids:
                            seen_ids.add(row_id)
                            # Write only the columns we care about
                            writer.writerow({col: row[col] for col in target_columns})
        else:
            # Print results directly if no output file is specified
            for row in reader:
                description = row.get('description', '').lower()
                if any(keyword.lower() in description for keyword in keywords):
                    row_id = row['id']
                    if row_id not in seen_ids:
                        seen_ids.add(row_id)
                        print(f"ID: {row['id']}, Year: {row['year']}, Cost: {row['cost']}, Description: {row['description']}")

# Usage example
if __name__ == "__main__":
    # Your target keywords (matches your SQL pattern)
    search_keywords = ['hotel', 'travel', 'taxi', 'food']
    # Path to your large CSV file
    input_csv = 'mydata.csv'
    # Optional: Path to save filtered results
    output_csv = 'filtered_results.csv'
    
    multi_keyword_search(input_csv, search_keywords, output_csv)

Key Details Explained #

• Memory Efficiency: csv.DictReader processes one row at a time, so even massive CSVs won't overwhelm your system.
• Distinct IDs: The seen_ids set ensures we only keep the first occurrence of each id, just like your PostgreSQL query's DISTINCT ON (id).
• Case Insensitivity: I added .lower() to both the description and keywords—remove this if you want strict case-sensitive matching (like your original SQL).
• Flexible Output: You can either print results to the console or save them to a new CSV for later use.

Alternative: Using Pandas (For Medium-Sized Files) #

If your CSV fits comfortably in memory, pandas can simplify the code even more:

import pandas as pd

# Load the CSV
df = pd.read_csv('mydata.csv')
# Define your keywords
keywords = ['hotel', 'travel', 'taxi', 'food']

# Filter rows where description contains any keyword
mask = df['description'].str.contains('|'.join(keywords), case=False)
filtered_df = df[mask]

# Keep only the first occurrence of each id
filtered_df = filtered_df.drop_duplicates(subset='id', keep='first')

# Select the columns you need and save/print
filtered_df = filtered_df[['id', 'year', 'cost', 'description']]
filtered_df.to_csv('filtered_results.csv', index=False)

Note: Pandas loads the entire CSV into memory, so skip this for files larger than your available RAM.

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