Integrating VegaFusion with distributed DataFrames

[OlegWock]

I use VegaFusion to chart big DataFrames. In addition to Pandas, I want to support PySpark DataFrames. But since they might be too big to fit into the memory of a single machine, I want to outsource actual aggregation from VegaFusion directly to Spark. I'm aware that VegaFusion doesn't support this natively (as far as I understand), but I was wondering if you could point me in a direction I should look into?

As I understand, in VF 1.x it could be done by using "SQL API" (SqlConnection​ and SqlDataset​), but it was removed from VF 2.x. 

Initially I was thinking about using build_pre_transform_spec_plan​ and then manually translating transformations from server spec into calls to PySpark API, but then I realized I'll be re-implementing a lot of functionality already implemented in VF (e.g. Vega expressions parser). I saw you mentioning new DataFusion's unparse API and that it would power new DuckDB integration. But as far as I understand, it's not the case currently, DuckDB tables are just converted to Arrow and passed to VF. Was there any issue with the unparse API, or did you decide to not implement it due to other factors?

[jonmmease]

Hi @OlegWock, as you noted we had some architecture in place in version 1 for generating SQL for different dialects. This was removed for v2 because it was a significant source of complexity (given how it worked) and wasn't being used beyond duckdb. And after the optimizations in VegaFusion 2 (and probably in DataFusion as well), the duckdb path wasn't faster, so I didn't put the work into exploring the unparse workflow after all.

That said, I'm definitely interested in bringing this functionality back if there's demand and we develop a cleaner architecture.

Here's what I'm thinking might work.

We still convert the Vega transforms to DataFusion DataFrame operations, but instead of only evaluating these to arrow at the end of the transform pipeline with collect_to_table() here, we would extract the DataFusion LogicalPlan from the DataFrame. Given this logical plan, there are a few things you could do with it for evaluation.

1. Convert to sql with DataFusion's unparser support (though I don't see support for the Spark dialect here right now)
2. Convert to subtrait, and look into using the substrait-java to create a spark plan.
3. Process the logical plan directly and map it to the Spark DataFrame API

I think what I'd propose is having VegaFusion restore the Datasets interface that was present in version 1. The sql dialects would be limited to what DataFusion's unparse supports, but there would be other variants that return the DataFusion logical plan, or substrait.

Take a look and see what you think!

[jonmmease]

Hi @OlegWock, that sounds like great progress!

I noticed that in some places VF works with raw bytes in a synchronous manner through the VegaTable struct.

Yes, today for each individual Vega dataset we build a query that consists of the pipeline of transforms. Then evaluate the result to arrow (The VegaFusionTable) and store it in the cache. Derived datasets then pull this value from the cache.

I think the cleanest extension would be to add DataFusion's LogicalPlan as a new enum variant to TaskValue. A LogicalPlan is the declarative data structure that a DataFrame builds up. DataFrame's add additional info about the connection. Importantly, LogicalPlan's are hashable, and so will be easier to integrate into TaskValue than DataFrames. You can create a LogicalPlan from a DataFrame and vice versa.

So the flow would still be that a derived dataset creates a DataFrame based on it's input, but this DataFrame could be based on either a VegaFusionTable (as it is today), or a LogicalPlan.

With this approach I had to introduce a lot of repeating code, the new flow was almost identical to the existing one, just operating on different types. So this made me think that two flows might be unified. VF could operate on DataFrames by default and return DataFrame(s) as a result of spec execution.

Each transform in VegaFusion operates on, and returns, a DataFrame. So the only code that should need to change is the code that currently creates a DataFrame from VegaFusionTable (as this will also need to handle LogicalPlans), and the code that currently uses df.collect() to extract the in-memory Arrow record batches (this will need to optionally extract the logical plan rather than arrow). We'll need to work out how to decide which approach to use, but we could probably start with something simple where we only use LogicalPlan's for remote DataFrames.

Couple of things I'm not sure about are VegaFusionCallable::Data functions. Those seem to operate directly on data, and in the case of outsourcing aggregations, there won't be any data. So if I understand correctly, this might require updating the signal compiler to work on DataFrames and potentially returning DataFrame as well (instead of scalars) which then can be materialized/converted to plan by calling function.

The data functions are typically only called on relatively small inline datasets (in particular for how Vega-Lite manages selections). I picture that we will continue to use DataFusion's regular arrow approach for datasets that don't original as a spark tables, and in this case I think it would be pretty unusual for a Vega specs to try to call a data function on a dataset derived from these. And when it does happen, we can decide whether it's an error, or it results in pulling data into memory.

I believe with a bit of tinkering with LogicalPlan (before converting to SQL) and/or making custom Spark dialect for DataFusion would make it work.

I'd prefer going down the path of contributing a custom Spark dialect to DataFusion. I'd like to keep dialect specific code out of VegaFusion if at all possible (this got overwhelming in VegaFusion 1.x). But we can certainly use a hack to get the sql working during development so that we can test things end-to-end.

What do you think of this approach?

I think your general approach above makes sense, and I'm impressed you got a POC working so quickly! See if my comments above make sense, and then we can work through a plan.

