Data Manipulation at Scale

import pandas as pd

# Load from various sources
df = pd.read_csv('sales_data.csv')

# Same data from other sources. Parquet reads much faster than CSV and
# preserves dtypes; read_sql pulls straight from a database connection.
# df = pd.read_parquet('sales_data.parquet')
# df = pd.read_sql('SELECT * FROM sales', connection)

# Basic info
print(df.shape)        # (1000000, 10) - 1M rows, 10 columns
print(df.dtypes)       # Column data types

# The rest print too much to reproduce here, but you will use them constantly:
#   df.info()                       memory usage and non-null counts
#   df.describe()                   statistical summary of numeric columns
#   df.head()                       first 5 rows
#   df.tail(10)                     last 10 rows
#   df.sample(5, random_state=42)   5 random rows; seed it to get the same
#                                   rows twice, otherwise it differs per run

# Continues from the block above: 'df' and 'pd' come from there

# Column selection
df['customer_id']              # Single column (Series)
df[['customer_id', 'price']]   # Multiple columns (DataFrame)

# Row selection
df.loc[0]                      # By label (first row)
df.iloc[0]                     # By position (first row)
df.loc[0:5]                    # Slice by label (inclusive)
df.iloc[0:5]                   # Slice by position (exclusive)

# Boolean filtering
high_value = df[df['price'] > 100]
recent = df[df['date'] > '2024-01-01']

# Multiple conditions
filtered = df[
    (df['price'] > 50) &
    (df['quantity'] > 1) &
    (df['customer_id'].isin([1001, 1002, 1003]))
]

# Query method (more readable for complex filters)
queried = df.query('price > 50 and quantity > 1')

print(f"rows total:             {len(df)}")
print(f"price > 100:            {len(high_value)}")
print(f"three-condition filter: {len(filtered)}")
print(f"query() == bracket form: {len(queried) == len(df[(df['price'] > 50) & (df['quantity'] > 1)])}")

# Continues from the block above: 'df' and 'pd' come from there

# Simple aggregation
total_sales = df['price'].sum()
avg_quantity = df['quantity'].mean()

# Group by single column
by_customer = df.groupby('customer_id').agg({
    'price': 'sum',
    'quantity': 'mean',
    'product_id': 'count'
})

# Group by multiple columns
by_customer_product = df.groupby(['customer_id', 'product_id']).agg({
    'price': ['sum', 'mean', 'std'],
    'quantity': ['sum', 'count']
})

# Custom aggregation functions. include_groups=False keeps the grouping
# column out of what the function sees, which is the 3.x default behaviour.
def calculate_revenue(group):
    return (group['price'] * group['quantity']).sum()

revenue_by_customer = df.groupby('customer_id').apply(
    calculate_revenue, include_groups=False
)

# Multiple aggregations with rename
summary = df.groupby('customer_id').agg(
    total_revenue=('price', 'sum'),
    avg_order_value=('price', 'mean'),
    order_count=('customer_id', 'count'),
    unique_products=('product_id', 'nunique')
)

print(summary.head(3).round(2))
print(f"\nrevenue for customer 1000: {revenue_by_customer.loc[1000]:,.2f}")

# Continues from the block above: 'df' and 'pd' come from there

# Create new columns
df['revenue'] = df['price'] * df['quantity']
df['discount_applied'] = df['price'] < 30

# Apply functions
df['price_category'] = df['price'].apply(
    lambda x: 'High' if x > 100 else 'Medium' if x > 50 else 'Low'
)

# Vectorized operations (much faster!)
df['total_cost'] = df['price'] * df['quantity'] * 1.1  # Add 10% tax

# Map values
category_map = {501: 'Electronics', 502: 'Clothing', 503: 'Food'}
df['category'] = df['product_id'].map(category_map)

# Merge datasets
customers = pd.read_csv('customers.csv')
merged = df.merge(
    customers,
    on='customer_id',
    how='left'  # left, right, inner, outer
)

