diff --git a/rust/datafusion/tests/sql.rs b/rust/datafusion/tests/sql.rs index bbadae6acb389..2532194d5fe07 100644 --- a/rust/datafusion/tests/sql.rs +++ b/rust/datafusion/tests/sql.rs @@ -99,9 +99,11 @@ fn parquet_single_nan_schema() { ctx.register_parquet("single_nan", &format!("{}/single_nan.parquet", testdata)) .unwrap(); let sql = "SELECT mycol FROM single_nan"; - let relation = ctx.sql(&sql, 1024 * 1024).unwrap(); - let mut results = relation.borrow_mut(); - while let Some(batch) = results.next().unwrap() { + let plan = ctx.create_logical_plan(&sql).unwrap(); + let plan = ctx.optimize(&plan).unwrap(); + let plan = ctx.create_physical_plan(&plan, DEFAULT_BATCH_SIZE).unwrap(); + let results = ctx.collect(plan.as_ref()).unwrap(); + for batch in results { assert_eq!(1, batch.num_rows()); assert_eq!(1, batch.num_columns()); }