[OlegWock]

I'd prefer going down the path of contributing a custom Spark dialect to DataFusion. I'd like to keep dialect specific code out of VegaFusion if at all possible (this got overwhelming in VegaFusion 1.x). But we can certainly use a hack to get the sql working during development so that we can test things end-to-end.

Oh yeah, definitely. We might want to include those tweaks in our fork of VF, but there are no reason to put it into upstream. I think it makes sense to support most generic usecase (i.e. logical plans), at least for the start

The approach with logical plan sounds good. I'll let you know how it goes. Thanks for your help.

[OlegWock]

Hello, @jonmmease . The approach of adding a logical plan as an option to task value worked really well. My code still needs a lot of polish and tests before I'm comfortable showing it to people, but overall it works surprisingly (to me at least) well. I have one more question, though 😅

I just noticed that bin transformation does eager evaluation of its "extent" signal here and then uses the resulting data to calculate bin parameters. This isn't compatible with the logical plan approach. We can create a logical plan to evaluate the original extent signal, but since its results are used to build subsequent logical plan (for binning), we need to either evaluate it or change how binning works.

Initially I thought maybe we can just rewrite calculate_bin_params into datafusion expressions, so instead of eagerly evaluating extent signal, it will just build on min and max expressions. The whole construction then will be included in the logical plan. But I quickly realized that calculate_bin_params is quite complex (branching, loop) and probably will yield very complicated SQL (if it can be translated into SQL at all).

Besides this approach, I had two ideas. One is to introduce a new parameter (callback) that will accept a logical plan and return a query result (perhaps as an arrow table). It will be stored on CompilationConfig, so transform like bin would be able to evaluate the extent signal in an "engine-agnostic" way. The default implementation will use DataFusion (same as currently), but users will be able to provide their own implementation to execute query in e.g. Spark. I'm not sure if this is a good approach tbh, it seems like too big of a band-aid for a small problem.

Another one was to make binning into UDF. It will accept field (which will be used for binning), extent, as well as other optional parameters like step or span and calculate which bin this row should go into. VF will provide implementation for DataFusion, but to execute the exported LogicalPlan, the user will be required to provide their own implementation in the engine (or rewrite the logical plan). This looks cleaner to me compared to the previous option, but I'm still not sure maybe there is a better option.

So I wanted to ask your opinion, maybe you have ideas how to work around this in the best way?

[jonmmease]

Hi @OlegWock, you're right that binning is a bit of special case in needing to perform a query to set up the transform. The pivot transform works this way as well.

I was picturing that we would end up with some logic to take a DataFusion DataFrame and evaluate it, either using the regular DataFusion path, or using the path that unparses to SQL and then evaluates with some external sql engine. We'll need to run this to convert the dataset to arrow at some point in the pipeline. Thanks to including LogicalPlan's as task values, we won't need to run this at the end of every Data task, because we'll be able to store the logical plan as the task value. But we'll still need to run this at the end of the final leaf datasets, so that we can set the arrow results to the front and, or embed them in the pre transformed vega spec.

Given that we'll need this function, I was picturing that we'd use it inside the bin transform, in place of the df.collect() we do today. So this would trigger an additional spark query while building the logical plan, but I think this would work.

An alternative is to try to implement the binning logic in lazy form using DataFusion dataframes. This is possible, and I've sketched out a way to do it using pure sql in #468. But it's not straightforward. And we'd also need to keep in mind that the bin transform can output a signal that contains a scalar value with the bin extend info. This is also possible to do lazily using DataFusion's Expr::Subquery. So I think both approaches are possible, but the former would be a good bit simpler, at the cost of requiring a query during logical plan construction.

[OlegWock]

the former would be a good bit simpler, at the cost of requiring a query during logical plan construction

This makes sense

I was picturing that we would end up with some logic to take a DataFusion DataFrame and evaluate it, either using the regular DataFusion path, or using the path that unparses to SQL and then evaluates with some external sql engine. We'll need to run this to convert the dataset to arrow at some point in the pipeline.

Besides adding Plan as variant to TaskValue enum, I updated transforms to pass DF/Plan around for as long as possible, potentially LogicalPlan being result of whole transform pipeline. Then, depending on method which evaluated this pipeline, it could call runtime.materialize_export_updates that would convert ExportUpdate (with TaskValue) into ExportUpdateArrow (MaterializedTaskValue) by executing query with DataFusion. If user wanted to execute query in different engine, there is pre_transform_logical_plan which just returns logical plan, which user can take and do whatever they want. But they would have to e.g. manually stitch it into spec, instead of using pre_transform_spec.

With introduction of this new callback, we can use it instead of collecting DataFrame directly, this will also allow for methods like runtime.pre_transform_spec to work with custom engines as well, that would be much better UX/DX I think.

One thing I'm concerned about is resulting data size. One of main reasons to use Spark is to work with data that doesn't effectively fits into single machine, so there is a risk of trying to pull too much data into VF when executing query. But as long as we use it only for things like extent or to materialize final results it should be fine.