Examples

Distributed sorting

from typing import List

import numpy as np

from pargraph import GraphEngine, delayed, graph


@delayed
def filter_array(array: np.ndarray, low: float, high: float) -> np.ndarray:
    """
    Filter an array to include only values within the specified range.

    :param array: The input array to be filtered
    :param low: The lower bound of the range
    :param high: The upper bound of the range
    :return: A new array containing only the values within the specified range
    """
    return array[(array >= low) & (array <= high)]


@delayed
def sort_array(array: np.ndarray) -> np.ndarray:
    """
    Sort array

    :param array: The input array to be sorted
    :return: A new array containing the sorted values
    """
    return np.sort(array)


@delayed
def reduce_arrays(*arrays: np.ndarray) -> np.ndarray:
    """
    Reduce arrays by concatenating them

    :param arrays: A variable number of arrays to be concatenated
    :return: A new array containing all the concatenated values
    """
    return np.concatenate(arrays)


@graph
def map_reduce_sort(array: np.ndarray, partition_count: int) -> np.ndarray:
    """
    Map reduce sort

    :param array: The input array to be sorted
    :param partition_count: The number of partitions to divide the array into
    :return: A new array containing the sorted values
    """
    return reduce_arrays(
        *(
            sort_array(filter_array(array, i / partition_count, (i + 1) / partition_count))
            for i in range(partition_count)
        )
    )


@graph
def map_reduce_sort_recursive(
    array: np.ndarray, partition_counts: List[int], _low: float = 0, _high: float = 1
) -> np.ndarray:
    """
    Map reduce sort recursively

    :param array: The input array to be sorted
    :param partition_counts: A list of partition counts for recursive sorting
    :param _low: The lower bound of the values range (not supposed to be used publicly)
    :param _high: The upper bound of the values range (not supposed to be used publicly)
    :return: A new array containing the sorted values
    """
    if len(partition_counts) == 0:
        return sort_array(array)

    partition_count, *partition_counts = partition_counts

    sorted_partitions = []
    for i in range(partition_count):
        low = _low + (_high - _low) * (i / partition_count)
        high = _low + (_high - _low) * ((i + 1) / partition_count)
        sorted_partitions.append(map_reduce_sort_recursive(filter_array(array, low, high), partition_counts, low, high))

    return reduce_arrays(*sorted_partitions)


if __name__ == "__main__":
    task_graph, keys = map_reduce_sort_recursive.to_graph(partition_counts=[2, 2, 2]).to_dict(
        array=np.random.rand(1_000_000)
    )

    graph_engine = GraphEngine()
    print(graph_engine.get(task_graph, keys)[0])

Aggregating NYC taxi data

import pandas as pd

from pargraph import GraphEngine, delayed, graph


@delayed
def get_total_amount_sum(file_url, unit):
    """
    Calculate the sum of the 'total_amount' column from a Parquet file and divide it by a given unit.

    :param file_url: The URL of the Parquet file
    :param unit: The unit to divide the total amount sum by
    :return: The sum of the 'total_amount' column divided by the given unit
    """
    return pd.read_parquet(file_url, columns=["total_amount"])["total_amount"].sum() / unit


@delayed
def add(*args):
    """
    Calculate the sum of all the arguments provided.

    :param args: A variable number of arguments to be summed
    :return: The sum of all arguments
    """
    return sum(args)


@delayed
def collect_result(*results):
    """
    Collect results, using every other element as the index and the remaining elements as the values.

    :param results: Arguments where even-indexed elements are used as the index and odd-indexed elements as the values
    :return: A pandas Series with the specified index and values
    """
    return pd.Series(results[1::2], index=results[0::2])


URL = "https://d37ci6vzurychx.cloudfront.net/trip-data/yellow_tripdata_{year}-{month:02}.parquet"


@graph
def nyc_yellow_taxi_fare_total_for_year(year, unit):
    """
    Calculate the total fare amount for a given year by summing the total amounts from each month.

    :param year: The year for which to calculate the total fare amount
    :param unit: The unit to divide the total amount sum by
    :return: Total fare amount for the given year
    """
    return add(*(get_total_amount_sum(URL.format(year=year, month=month), unit) for month in range(1, 13)))


@graph
def nyc_yellow_taxi_fare_total_by_year(years, unit):
    """
    Calculate the total fare amount for multiple years.

    :param years: A list of years for which to calculate the total fare amount
    :param unit: The unit to divide the total amount sum by
    :return: Total fare amounts for the given years
    """
    results = []

    for year in years:
        results.append(year)
        results.append(nyc_yellow_taxi_fare_total_for_year(year, unit))

    return collect_result(*results)


if __name__ == "__main__":
    task_graph, keys = nyc_yellow_taxi_fare_total_by_year.to_graph(years=range(2024, 2025)).to_dict(unit=1_000_000)

    graph_engine = GraphEngine()
    print(graph_engine.get(task_graph, keys)[0])

Pythagoras theorem using Python operations

from pargraph import GraphEngine, graph


@graph
def pythagoras(a: float, b: float) -> float:
    """
    Calculate the length of the hypotenuse of a right triangle given the lengths of the other two sides.

    :param a: Length of one side of the triangle.
    :param b: Length of the other side of the triangle.
    :return: Length of the hypotenuse.
    """
    return (a**2 + b**2) ** 0.5


if __name__ == "__main__":
    task_graph, keys = pythagoras.to_graph().to_dict(a=3, b=4)

    graph_engine = GraphEngine()
    print(graph_engine.get(task_graph, keys)[0])