Fast queries on Parquet data in Rust
Several projects in banking and consumer products have made use of data stored in data lakes in Parquet, a compact and efficient column-oriented format. Python with Pandas is a feasible but very inefficient choice for this, and I have typically used Spark here. Recently revisting Rust for long-lived MCP servers, I have been exploring DataFusion for direct queries on Parquet data from Rust, and am very impressed.
DataFusion describes itself as an extensible query engine written in Rust that uses Apache Arrow as its in-memory format. The web site explains that Out of the box, DataFusion offers SQL and Dataframe APIs, excellent performance, built-in support for CSV, Parquet, JSON, and Avro, extensive customization, and a great community. It is basically a library to access these data formats and query them, either via chained declarative function calls, or via strings with SQL queries.
To test it, I used two years’ worth of the New York City taxi rides data set, which is made available in Parquet format. Download any number of months of the Parquet files, and put them into a directory called ‘data’. I wanted to test not just the ability to read and access the files, but to perform queries using a declarative syntax such as SQL, since that would be the enabler of MCP access to the data.
You can DataFusion into your projects by adding the following dependencies to your Cargo.toml:
[dependencies]
datafusion = "55.0.0"
tokio = { version = "1.0", features = ["rt-multi-thread"] }
The functions can be imported in your Rust code as follows:
use datafusion::arrow::array::RecordBatch;
use datafusion::arrow::util::pretty::pretty_format_batches;
use datafusion::error::Result;
use datafusion::prelude::*;
Then, the following code will register the tables, including reading the schema and other metadata (but not yet the data itself):
// Create a context and register the table, assumes data is in 'data' subdirectory of parent
let ctx = SessionContext::new();
ctx.register_parquet("rides", "../data", ParquetReadOptions::new())
.await?;
To run the query, execute the sql() function on the context, and either collect() or iterate the resulting rows:
// Create a plan to run a SQL query
let df = ctx.sql(q).await?;
// Execute and collect the results
let batches = df.collect().await?;
The above code collects the entire result, which makes sense for an MCP server, which returns the result of a request. If the results are very large, MCP allows pagination. And within Rust, you can also iterate row-by-row for programmatic aggregation or other processing.
Putting it together, here is a Rust command-line program that takes a query as the argument (must be in quotes), runs the query, and shows the result:
// Execute query on command line against Parquet files in the data directory,
// which is assumed to be under the parent directory.
//
// Sample usage: cargo run "select avg(fare_amount) from rides"
//
// Sample output:
// Query: select avg(fare_amount) from rides
// +------------------------+
// | avg(rides.fare_amount) |
// +------------------------+
// | 19.219927893175566 |
// +------------------------+
use datafusion::arrow::array::RecordBatch;
use datafusion::arrow::util::pretty::pretty_format_batches;
use datafusion::error::Result;
use datafusion::prelude::*;
#[tokio::main]
async fn main() -> Result<()> {
// Get query from the command line
let args: Vec<_> = std::env::args().collect();
if args.len() != 2 {
println!("Missing query or extra arguments");
std::process::exit(1);
}
let q = args[1].as_str();
println!("Query: {}", q);
// Execute query and show result
let batches = query(&q).await?;
println!("{}", pretty_format_batches(&batches)?);
Ok(())
}
// Execute a query, returning the resulting record batches
async fn query(q: &str) -> Result<Vec<RecordBatch>> {
// Register the table
let ctx = SessionContext::new();
ctx.register_parquet("rides", "../data", ParquetReadOptions::new())
.await?;
// Create a plan to run a SQL query
let df = ctx.sql(q).await?;
// Execute and collect the results
let batches = df.collect().await?;
Ok(batches)
}
In the next post, we’ll put this into an MCP server, to make fast Parquet data lake queries available as a tool to an LLM.