USING DATAFUSION'S UDAFS TO DO ML TRAINING

Introduction

User-Defined Aggregate Functions (udaf) are emerging as essential tools for sophisticated data analytics on a subset of data. In this post, we use XGboost model training as an example of a complex operation that is well-suited for the UDAFs computation. However, there are two challenges:

1. Creating a UDAF often involves hardcoding the aggregation computation process.
2. UDAF’s operations are generally not fully transparent to the query processor, limiting the system’s ability to optimize queries effectively. These challenges underscore the need for a more flexible and efficient approach to aggregating data, which we aim to address through our work with DataFusion.

Many systems provide udaf functionality via accumulators which receive incremental subsets of the data rather than all of it once. This can be more efficient if your operation is expressible as an accumulation (think of something like a sufficient statistic). However, for many modeling use cases, we need simultaneous access to all of the data to be processed. Our use case is gradient boosting via XGBoost, which does indeed need all the data at once.

Problem Statement

We have some loan data and want to find how the most salient features for predicting a target vary over time: training one XGBoost model per year to get the n “best features” per year.

We’d like the user experience to look as close as possible to pandas. Below is our pseudo-code target for how we might do it in pure pandas, how it changes when we do it in pandas via Ibis, and how we’d like to be able to do it in DataFusion via Ibis. We’ll use this as a guide to build our library.

And here’s code for a toy model that fleshes out the pure pandas pseudo code¹

Our goal is to teach DataFusion how to use the same curried_calc_best_features function for DataFusion.

Interlude

We’ll use Ibis to target both pandas and DataFusion. Ibis aims to be the portable Python dataframe library, allowing you to specify your computation once and target multiple engines, including query engines like DataFusion, and pandas. It is designed to be a high-level API for data manipulation and analysis, and can be used to build complex data processing pipelines.

In our case, we are using Ibis to register UDAFs with DataFusion and pandas, generate a deferred expression in a similar API to pandas and execute them on multiple backends. This requires us to register a udaf with Ibis for each backend.

For the pandas backend, we can automate udaf registration for a function with a known input table like this

We’d like to be able to create the udaf and invoke Ibis with the DataFusion backend in a similar way.

Implementation

We know what we want the user experience to feel like, but what is the DataFusion baseline?

Per the DataFusion example of a udaf, first, you define a class that inherits from Accumulator

You then register this class with DataFusion, specifying additional information (some types and volatility)

We can’t / don’t want to update and merge models, so what’s to be done? We could try to accumulate the rows into a ListArray<Struct>, but this would be computationally costly (think of it like doing a transpose, only worse). Instead, we’ll accumulate the data as serialized RecordBatchs. This is surprisingly cheap because of the underlying Arrow representation.

A first pass at our Accumulator might look something like this

What’s happening here?

• _states is a list of serialized StructArray
• state simply returns the accumulator’s list of serialized StructArray, _states, cast into its pyarrow type
• update combines the columns’ RecordBatchs into a single StructArray, serializes it with python’s pickle and extends _states
• merge simply extends _states with other Accumulators’ _states
• pystate converts all the data in _states into a single pandas DataFrame
• pyevaluate is where the computation actually happens: invoke our pure pandas style function on the DataFrame returned by pystate
• evaluate invokes pyevaluate and casts it’s return value to return_type

Stay DRY, my friends

Note that almost everything here is boiler-plate with exception of

• Our aggregation function: curried_calc_best_features
• The pyarrow type used to cast in evaluate
• The variable column_names used by update

Note also that we have to coordinate names and types between the Ibis table, udaf definition, and udaf invocation.

Let’s combine all of this into a single function that dynamically generates a class and registers it with both DataFusion and Ibis.

The Datafusion related section of our final library code will look like this

Now, via make_datafusion_udaf, we only need to provide

• the Ibis table the udaf will be invoked on (used to extract names and types)
• the function to invoke on each group’s pandas DataFrame
• the return type of the function

We can optionally provide a name for the udaf that DataFusion uses internally and Ibis will use for the SQL string it generates for the DataFusion backend.

But, does it work?

Ok, let’s invoke this on some real data! We’ll use the LendingClub data from Kaggle and try to determine which features are predictive of a loan not being “Fully Paid” by the end of its standard term (which we’ve already massaged into the data our target variable called event_occurred as well as some other derived variables like dti and cr_age_days).