Distributed Python Tasks - Teradata AI Studio

Use Ray remote functions to parallelize Python work across workers.

import ray
import polars as pl

RAY_ADDRESS = "ray://sit-demo-ray-cluster-head-svc.ray.svc.cluster.local:10001"
ray.init(RAY_ADDRESS, ignore_reinit_error=True)

# Public NYC Taxi parquet files
PARTITION_PATHS = [
    "https://d37ci6vzurychx.cloudfront.net/trip-data/yellow_tripdata_2024-01.parquet",
    "https://d37ci6vzurychx.cloudfront.net/trip-data/yellow_tripdata_2024-02.parquet",
    "https://d37ci6vzurychx.cloudfront.net/trip-data/yellow_tripdata_2024-03.parquet",
]

@ray.remote
def analyze_taxi_monthly_data(url: str) -> dict:
    """
    Runs inside a Ray worker process.
    Downloads one month of NYC taxi data and computes trip metrics.
    """
    import polars as pl

    df = pl.read_parquet(url)

    result = (
        df.filter(
            (pl.col("trip_distance") > 0) &
            (pl.col("fare_amount") > 0) &
            (pl.col("passenger_count") > 0)
        )
        .with_columns(
            pl.col("tpep_pickup_datetime").dt.month().alias("month")
        )
        .group_by("month")
        .agg([
            pl.col("fare_amount").mean().round(2).alias("avg_fare"),
            pl.col("trip_distance").mean().round(2).alias("avg_distance_miles"),
            pl.col("passenger_count").mean().round(2).alias("avg_passengers"),
            pl.col("fare_amount").sum().round(2).alias("total_revenue"),
            pl.len().alias("total_trips"),
        ])
    )

    return result.to_dicts()


# process data for all 3 months in parallel
futures = [analyze_taxi_monthly_data.remote(url) for url in PARTITION_PATHS]

# wait for all workers to complete
results = ray.get(futures)

# display results
final = (
    pl.from_dicts([row for month in results for row in month])
    .sort("month")
)

print(final)