round trip TPCH queries in tests (#630)

* honor table name for csv/parquet scan in ballista plan serde

* disable query 7,8,9 in ballista integration test

* add tpch query ballista roundtrip test

* also roud trip physical plan

* fix clippy

* simplify test code

Author: QP Hou

benchmarks/Cargo.toml

@@ -39,3 +39,6 @@ futures = "0.3"
 env_logger = "^0.8"
 mimalloc = { version = "0.1", optional = true, default-features = false }
 snmalloc-rs = {version = "0.2", optional = true, features= ["cache-friendly"] }
+
+[dev-dependencies]
+ballista-core = { path = "../ballista/rust/core" }

benchmarks/src/bin/tpch.rs

@@ -573,7 +573,6 @@ mod tests {
 
     use datafusion::arrow::array::*;
     use datafusion::arrow::util::display::array_value_to_string;
-
     use datafusion::logical_plan::Expr;
     use datafusion::logical_plan::Expr::Cast;
 
@@ -1042,4 +1041,88 @@ mod tests {
 
         Ok(())
     }
+
+    mod ballista_round_trip {
+        use super::*;
+        use ballista_core::serde::protobuf;
+        use datafusion::physical_plan::ExecutionPlan;
+        use std::convert::TryInto;
+
+        fn round_trip_query(n: usize) -> Result<()> {
+            let config = ExecutionConfig::new()
+                .with_concurrency(1)
+                .with_batch_size(10);
+            let mut ctx = ExecutionContext::with_config(config);
+
+            // set tpch_data_path to dummy value and skip physical plan serde test when TPCH_DATA
+            // is not set.
+            let tpch_data_path =
+                env::var("TPCH_DATA").unwrap_or_else(|_| "./".to_string());
+
+            for &table in TABLES {
+                let schema = get_schema(table);
+                let options = CsvReadOptions::new()
+                    .schema(&schema)
+                    .delimiter(b'|')
+                    .has_header(false)
+                    .file_extension(".tbl");
+                let provider = CsvFile::try_new(
+                    &format!("{}/{}.tbl", tpch_data_path, table),
+                    options,
+                )?;
+                ctx.register_table(table, Arc::new(provider))?;
+            }
+
+            // test logical plan round trip
+            let plan = create_logical_plan(&mut ctx, n)?;
+            let proto: protobuf::LogicalPlanNode = (&plan).try_into().unwrap();
+            let round_trip: LogicalPlan = (&proto).try_into().unwrap();
+            assert_eq!(
+                format!("{:?}", plan),
+                format!("{:?}", round_trip),
+                "logical plan round trip failed"
+            );
+
+            // test optimized logical plan round trip
+            let plan = ctx.optimize(&plan)?;
+            let proto: protobuf::LogicalPlanNode = (&plan).try_into().unwrap();
+            let round_trip: LogicalPlan = (&proto).try_into().unwrap();
+            assert_eq!(
+                format!("{:?}", plan),
+                format!("{:?}", round_trip),
+                "opitmized logical plan round trip failed"
+            );
+
+            // test physical plan roundtrip
+            if env::var("TPCH_DATA").is_ok() {
+                let physical_plan = ctx.create_physical_plan(&plan)?;
+                let proto: protobuf::PhysicalPlanNode =
+                    (physical_plan.clone()).try_into().unwrap();
+                let round_trip: Arc<dyn ExecutionPlan> = (&proto).try_into().unwrap();
+                assert_eq!(
+                    format!("{:?}", physical_plan),
+                    format!("{:?}", round_trip),
+                    "physical plan round trip failed"
+                );
+            }
+
+            Ok(())
+        }
+
+        macro_rules! test_round_trip {
+            ($tn:ident, $query:expr) => {
+                #[test]
+                fn $tn() -> Result<()> {
+                    round_trip_query($query)
+                }
+            };
+        }
+
+        test_round_trip!(q1, 1);
+        test_round_trip!(q3, 3);
+        test_round_trip!(q5, 5);
+        test_round_trip!(q6, 6);
+        test_round_trip!(q10, 10);
+        test_round_trip!(q12, 12);
+    }
 }

datafusion/src/datasource/mod.rs

@@ -28,6 +28,7 @@ pub use self::csv::{CsvFile, CsvReadOptions};
 pub use self::datasource::{TableProvider, TableType};
 pub use self::memory::MemTable;
 
+/// Source for table input data
 pub(crate) enum Source<R = Box<dyn std::io::Read + Send + Sync + 'static>> {
     /// Path to a single file or a directory containing one of more files
     Path(String),