Building a DataFusion CSV reader with arrow-extendr

Want to skip to the end with the source code? Click here.

Goal

For this tutorial we’re going to create a very simple Rust-based R package using extendr and arrow-extendr. The package will use the new and very powerful DataFusion crate to create a csv reader.

We’ll learn a little bit about how extendr, Rust, and arrow-extendr works along the way.

Create a new R package

We will use {usethis} to create a new R package called dfusionrdr (prnounced d-fusion reader).

usethis::create_package("dfusionrdr")

This will open a new R project with the scaffolding of an R package. From here, we need to make the R package into an extendr R package. To do so we userextendr::use_extendr().

rextendr::use_extendr()

Building your package

Once you’ve initialized extendr in your package, we can check to see if everything worked by running the hello_world() function that is included. To do so, we can build our package, then document it.

To make R functions from Rust available into R, we run rextendr::document().

Run devtools::load_all() to bring the documented functions into scope and run the function!

devtools::load_all()
hello_world()
#> [1] "Hello world!"

We’ve now ran Rust code directly from R. Pretty simple Rust, but Rust nonetheless.

Adding Rust dependencies

Much like how we like to use R packages to make our lives easier, we can use Rust crates (libraries) to make do crafty things. To do so, we will open up our Rust crate in our preferred editor. I prefer VS Code.

Open src/rust/ in VS Code. Then we will add 3 additional dependencies. These are

• datafusion
  • a powerful Arrow-based DataFrame library (like Polars but different)
• tokio
  • which will give us the ability to run code lazily and asynchronously which is required by datafusion
• arrow_extendr
  • this is a crate I built that lets us send Arrow data from Rust to R and back

In the terminal run the following

cargo add datafusion
cargo add tokio
cargo add arrow_extendr --git https://github.com/JosiahParry/arrow-extendr

Making R work with DataFusion

DataFusion requires one additional C library that we need to use we need to add it to our Makevars. This is not something you typically have to do, but DataFusion requires it from us.

Open Makevars and Makevars.win. One the line that starts with PKG_LIBS add -llzma to the end.

Again, this is not a common thing you have to do. This is specifically for our use case.

Building our CSV Reader

Open src/lib.rs. This is where your R package is defined. For larger packages you may want to break it up into multiple smaller files. But our use case is relatively small (and frankly, not that simple, lol!).

Let’s first start by removing our hello_world example from our code. Delete the hello world function (lines 3-8) and remove it from the module declaration under mod dfusionrdr.

Let’s create the scaffolding for our first function read_csv_dfusion()

#[extendr]
/// @export 
fn read_csv(csv_path: &str)  {
  rprintln!("We will read the csv file at: `{csv_path}`");
}

This function prints a message indicating we will read a CSV at the path provided. It takes one argument csv_path which is an &str. A &str in Rust is a like a scalar character in R e.g. "my-file.csv"

Next we need to make sure the function is available to R in the module.

extendr_module! {
    mod dfusionrdr;
    fn read_csv_dfusion;
}

From RStudio, let’s build, document, and load again.

devtools::build()    # 1. 
rextendr::document() # 2. 
devtools::load_all() # 3.

Import dependencies

In order to use DataFusion to read dependencies we need to import it. A lot of Rust libraries have something called a prelude. The prelude is a special module that contains common structs, traits, enums, etc that are very useful for the crate. Notice that the top of your lib.rs includes use extendr_api::prelude::*; this brings all of the Rust based R objects into scope such as Robj, Doubles, Integers etc.

DataFusion also has a useful prelude that we want to bring into scope. We will add use datafusion::prelude::*; to the top of our file (much like adding library()). This brings important objects into scope for us. We will also need tokio::runtime::Runtime as well.

The first 3 lines of your lib.rs should look like this:

use extendr_api::prelude::*;
use datafusion::prelude::*;
use tokio::runtime::Runtime;

Context and Runtime

DataFusion requires something called a SessionContext. The session context

“maintains the state of the connection between a user and an instance of the DataFusion engine.”

We need to instantiate this struct inside of our function.

fn read_csv_dfusion(csv_path: &str) {
    let ctx = SessionContext::new();
}

We now have a ctx object which we can use to read our csv. It has a method called read_csv(). It requires the path of a csv to read as well as a struct called CsvReadOptions which determines how it will be read into memory. We will pass csv_path to the first argument and create a default options struct with the new() method.

