Custom Data Aggregators

In this tutorial, we will guide you through the process of building custom data aggregator and integrating them into the data flow. Specifically, we will cover how to:

1. Create and configure a custom aggregator to compute dataset-wide statistics.

2. Apply the custom aggregator to a dataset.

%config InlineBackend.figure_format = 'svg'

Imports

In this section, we will import all the necessary libraries and components that we need for creating and applying a custom data aggregator.

import datasets
from hyped import DataFlow
from hyped.typing import Float, Int
from hyped.core import BaseDataAggregator, BaseDataAggregatorConfig, RunContext, process_mode

Creating the Dataset and Data Flow instance

In this section, we will create a sample dataset and instantiate the DataFlow object that will manage the data transformations.

# Create a sample dataset
ds = datasets.Dataset.from_dict({"x": [0, 1, 2, 3, 4]})

# Create a data flow instance
flow = DataFlow(ds.features)

Implementing the Custom Aggregator

In this section, we’ll implement a custom data aggregator using hyped. The example below demonstrates how to create a custom aggregator that calculates the running average of a dataset.

# Define configuration for the custom aggregater
class CustomAggregatorConfig(BaseDataAggregatorConfig):
    ...

# Implement the custom aggregater
class CustomAggregator(BaseDataAggregator[CustomAggregatorConfig]):

    def seed(self, ctx: RunContext) -> tuple[Float, int]:
        return 0.0, 0

    async def extract(self, ctx: RunContext, x: Int) -> tuple[int, int]:
        # receives a batch of data and extracts the relevant aggregation data for the update
        return sum(x), len(x)

    async def update(self, ctx: RunContext, val: Float, state: int, extracted: tuple[int, int]) -> tuple[Float, int]:
        # update the aggregated value and state
        s, n = extracted
        return (val * state + s) / (state + n), state + n

The initialize method sets up the aggregator’s initial state. In this case:

• The first value represents the aggregated result (initially 0.0 for the average).

• The second value, refered to as state, tracks the total count of elements processed (initially 0).

The extract method processes a batch of input data, preparing it for aggregation. It computes

• the sum (sum(x)) of the batch

• the count (len(x)) of the batch

both will be used to update the aggregated state.

The update method updates the aggregation state with the extracted data.

• It calculates the new average by incorporating the previous aggregate (val * state), the new sum(s), and their respective weights (state and n).

• It also updates the count (state -> state + n) for future calculations.

Applying the Custom Aggregator

In this section, we will apply the custom aggregator to the dataset within the data flow graph and observe the results.

# Apply the custom aggregator
avg = CustomAggregator().call(x=flow.source["x"])

# Apply the aggregation to the dataset
out_ds, aggregates = flow.apply(ds, collect=flow.source, aggregate={"avg(x)": avg})

# Display the aggregated value
print(aggregates)

{'avg(x)': 2.0}

The process_mode for Aggregators

The process_mode decorator is a powerful feature in hyped that allows you to control the data format passed to the extract and update methods in an aggregator. By default:

• The extract method operates on batches of data, allowing for efficient batch-level processing.

• The update method operates on single instances extracted from the batch, ensuring precise, step-by-step updates to the aggregated state.

This default behavior strikes a balance between efficiency and thread safety. However, you can override this behavior using the process_mode decorator to customize how these methods handle data. Here’s an example:

# Define configuration for the custom aggregator
class CustomAggregatorConfigV2(BaseDataAggregatorConfig):
    ...

# Implement the custom aggregator
class CustomAggregatorV2(BaseDataAggregator[CustomAggregatorConfigV2]):

    def seed(self, ctx: RunContext) -> tuple[Float, int]:
        return 0.0, 0

    @process_mode(batched=False)
    async def extract(self, ctx: RunContext, x: Int) -> tuple[int, int]:
        # receives a batch of data and extracts the relevant aggregation data for the update
        return x, 1

    async def update(self, ctx: RunContext, val: Float, state: int, extracted: tuple[int, int]) -> tuple[Float, int]:
        # update the aggregated value and state
        s, n = extracted
        return (val * state + s) / (state + n), state + n

The @process_mode(batched=False) decorator is applied to the extract method, forcing it to process individual samples rather than batches.

The update method still operates on single instances, meaning it will be called for each extracted value from the input.

Applying this implementation of the aggregator leads to the exact same result as before:

# Apply the custom aggregator
avg = CustomAggregatorV2().call(x=flow.source["x"])

# Apply the aggregation to the dataset
out_ds, aggregates = flow.apply(ds, collect=flow.source, aggregate={"avg(x)": avg})

# Display the aggregated value
print(aggregates)

{'avg(x)': 2.0}

Note on Changing the ``process_mode`` of the ``update`` function

It’s worth noting that changing the process_mode of the update function doesn’t have much practical effect in most scenarios. Even if you set @process_mode(batched=True) for the update function, it will receive batches containing a single sample. This behavior arises because the extract method produces individual extracted results for each input sample by default.