# Pivot tables (wide format). 100 product ids would be unreadable, so
# pivot on the four regions instead.
pivot = df.pivot_table(
    values='revenue',
    index='region',
    columns='status',
    aggfunc='sum',
    fill_value=0
)

# Melt (long format)
long_format = pivot.reset_index().melt(
    id_vars=['region'],
    var_name='status',
    value_name='revenue'
)

print(f"columns after merge: {len(df.columns)} -> {len(merged.columns)}")
print(pivot.round(0))
print(f"\nmelted back to {long_format.shape[0]} rows x {long_format.shape[1]} columns")

import dask.dataframe as dd
import pandas as pd

# Load data (lazy - doesn't load immediately). read_csv also accepts a glob
# such as 'large_sales_*.csv' to treat many files as one frame. blocksize
# controls how the file is cut into partitions.
ddf = dd.read_csv('sales_data.csv', blocksize='16MB')

# Partitions are how Dask splits the data. Each one is a Pandas DataFrame,
# and Dask works on them in parallel across cores.
print(f"partitions: {ddf.npartitions}")

# All operations are lazy (no computation yet)
filtered = ddf[ddf['price'] > 50]
grouped = filtered.groupby('customer_id')['price'].sum()
result = grouped / 1000  # Still no computation!

# Nothing has been read yet: this is a plan, not a value
print(f"lazy: {hasattr(result, 'compute')}")

# Trigger computation
final_result = result.compute()  # Now it executes in parallel
print(f"computed into a pandas Series: {isinstance(final_result, pd.Series)}")
print(final_result.head(3).round(3))

from dask.distributed import Client
import dask.dataframe as dd

# A Client spawns worker processes, so the script needs the main guard
if __name__ == '__main__':
    # Start a local cluster. Client() alone uses all cores and opens a
    # dashboard at http://localhost:8787; pinned here so the run is
    # reproducible.
    client = Client(n_workers=4, threads_per_worker=1, dashboard_address=None)

    ddf = dd.read_parquet('sales_data.parquet')

    # Complex pipeline (all operations are lazy).
    # Note what is NOT here: 'product_id': 'nunique'. Dask's groupby().agg()
    # rejects it with ValueError: unknown aggregate nunique. It is one of the
    # places the Dask API stops matching Pandas, and only running it tells you.
    result = (
        ddf
        .query('price > 50 and quantity > 1')
        .groupby('customer_id')
        .agg({
            'price': ['sum', 'mean'],
            'quantity': 'sum'
        })
        .compute()  # Execute in parallel
    )
    print(f"aggregated rows: {result.shape[0]}")

    # Persist in memory for reuse
    ddf_filtered = ddf[ddf['price'] > 50].persist()

    # Now operations on ddf_filtered are fast (already in memory).
    # nunique is fine on a Series, just not inside an agg() dict.
    high_value_customers = ddf_filtered.groupby('customer_id')['price'].sum()
    product_diversity = ddf_filtered.groupby('customer_id')['product_id'].nunique()

    # Both use the same persisted data
    result1 = high_value_customers.compute()
    result2 = product_diversity.compute()
    print(f"products bought by customer 1000: {result2.loc[1000]}")

    # Custom functions with map_partitions
    def custom_analysis(partition):
        """Runs on each partition (Pandas DataFrame)"""
        # Do whatever you want with Pandas
        return partition[partition['price'] > partition['price'].mean()]

    above_avg = ddf.map_partitions(custom_analysis).compute()
    print(f"rows above their partition mean: {len(above_avg)}")
    print(f"workers: {len(client.scheduler_info()['workers'])}")

    client.close()

# PANDAS VERSION (works for small data)
import pandas as pd

df = pd.read_csv('sales_data.csv')
pandas_result = (
    df[df['price'] > 50]
    .groupby('customer_id')['price']
    .mean()
)

# DASK VERSION (works for large data) - minimal changes!
import dask.dataframe as dd

