Migrate Pandas to Polars in 20 Minutes Using AI

pip install polars anthropic --break-system-packages

export ANTHROPIC_API_KEY='your-key-here'

import anthropic
import os
import sys

def convert_pandas_to_polars(pandas_code: str) -> str:
    """Convert Pandas code to Polars using Claude API."""
    
    client = anthropic.Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY"))
    
    # System prompt defines conversion rules
    system_prompt = """You are a Python expert converting Pandas code to Polars.

Rules:
1. Use polars.scan_csv() for lazy loading (not read_csv)
2. Replace df.groupby().agg() with df.group_by().agg()
3. Use pl.col() for column references in expressions
4. Chain operations with .pipe() when needed
5. Replace inplace=True operations with reassignment
6. Use .collect() at the end to execute lazy operations
7. Keep imports minimal: import polars as pl

Preserve:
- Variable names
- Comments
- Logic flow
- Error handling

Return ONLY the converted Python code, no explanations."""

    message = client.messages.create(
        model="claude-sonnet-4-20250514",
        max_tokens=4000,
        system=system_prompt,
        messages=[{
            "role": "user",
            "content": f"Convert this Pandas code to Polars:\n\n```python\n{pandas_code}\n```"
        }]
    )
    
    # Extract code from response
    response_text = message.content[0].text
    
    # Remove markdown code fences if present
    if "```python" in response_text:
        code = response_text.split("```python")[1].split("```")[0].strip()
    elif "```" in response_text:
        code = response_text.split("```")[1].split("```")[0].strip()
    else:
        code = response_text.strip()
    
    return code


def main():
    if len(sys.argv) != 2:
        print("Usage: python pandas_to_polars.py <input_file.py>")
        sys.exit(1)
    
    input_file = sys.argv[1]
    output_file = input_file.replace('.py', '_polars.py')
    
    # Read original code
    with open(input_file, 'r') as f:
        pandas_code = f.read()
    
    print(f"Converting {input_file}...")
    
    # Convert using AI
    polars_code = convert_pandas_to_polars(pandas_code)
    
    # Write converted code
    with open(output_file, 'w') as f:
        f.write(polars_code)
    
    print(f"✓ Converted code saved to {output_file}")
    print(f"✓ Review the output and test before using in production")


if __name__ == "__main__":
    main()

import pandas as pd
import numpy as np

# Load data
df = pd.read_csv('sales_data.csv')

# Clean data
df['revenue'] = df['quantity'] * df['price']
df = df[df['revenue'] > 0]

# Aggregate by category
summary = df.groupby('category').agg({
    'revenue': ['sum', 'mean'],
    'quantity': 'sum'
}).reset_index()

# Add percentage
summary['pct_of_total'] = summary[('revenue', 'sum')] / summary[('revenue', 'sum')].sum()

# Save results
summary.to_csv('summary.csv', index=False)
print(f"Processed {len(df)} rows")

python pandas_to_polars.py sample_pandas.py

import polars as pl

# Load data lazily
df = pl.scan_csv('sales_data.csv')

# Clean data
df = df.with_columns([
    (pl.col('quantity') * pl.col('price')).alias('revenue')
]).filter(pl.col('revenue') > 0)

# Aggregate by category
summary = df.group_by('category').agg([
    pl.col('revenue').sum().alias('revenue_sum'),
    pl.col('revenue').mean().alias('revenue_mean'),
    pl.col('quantity').sum().alias('quantity_sum')
]).collect()

# Add percentage
total_revenue = summary['revenue_sum'].sum()
summary = summary.with_columns([
    (pl.col('revenue_sum') / total_revenue).alias('pct_of_total')
])

# Save results
summary.write_csv('summary.csv')
print(f"Processed {len(df.collect())} rows")

# Pandas
df['date'] = pd.to_datetime(df['date'])

# Polars (check conversion)
df = df.with_columns(pl.col('date').str.strptime(pl.Date, '%Y-%m-%d'))

# Pandas
df.loc[df['age'] > 18, 'category'] = 'adult'

# Polars (AI converts to when-then)
df = df.with_columns(
    pl.when(pl.col('age') > 18)
      .then(pl.lit('adult'))
      .otherwise(pl.col('category'))
      .alias('category')
)

# Pandas
df.drop_duplicates(inplace=True)

# Polars (no inplace, always reassign)
df = df.unique()

import pandas as pd
import polars as pl
import numpy as np

# Generate test data
np.random.seed(42)
data = {
    'category': np.random.choice(['A', 'B', 'C'], 1000),
    'quantity': np.random.randint(1, 100, 1000),
    'price': np.random.uniform(10, 100, 1000)
}

# Test Pandas version
df_pd = pd.DataFrame(data)
df_pd['revenue'] = df_pd['quantity'] * df_pd['price']
result_pd = df_pd.groupby('category')['revenue'].sum().sort_index()

# Test Polars version
df_pl = pl.DataFrame(data)
df_pl = df_pl.with_columns([
    (pl.col('quantity') * pl.col('price')).alias('revenue')
])
result_pl = df_pl.group_by('category').agg(
    pl.col('revenue').sum()
).sort('category')

# Compare results
print("Pandas result:")
print(result_pd)
print("\nPolars result:")
print(result_pl)

# Check if values match (within floating point precision)
assert np.allclose(result_pd.values, result_pl['revenue'].to_numpy())
print("\n✓ Results match!")

python verify_conversion.py

import time
import polars as pl
import pandas as pd

# Create test dataset
df_pd = pd.DataFrame({
    'id': range(1_000_000),
    'category': np.random.choice(['A', 'B', 'C', 'D'], 1_000_000),
    'value': np.random.randn(1_000_000)
})

# Pandas timing
start = time.time()
result_pd = df_pd.groupby('category').agg({'value': ['mean', 'std', 'count']})
pandas_time = time.time() - start

# Polars timing
df_pl = pl.DataFrame(df_pd)
start = time.time()
result_pl = df_pl.group_by('category').agg([
    pl.col('value').mean(),
    pl.col('value').std(),
    pl.col('value').count()
])
polars_time = time.time() - start

print(f"Pandas: {pandas_time:.3f}s")
print(f"Polars: {polars_time:.3f}s")
print(f"Speedup: {pandas_time/polars_time:.1f}x")

import os
from pathlib import Path

def convert_project(src_dir: str, out_dir: str):
    """Convert all Python files in a directory."""
    
    Path(out_dir).mkdir(exist_ok=True)
    
    for root, dirs, files in os.walk(src_dir):
        for file in files:
            if file.endswith('.py'):
                input_path = os.path.join(root, file)
                rel_path = os.path.relpath(input_path, src_dir)
                output_path = os.path.join(out_dir, rel_path)
                
                # Ensure output directory exists
                os.makedirs(os.path.dirname(output_path), exist_ok=True)
                
                with open(input_path, 'r') as f:
                    pandas_code = f.read()
                
                # Skip files without pandas imports
                if 'import pandas' not in pandas_code:
                    continue
                
                print(f"Converting {rel_path}...")
                polars_code = convert_pandas_to_polars(pandas_code)
                
                with open(output_path, 'w') as f:
                    f.write(polars_code)

# Usage
convert_project('src/', 'src_polars/')

# Check if key is set
echo $ANTHROPIC_API_KEY

# If empty, re-export
export ANTHROPIC_API_KEY='sk-ant-...'

# Explicitly cast when needed
df = df.with_columns([
    pl.col('date').cast(pl.Date),
    pl.col('amount').cast(pl.Float64)
])