Some last minute things

Signed-off-by: Adam Gutglick <[email protected]>

Author: AdamGS

bench-vortex/src/engines/df/mod.rs

@@ -11,7 +11,6 @@ use datafusion::execution::SessionStateBuilder;
 use datafusion::execution::cache::cache_manager::CacheManagerConfig;
 use datafusion::execution::cache::cache_unit::{DefaultFileStatisticsCache, DefaultListFilesCache};
 use datafusion::execution::runtime_env::RuntimeEnvBuilder;
-use datafusion::physical_plan::collect;
 use datafusion::prelude::{SessionConfig, SessionContext};
 use datafusion_common::GetExt;
 use datafusion_physical_plan::ExecutionPlan;
@@ -135,11 +134,10 @@ pub async fn execute_query(
     ctx: &SessionContext,
     query: &str,
 ) -> anyhow::Result<(Vec<RecordBatch>, Arc<dyn ExecutionPlan>)> {
-    let plan = ctx.sql(query).await?;
-    let (state, plan) = plan.into_parts();
-    let physical_plan = state.create_physical_plan(&plan).await?;
+    let df = ctx.sql(query).await?;
+    let physical_plan = df.clone().create_physical_plan().await?;
+    let result = df.collect().await?;
 
-    let result = collect(physical_plan.clone(), state.task_ctx()).await?;
     Ok((result, physical_plan))
 }
 

vortex-datafusion/src/persistent/mod.rs

@@ -205,36 +205,28 @@ mod tests {
 
         session
             .sql(&format!(
-                "CREATE EXTERNAL TABLE my_tbl_vx \
+                "CREATE EXTERNAL TABLE my_tbl \
                 (c1 VARCHAR NOT NULL, c2 INT NOT NULL) \
-                STORED AS vortex  \
+                STORED AS VORTEX  \
                 WITH ORDER (c1 ASC)
-                LOCATION '{}/vx/'",
+                LOCATION '{}/'",
                 dir.path().to_str().unwrap()
             ))
             .await?;
 
         session
-            .sql("INSERT INTO my_tbl_vx VALUES ('air', 5), ('balloon', 42)")
+            .sql("INSERT INTO my_tbl VALUES ('air', 10), ('alabama', 20), ('balloon', 30)")
             .await?
             .collect()
             .await?;
 
         session
-            .sql("INSERT INTO my_tbl_vx VALUES ('zebra', 5)")
+            .sql("INSERT INTO my_tbl VALUES ('kangaroo', 11), ('zebra', 21)")
             .await?
             .collect()
             .await?;
 
-        session
-            .sql("INSERT INTO my_tbl_vx VALUES ('texas', 2000), ('alabama', 2000)")
-            .await?
-            .collect()
-            .await?;
-
-        let df = session
-            .sql("SELECT * FROM my_tbl_vx ORDER BY c1 ASC limit 3")
-            .await?;
+        let df = session.sql("SELECT * FROM my_tbl ORDER BY c1 ASC").await?;
 
         let physical_plan = df.clone().create_physical_plan().await?;
 
@@ -243,27 +235,28 @@ mod tests {
         ┌───────────────────────────┐
         │  SortPreservingMergeExec  │
         │    --------------------   │
-        │  c1 ASC NULLS LASTlimit:  │
-        │             3             │
+        │     c1 ASC NULLS LAST     │
         └─────────────┬─────────────┘
         ┌─────────────┴─────────────┐
         │       DataSourceExec      │
         │    --------------------   │
-        │          files: 3         │
+        │          files: 2         │
         │       format: vortex      │
         └───────────────────────────┘
         ");
 
         let r = df.collect().await?;
 
         insta::assert_snapshot!(pretty_format_batches(&r)?.to_string(), @r"
-        +---------+------+
-        | c1      | c2   |
-        +---------+------+
-        | air     | 5    |
-        | alabama | 2000 |
-        | balloon | 42   |
-        +---------+------+
+        +----------+----+
+        | c1       | c2 |
+        +----------+----+
+        | air      | 10 |
+        | alabama  | 20 |
+        | balloon  | 30 |
+        | kangaroo | 11 |
+        | zebra    | 21 |
+        +----------+----+
         ");
 
         Ok(())