ddf = dd.read_csv('sales_data.csv', blocksize='16MB')  # Changed pd to dd
dask_result = (
    ddf[ddf['price'] > 50]
    .groupby('customer_id')['price']
    .mean()
    .compute()  # Added .compute()
)

# Same numbers, two engines
print(f"identical: {pandas_result.round(9).equals(dask_result.round(9))}")
print(pandas_result.head(2).round(3))

# What Dask doesn't support well:
# - Operations requiring full dataset sorting
# - Complex window functions
# - Some merge operations with shuffling
# - Operations that break partitioning

# For these, might need to use .compute() to go back to Pandas
as_pandas = ddf.compute()  # Convert to Pandas
ranked = as_pandas.sort_values('price', ascending=False)
print(f"most expensive order: {ranked.iloc[0]['price']}")

# Then convert back to Dask if needed
ddf_result = dd.from_pandas(ranked, npartitions=4)
print(f"back to dask in {ddf_result.npartitions} partitions")

import time

import vaex

# Open dataset (instant, no loading!). vaex.open also takes a glob such as
# 'data/*.csv', or an .hdf5 file, and memory-maps rather than reads.
start = time.perf_counter()
df = vaex.open('sales_data.parquet')
print(f"opened {df.shape[0]:,} rows in {time.perf_counter() - start:.4f}s")

# Only show what's needed (still not loading all data)
print(df[['order_id', 'customer_id', 'price', 'quantity']].head(2))

# Statistics are computed on-demand
print(f"mean:  {df.price.mean():.4f}")   # Fast! Uses efficient algorithms
print(f"std:   {df.price.std():.4f}")    # Doesn't load all data
print(f"customers: {df.customer_id.nunique()}")

# Convert to Pandas (loads into memory - be careful!)
small_subset = df[df['price'] > 499]
pandas_df = small_subset.to_pandas_df()  # Now it's in RAM
print(f"subset in RAM: {len(pandas_df)} rows")

import vaex

df = vaex.open('sales_data.parquet')
print(f"columns on disk: {len(df.get_column_names())}")

# Create virtual columns (no memory used!)
df['revenue'] = df.price * df.quantity
df['is_high_value'] = df.price > 100

# Virtual columns computed on-demand
print(f"mean revenue: {df.revenue.mean():.4f}")   # Computed on the fly
print(f"high value:   {df.is_high_value.sum():,}")

# Filter with virtual columns
high_revenue = df[df.revenue > 1000]
print(f"revenue > 1000: {len(high_revenue):,} rows")

# Expressions are also virtual. Only real columns are available here:
# there is no 'tax' column, so build the tax from the discount instead.
df['net_price'] = (df.price * df.quantity) * (1 - df.discount)

# Chain operations (all virtual, no memory)
df['normalized_price'] = (df.price - df.price.mean()) / df.price.std()

# Nothing above was stored: the frame still owns only its original columns
print(f"virtual columns: {len(df.virtual_columns)}")
print(f"still on disk:   {len(df.get_column_names(virtual=False))}")

# Export to materialize virtual columns
df.export('sales_with_revenue.parquet')  # Now they are real columns
print(f"after export:    {len(vaex.open('sales_with_revenue.parquet').get_column_names())}")

import vaex

df = vaex.open('sales_data.parquet')

# Simple aggregations (very fast!)
print(f"mean:     {df.price.mean():.4f}")
print(f"sum:      {df.price.sum():,.2f}")
print(f"std:      {df.price.std():.4f}")
print(f"nunique:  {df.customer_id.nunique()}")

# Multiple aggregations, named
grouped = df.groupby('customer_id', agg={
    'total_revenue': vaex.agg.sum('price'),
    'avg_price': vaex.agg.mean('price'),
    'order_count': vaex.agg.count(),
    'unique_products': vaex.agg.nunique('product_id')
})
print(grouped.sort('customer_id').head(2))