fn read_csv_dfusion(csv_path: &str) {
    let ctx = SessionContext::new();
    let csv = ctx.read_csv(
        csv_path,
        CsvReadOptions::new()
    );
}

This will compile with a bunch of warnings about unused variables. But, more importantly, the csv variable we created is special. If you have your Rust analyzer configured you should see that it is of type impl Future<Output = Result <..., ...>>. That right there is problematic!

When you see impl Future<...> that tells us it is an asynchronous result that needs to be polled and executed. async functions are lazy. They don’t do anything until you ask it to. The way to do this is by calling the .await attribute. We can then unwrap() the results and store it into another variable.

fn read_csv_dfusion(csv_path: &str) {
    let ctx = SessionContext::new();
    let csv = ctx.read_csv(
        "sdf",
        CsvReadOptions::new()
    );

    let csv_result = csv.await.unwrap();
}

If we run cargo check in our terminal we will get the message:

error[E0728]: `await` is only allowed inside `async` functions and blocks

One way to get this to work would be to add async fn instead of fn but that isn’t supported by extendr since R is single threaded and doesn’t support async. So how do we get around this?

async with extendr and tokio

In order to run async functions we need to execute it in a runtime. tokio provides this for us with the Runtime struct. It lets us run impl Future<...> in a non async function!

We’ll modify our function definition to

fn read_csv_dfusion(csv_path: &str) {
    let rt = Runtime::new().unwrap();
    let ctx = SessionContext::new();
    let csv = ctx.read_csv(
        csv_path,
        CsvReadOptions::new()
    );
}

With the Runtime object rt we can call the block_on() method which takes a Future and runs it until it has completed. This means that we don’t get to use async functionality—e.g. executing 2 or more things at the same time—but we still get to take the result!

Let’s read the csv into an object called df using the block_on() method.

fn read_csv_dfusion(csv_path: &str) {
    let rt = Runtime::new().unwrap();
    let ctx = SessionContext::new();
    let csv = ctx.read_csv(
        csv_path,
        CsvReadOptions::new()
    );

    let df = rt.block_on(csv).unwrap();
}

The analyzer shows that this is a DataFrame. Awesome! Now, how can we get this into memory?

Sending DataFrames to R with arrow-extendr

This is where arrow-extendr comes into play. arrow-extendr provides a couple of traits which allow us to convert a number of arrow-rs types into an Robj.

The IntoArrowRobj trait can convert a Vec<RecordBatch> into an Robj. The R documentation for a RecordBatch says

“A record batch is a collection of equal-length arrays matching a particular Schema. It is a table-like data structure that is semantically a sequence of fields, each a contiguous Arrow Array.”

Based on that, a Vec<RecordBatch> is a collection of chunks of a table-like data structures.

DataFrames have a method .collect() which creates a Vec<RecordBatch>.

Let’s modify our function to turn the DataFrame into a Vec<RecordBatch>.

fn read_csv_dfusion(csv_path: &str) {
    let rt = Runtime::new().unwrap();
    let ctx = SessionContext::new();
    let csv = ctx.read_csv(
        csv_path,
        CsvReadOptions::new()
    );

    // create a dataframe from the csv
    let df = rt.block_on(csv).unwrap();

    // collect the results into record batches
    let res = rt.block_on(df.collect()).unwrap();
}

With this, we can send the results to R with the into_arrow_robj() method! First we need to add use arrow_extendr::to::IntoArrowRobj; to the top of our script to bring the trait into scope.

Then in our function we need to specify the return type as Robj (see the first line of the definition -> Robj) and then turn res into an Robj

fn read_csv_dfusion(csv_path: &str) -> Robj {
    let rt = Runtime::new().unwrap();
    let ctx = SessionContext::new();
    let csv = ctx.read_csv(
        csv_path,
        CsvReadOptions::new()
    );

    // create a dataframe from the csv
    let df = rt.block_on(csv).unwrap();

    // collect the results into record batches
    let res = rt.block_on(df.collect()).unwrap();

    res.into_arrow_robj().unwrap()
}

Handling arrow-rs from R

Let’s rebuild and document our function again.

I’ve added a csv of {palmerpenguins} to the inst/ folder of our package for testing. Let’ try reading this in.

res <- read_csv_dfusion("inst/penguins.csv")
res
#> <nanoarrow_array_stream struct<rowid: int64, species: string, island: string, bill_length_mm: string, bill_depth_mm: string, flipper_length_mm: string, body_mass_g: string, sex: string, year: int64>>
#>  $ get_schema:function ()  
#>  $ get_next  :function (schema = x$get_schema(), validate = TRUE)  
#>  $ release   :function ()  

