Skip to content

Commit

Permalink
feat: enable SQL query support in datafusion example (#169)
Browse files Browse the repository at this point in the history
Added support for executing SQL queries using DataFusion's parser and physical plan execution. This enhancement allows querying the "music" table with SQL statements, improving flexibility and functionality.
  • Loading branch information
loloxwg authored Sep 27, 2024
1 parent 837cc3f commit 6fb5a89
Showing 1 changed file with 28 additions and 5 deletions.
33 changes: 28 additions & 5 deletions examples/datafusion.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,8 +17,11 @@ use datafusion::{
error::{DataFusionError, Result},
execution::{RecordBatchStream, SendableRecordBatchStream, TaskContext},
physical_expr::EquivalenceProperties,
physical_plan::{DisplayAs, DisplayFormatType, ExecutionMode, ExecutionPlan, PlanProperties},
physical_plan::{
execute_stream, DisplayAs, DisplayFormatType, ExecutionMode, ExecutionPlan, PlanProperties,
},
prelude::*,
sql::parser::DFParser,
};
use fusio::{local::TokioFs, path::Path};
use futures_core::Stream;
Expand Down Expand Up @@ -225,9 +228,29 @@ async fn main() -> Result<()> {
let provider = MusicProvider { db: Arc::new(db) };
ctx.register_table("music", Arc::new(provider))?;

let df = ctx.table("music").await?;
let df = df.select(vec![col("name")])?;
let batches = df.collect().await?;
pretty::print_batches(&batches).unwrap();
{
let df = ctx.table("music").await?;
let df = df.select(vec![col("name")])?;
let batches = df.collect().await?;
pretty::print_batches(&batches).unwrap();
}

{
// support sql query for tonbo
let statements = DFParser::parse_sql("select * from music")?;
let plan = ctx
.state()
.statement_to_plan(statements.front().cloned().unwrap())
.await?;
ctx.execute_logical_plan(plan).await?;
let df = ctx.table("music").await?;
let physical_plan = df.create_physical_plan().await?;
let mut stream = execute_stream(physical_plan, ctx.task_ctx())?;
while let Some(maybe_batch) = stream.next().await {
let batch = maybe_batch?;
pretty::print_batches(&[batch]).unwrap();
}
}

Ok(())
}

0 comments on commit 6fb5a89

Please # to comment.