# Binning (histograms) - super fast
counts = df.count(binby='price', limits=[0, 500], shape=5)
print(f"price histogram, 5 bins over 0-500: {counts.astype(int).tolist()}")

# 2D histograms
counts_2d = df.count(
    binby=[df.price, df.quantity],
    limits=[[0, 500], [0, 10]],
    shape=[5, 3]
)
print(f"2D histogram shape: {counts_2d.shape}, total binned: {int(counts_2d.sum()):,}")

import polars as pl

# EAGER MODE (executes immediately, like Pandas)
df = pl.read_csv('sales_data.csv')
df_result = (
    df
    .filter(pl.col('price') > 50)
    .group_by('customer_id')
    .agg(pl.col('price').sum())
)

# LAZY MODE (builds query plan, executes when needed)
df_lazy = pl.scan_csv('sales_data.csv')  # scan_ for lazy
result_lazy = (
    df_lazy
    .filter(pl.col('price') > 50)
    .group_by('customer_id')
    .agg(pl.col('price').sum())
    # Nothing executed yet!
)

# Optimize and execute
result = result_lazy.collect()  # Now it runs (optimized)
print(f"eager and lazy agree: {result.sort('customer_id').equals(df_result.sort('customer_id'))}")

# Example: Lazy is much faster for complex pipelines.
# Polars optimizes the entire chain before executing.
df_lazy = (
    pl.scan_parquet('sales_data.parquet')
    .select(['customer_id', 'price'])   # Only reads these columns
    .filter(pl.col('price') > 50)
    .group_by('customer_id')
    .agg(pl.col('price').mean())
)

# See the query plan. Read it bottom-up: the filter and the column list
# have been pushed down into the Parquet scan, so the other 8 columns are
# never read off disk.
plan = df_lazy.explain()
print(f"projection pushdown in plan: {'PROJECT' in plan}")
print(f"predicate pushdown in plan:  {'SELECTION' in plan or 'FILTER' in plan}")

result = df_lazy.collect()  # Optimized execution
print(result.sort('customer_id').head(2))

import polars as pl

df = pl.read_csv('sales_data.csv')

# Expression API - chainable, optimizable
result = df.select([
    pl.col('customer_id'),
    pl.col('price').alias('original_price'),
    (pl.col('price') * pl.col('quantity')).alias('revenue'),
    pl.col('price').mean().over('customer_id').alias('avg_customer_price'),
    # pl.lit() marks these as literal strings. A bare 'High' would be read
    # as a COLUMN name and raise ColumnNotFoundError.
    pl.when(pl.col('price') > 100)
      .then(pl.lit('High'))
      .otherwise(pl.lit('Low'))
      .alias('price_category')
])

# Powerful aggregations
summary = df.group_by('customer_id').agg([
    pl.col('price').sum().alias('total_spent'),
    pl.col('price').mean().alias('avg_price'),
    pl.col('product_id').n_unique().alias('unique_products'),
    pl.col('quantity').sum().alias('total_quantity'),
    # Conditional aggregation
    pl.col('price').filter(pl.col('price') > 100).sum().alias('high_value_total')
])

# Window functions (over)
df_with_ranks = df.select([
    '*',  # All existing columns
    pl.col('price').rank().over('customer_id').alias('price_rank'),
    pl.col('price').cum_sum().over('customer_id').alias('cumulative_spent')
])

# String operations, on the text columns this dataset actually has
df = df.with_columns([
    pl.col('region').str.to_uppercase().alias('region_upper'),
    pl.col('payment_method').str.slice(0, 4).alias('payment_short'),
    pl.col('status').str.starts_with('ship').alias('is_shipped')
])

print(result.head(3))

import time

import pandas as pd
import polars as pl

def fastest(fn, runs=3):
    """Best of N: the minimum is the least noisy estimate"""
    times = []
    for _ in range(runs):
        start = time.perf_counter()
        fn()
        times.append(time.perf_counter() - start)
    return min(times)