Now, this doesn’t look very familiar to most R users. This is an object from the {nanoarrow} R package called "nanoarrow_array_stream". This is how data is received from Rust in R. We can process batches from this “stream” using the method get_next(). But there’s a handy as.data.frame() method for it.

res <- read_csv_dfusion("inst/penguins.csv")
as.data.frame(res) |> 
  head()
#>   rowid species    island bill_length_mm bill_depth_mm flipper_length_mm
#> 1     1  Adelie Torgersen           39.1          18.7               181
#> 2     2  Adelie Torgersen           39.5          17.4               186
#> 3     3  Adelie Torgersen           40.3            18               195
#> 4     4  Adelie Torgersen             NA            NA                NA
#> 5     5  Adelie Torgersen           36.7          19.3               193
#> 6     6  Adelie Torgersen           39.3          20.6               190
#>   body_mass_g    sex year
#> 1        3750   male 2007
#> 2        3800 female 2007
#> 3        3250 female 2007
#> 4          NA     NA 2007
#> 5        3450 female 2007
#> 6        3650   male 2007

Boom! We’ve written ourselves a reader! Let’s do a simple bench mark comparing it to readr.

library(dfusionrdr)
bench::mark(
  datafusion = as.data.frame(read_csv_dfusion("inst/penguins.csv")),
  readr = readr::read_csv("inst/penguins.csv"),
  check = FALSE
)
#> # A tibble: 2 × 6
#>   expression      min   median `itr/sec` mem_alloc `gc/sec`
#>   <bch:expr> <bch:tm> <bch:tm>     <dbl> <bch:byt>    <dbl>
#> 1 datafusion  922.3µs   1.04ms     915.     1.07MB      0  
#> 2 readr        14.6ms  15.31ms      65.6    13.2MB     80.8

Insanely fast!

Addendum

The source code for the entire package is below. It also includes a function read_sql_csv_dfusion() which takes a SQL statement and reads it into memory if you want to explore that. For example:

query <- 'SELECT count(*) as "n", "species" FROM csv GROUP BY "species"'

as.data.frame(
  read_sql_csv_dfusion(query, "inst/penguins.csv") 
)
#>     n   species
#> 1  68 Chinstrap
#> 2 152    Adelie
#> 3 124    Gentoo

Source code:

use arrow_extendr::to::IntoArrowRobj;
use extendr_api::prelude::*;
use datafusion::prelude::*;
use tokio::runtime::Runtime;
use std::result::Result;

#[extendr]
/// @export
fn read_csv_dfusion(csv_path: &str) -> Robj {
    let rt = Runtime::new().unwrap();
    let ctx = SessionContext::new();
    let csv = ctx.read_csv(
        csv_path,
        CsvReadOptions::new()
    );

    // create a dataframe from the csv
    let df = rt.block_on(csv).unwrap();

    // collect the results into record batches
    let res = rt.block_on(df.collect()).unwrap();

    res.into_arrow_robj().unwrap()
}

#[extendr]
fn read_sql_csv_dfusion(sql_query: &str, csv_path: &str) -> Result<Robj, Error> {
    let rt = Runtime::new().unwrap();
    let ctx = SessionContext::new();

    let csv_tbl_fut = ctx.register_csv(
        "csv",
        csv_path,
        CsvReadOptions::new(),
    );

    // // we dont assign it to anything because we're just ensuring that it gets ran
    let _ = rt.block_on(csv_tbl_fut).unwrap();

    // // create a plan
    let df = rt.block_on(ctx.sql(sql_query)).unwrap();

    // // collect into a vector of record batches
    let res = rt.block_on(df.collect()).unwrap();

    // return the result to R
    res.into_arrow_robj()
}

// Macro to generate exports.

// This ensures exported functions are registered with R.
// See corresponding C code in `entrypoint.c`.
extendr_module! {
    mod dfusionrdr;
    fn read_sql_csv_dfusion;
    fn read_csv_dfusion;
}

[package]
name = 'dfusionrdr'
publish = false
version = '0.1.0'
edition = '2021'

[lib]
crate-type = [ 'staticlib' ]
name = 'dfusionrdr'

[dependencies]
arrow_extendr = { git = "https://github.com/JosiahParry/arrow-extendr" }
datafusion = "33.0.0"
extendr-api = '*'
tokio = "1.34.0"