feat: Support serde for JsonSource PhysicalPlan (#15311)

Author: westhide

datafusion/proto/proto/datafusion.proto

@@ -273,7 +273,6 @@ message DmlNode{
     INSERT_APPEND = 3;
     INSERT_OVERWRITE = 4;
     INSERT_REPLACE = 5;
-    
   }
   Type dml_type = 1;
   LogicalPlanNode input = 2;
@@ -726,6 +725,7 @@ message PhysicalPlanNode {
     CsvSinkExecNode csv_sink = 28;
     ParquetSinkExecNode parquet_sink = 29;
     UnnestExecNode unnest = 30;
+    JsonScanExecNode json_scan = 31;
   }
 }
 
@@ -1024,6 +1024,10 @@ message CsvScanExecNode {
   bool newlines_in_values = 7;
 }
 
+message JsonScanExecNode {
+  FileScanExecConf base_conf = 1;
+}
+
 message AvroScanExecNode {
   FileScanExecConf base_conf = 1;
 }

datafusion/proto/src/generated/pbjson.rs

@@ -33,7 +33,7 @@ use datafusion::datasource::file_format::parquet::ParquetSink;
 use datafusion::datasource::physical_plan::AvroSource;
 #[cfg(feature = "parquet")]
 use datafusion::datasource::physical_plan::ParquetSource;
-use datafusion::datasource::physical_plan::{CsvSource, FileScanConfig};
+use datafusion::datasource::physical_plan::{CsvSource, FileScanConfig, JsonSource};
 use datafusion::datasource::source::DataSourceExec;
 use datafusion::execution::runtime_env::RuntimeEnv;
 use datafusion::execution::FunctionRegistry;
@@ -247,6 +247,15 @@ impl AsExecutionPlan for protobuf::PhysicalPlanNode {
                 .with_file_compression_type(FileCompressionType::UNCOMPRESSED);
                 Ok(conf.build())
             }
+            PhysicalPlanType::JsonScan(scan) => {
+                let scan_conf = parse_protobuf_file_scan_config(
+                    scan.base_conf.as_ref().unwrap(),
+                    registry,
+                    extension_codec,
+                    Arc::new(JsonSource::new()),
+                )?;
+                Ok(scan_conf.build())
+            }
             #[cfg_attr(not(feature = "parquet"), allow(unused_variables))]
             PhysicalPlanType::ParquetScan(scan) => {
                 #[cfg(feature = "parquet")]
@@ -1684,6 +1693,26 @@ impl AsExecutionPlan for protobuf::PhysicalPlanNode {
             }
         }
 
+        if let Some(data_source_exec) = plan.downcast_ref::<DataSourceExec>() {
+            let data_source = data_source_exec.data_source();
+            if let Some(scan_conf) = data_source.as_any().downcast_ref::<FileScanConfig>()
+            {
+                let source = scan_conf.file_source();
+                if let Some(_json_source) = source.as_any().downcast_ref::<JsonSource>() {
+                    return Ok(protobuf::PhysicalPlanNode {
+                        physical_plan_type: Some(PhysicalPlanType::JsonScan(
+                            protobuf::JsonScanExecNode {
+                                base_conf: Some(serialize_file_scan_config(
+                                    scan_conf,
+                                    extension_codec,
+                                )?),
+                            },
+                        )),
+                    });
+                }
+            }
+        }
+
         #[cfg(feature = "parquet")]
         if let Some(exec) = plan.downcast_ref::<DataSourceExec>() {
             let data_source_exec = exec.data_source();

datafusion/proto/src/generated/prost.rs

@@ -1277,6 +1277,15 @@ fn roundtrip_analyze() -> Result<()> {
     )))
 }
 
+#[tokio::test]
+async fn roundtrip_json_source() -> Result<()> {
+    let ctx = SessionContext::new();
+    ctx.register_json("t1", "../core/tests/data/1.json", Default::default())
+        .await?;
+    let plan = ctx.table("t1").await?.create_physical_plan().await?;
+    roundtrip_test(plan)
+}
+
 #[test]
 fn roundtrip_json_sink() -> Result<()> {
     let field_a = Field::new("plan_type", DataType::Utf8, false);