# The same aggregation, two engines. Polars parallelises across cores by
# default; pandas runs it on one.
pdf = pd.read_csv('sales_data.csv')
pldf = pl.read_csv('sales_data.csv')

pandas_agg = fastest(lambda: pdf.groupby('customer_id').agg(
    price_sum=('price', 'sum'),
    quantity_mean=('quantity', 'mean'),
    products=('product_id', 'nunique'),
))
polars_agg = fastest(lambda: pldf.group_by('customer_id').agg([
    pl.col('price').sum(),
    pl.col('quantity').mean(),
    pl.col('product_id').n_unique(),
]))
print(f"groupby   pandas {pandas_agg * 1000:7.1f} ms   "
      f"polars {polars_agg * 1000:6.1f} ms   {pandas_agg / polars_agg:4.1f}x")

# Reading the same 61 MB CSV
pandas_read = fastest(lambda: pd.read_csv('sales_data.csv'))
polars_read = fastest(lambda: pl.read_csv('sales_data.csv'))
print(f"read_csv  pandas {pandas_read * 1000:7.1f} ms   "
      f"polars {polars_read * 1000:6.1f} ms   {pandas_read / polars_read:4.1f}x")

# Streaming for larger-than-memory data
streamed = (
    pl.scan_csv('sales_data.csv')
    .filter(pl.col('price') > 50)
    .group_by('customer_id')
    .agg(pl.col('price').sum())
    .collect(engine='streaming')  # Process in chunks
)
print(f"streamed groups: {streamed.height}")

# Projection pushdown (read only needed columns)
narrow = pl.scan_csv('sales_data.csv').select(['customer_id', 'price']).collect()
print(f"projection pushdown read {narrow.width} of {pldf.width} columns")

import resource
import time

import pandas as pd

start = time.perf_counter()

df = pd.read_parquet('sales_large.parquet')
filtered = df[df['price'] > 100]
result = filtered.groupby('customer_id').agg(
    revenue=('price', 'sum'),
    order_count=('price', 'count'),
    avg_order_value=('price', 'mean'),
)
result.to_csv('high_value_customers.csv')

print(f"time:         {time.perf_counter() - start:.1f}s")
print(f"peak RSS:     {resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024**2:.1f} GB")
print(f"rows written: {len(result)}")

import resource
import time

import dask.dataframe as dd

start = time.perf_counter()

ddf = dd.read_parquet('sales_large.parquet')
filtered = ddf[ddf['price'] > 100]
result = filtered.groupby('customer_id').agg({
    'price': ['sum', 'count', 'mean']
}).compute()
result.to_csv('high_value_customers.csv')

print(f"time:         {time.perf_counter() - start:.1f}s")
print(f"peak RSS:     {resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024**2:.1f} GB")
print(f"rows written: {len(result)}")

import resource
import time

import vaex

start = time.perf_counter()

df = vaex.open('sales_large.parquet')  # Instant, memory-mapped
filtered = df[df.price > 100]
result = filtered.groupby('customer_id', agg={
    'revenue': vaex.agg.sum('price'),
    'order_count': vaex.agg.count(),
    'avg_order_value': vaex.agg.mean('price'),
})
result.export('high_value_customers.csv')

print(f"time:         {time.perf_counter() - start:.1f}s")
print(f"peak RSS:     {resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024**2:.1f} GB")
print(f"rows written: {len(result)}")

import resource
import time

import polars as pl

start = time.perf_counter()

df = pl.scan_parquet('sales_large.parquet')
result = (
    df
    .filter(pl.col('price') > 100)
    .group_by('customer_id')
    .agg([
        pl.col('price').sum().alias('revenue'),
        pl.len().alias('order_count'),
        pl.col('price').mean().alias('avg_order_value')
    ])
    .collect()
)
result.write_csv('high_value_customers.csv')

print(f"time:         {time.perf_counter() - start:.1f}s")
print(f"peak RSS:     {resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024**2:.1f} GB")
print(f"rows written: {result.height}")