You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by al...@apache.org on 2022/03/23 09:29:58 UTC
[arrow-datafusion] branch master updated: Add `write_json`, `read_json`, `register_json`, and `JsonFormat` to `CREATE EXTERNAL TABLE` functionality (#2023)
This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git
The following commit(s) were added to refs/heads/master by this push:
new 59c6d93 Add `write_json`, `read_json`, `register_json`, and `JsonFormat` to `CREATE EXTERNAL TABLE` functionality (#2023)
59c6d93 is described below
commit 59c6d93a4a5ddb5edeb2a92b9ce00c046cedee50
Author: Matthew Turner <ma...@outlook.com>
AuthorDate: Wed Mar 23 05:29:53 2022 -0400
Add `write_json`, `read_json`, `register_json`, and `JsonFormat` to `CREATE EXTERNAL TABLE` functionality (#2023)
* Add json exec functionality
* Add write_json
* First commit
* Rebase and update tests
* Cargo fmt
* Cleanup
---
datafusion/src/dataframe.rs | 9 +-
datafusion/src/execution/context.rs | 77 ++++++++++++++---
datafusion/src/execution/options.rs | 16 +++-
datafusion/src/logical_plan/builder.rs | 49 +++++++++++
datafusion/src/physical_plan/file_format/json.rs | 95 ++++++++++++++++++++-
datafusion/src/physical_plan/file_format/mod.rs | 1 +
datafusion/tests/jsons/3.json | 1 +
datafusion/tests/sql/json.rs | 103 +++++++++++++++++++++++
datafusion/tests/sql/mod.rs | 1 +
9 files changed, 335 insertions(+), 17 deletions(-)
diff --git a/datafusion/src/dataframe.rs b/datafusion/src/dataframe.rs
index ae2cc76..28f845d 100644
--- a/datafusion/src/dataframe.rs
+++ b/datafusion/src/dataframe.rs
@@ -35,7 +35,7 @@ use crate::arrow::util::pretty;
use crate::datasource::TableProvider;
use crate::datasource::TableType;
use crate::execution::context::{SessionState, TaskContext};
-use crate::physical_plan::file_format::{plan_to_csv, plan_to_parquet};
+use crate::physical_plan::file_format::{plan_to_csv, plan_to_json, plan_to_parquet};
use crate::physical_plan::{collect, collect_partitioned};
use crate::physical_plan::{execute_stream, execute_stream_partitioned, ExecutionPlan};
use crate::scalar::ScalarValue;
@@ -576,6 +576,13 @@ impl DataFrame {
let state = self.session_state.read().clone();
plan_to_parquet(&state, plan, path, writer_properties).await
}
+
+ /// Executes a query and writes the results to a partitioned JSON file.
+ pub async fn write_json(&self, path: impl AsRef<str>) -> Result<()> {
+ let plan = self.create_physical_plan().await?;
+ let state = self.session_state.read().clone();
+ plan_to_json(&state, plan, path).await
+ }
}
#[async_trait]
diff --git a/datafusion/src/execution/context.rs b/datafusion/src/execution/context.rs
index 7863eb4..fdf5a64 100644
--- a/datafusion/src/execution/context.rs
+++ b/datafusion/src/execution/context.rs
@@ -26,6 +26,7 @@ use crate::{
file_format::{
avro::{AvroFormat, DEFAULT_AVRO_EXTENSION},
csv::{CsvFormat, DEFAULT_CSV_EXTENSION},
+ json::{JsonFormat, DEFAULT_JSON_EXTENSION},
parquet::{ParquetFormat, DEFAULT_PARQUET_EXTENSION},
FileFormat,
},
@@ -75,7 +76,7 @@ use crate::physical_optimizer::repartition::Repartition;
use crate::execution::runtime_env::{RuntimeConfig, RuntimeEnv};
use crate::logical_plan::plan::Explain;
-use crate::physical_plan::file_format::{plan_to_csv, plan_to_parquet};
+use crate::physical_plan::file_format::{plan_to_csv, plan_to_json, plan_to_parquet};
use crate::physical_plan::planner::DefaultPhysicalPlanner;
use crate::physical_plan::udaf::AggregateUDF;
use crate::physical_plan::udf::ScalarUDF;
@@ -91,7 +92,7 @@ use chrono::{DateTime, Utc};
use parquet::file::properties::WriterProperties;
use uuid::Uuid;
-use super::options::{AvroReadOptions, CsvReadOptions};
+use super::options::{AvroReadOptions, CsvReadOptions, NdJsonReadOptions};
/// The default catalog name - this impacts what SQL queries use if not specified
const DEFAULT_CATALOG: &str = "datafusion";
@@ -211,24 +212,24 @@ impl SessionContext {
ref has_header,
}) => {
let (file_format, file_extension) = match file_type {
- FileType::CSV => Ok((
+ FileType::CSV => (
Arc::new(CsvFormat::default().with_has_header(*has_header))
as Arc<dyn FileFormat>,
DEFAULT_CSV_EXTENSION,
- )),
- FileType::Parquet => Ok((
+ ),
+ FileType::Parquet => (
Arc::new(ParquetFormat::default()) as Arc<dyn FileFormat>,
DEFAULT_PARQUET_EXTENSION,
- )),
- FileType::Avro => Ok((
+ ),
+ FileType::Avro => (
Arc::new(AvroFormat::default()) as Arc<dyn FileFormat>,
DEFAULT_AVRO_EXTENSION,
- )),
- _ => Err(DataFusionError::NotImplemented(format!(
- "Unsupported file type {:?}.",
- file_type
- ))),
- }?;
+ ),
+ FileType::NdJson => (
+ Arc::new(JsonFormat::default()) as Arc<dyn FileFormat>,
+ DEFAULT_JSON_EXTENSION,
+ ),
+ };
let options = ListingOptions {
format: file_format,
@@ -379,7 +380,6 @@ impl SessionContext {
}
/// Creates a DataFrame for reading an Avro data source.
-
pub async fn read_avro(
&mut self,
uri: impl Into<String>,
@@ -402,6 +402,29 @@ impl SessionContext {
)))
}
+ /// Creates a DataFrame for reading an Json data source.
+ pub async fn read_json(
+ &mut self,
+ uri: impl Into<String>,
+ options: NdJsonReadOptions<'_>,
+ ) -> Result<Arc<DataFrame>> {
+ let uri: String = uri.into();
+ let (object_store, path) = self.runtime_env().object_store(&uri)?;
+ let target_partitions = self.copied_config().target_partitions;
+ Ok(Arc::new(DataFrame::new(
+ self.state.clone(),
+ &LogicalPlanBuilder::scan_json(
+ object_store,
+ path,
+ options,
+ None,
+ target_partitions,
+ )
+ .await?
+ .build()?,
+ )))
+ }
+
/// Creates an empty DataFrame.
pub fn read_empty(&self) -> Result<Arc<DataFrame>> {
Ok(Arc::new(DataFrame::new(
@@ -508,6 +531,22 @@ impl SessionContext {
Ok(())
}
+ // Registers a Json data source so that it can be referenced from SQL statements
+ /// executed against this context.
+ pub async fn register_json(
+ &mut self,
+ name: &str,
+ uri: &str,
+ options: NdJsonReadOptions<'_>,
+ ) -> Result<()> {
+ let listing_options =
+ options.to_listing_options(self.copied_config().target_partitions);
+
+ self.register_listing_table(name, uri, listing_options, options.schema)
+ .await?;
+ Ok(())
+ }
+
/// Registers a Parquet data source so that it can be referenced from SQL statements
/// executed against this context.
pub async fn register_parquet(&mut self, name: &str, uri: &str) -> Result<()> {
@@ -709,6 +748,16 @@ impl SessionContext {
plan_to_csv(&state, plan, path).await
}
+ /// Executes a query and writes the results to a partitioned JSON file.
+ pub async fn write_json(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ path: impl AsRef<str>,
+ ) -> Result<()> {
+ let state = self.state.read().clone();
+ plan_to_json(&state, plan, path).await
+ }
+
/// Executes a query and writes the results to a partitioned Parquet file.
pub async fn write_parquet(
&self,
diff --git a/datafusion/src/execution/options.rs b/datafusion/src/execution/options.rs
index 79b0753..d9a0304 100644
--- a/datafusion/src/execution/options.rs
+++ b/datafusion/src/execution/options.rs
@@ -23,7 +23,7 @@ use arrow::datatypes::{Schema, SchemaRef};
use crate::datasource::file_format::json::DEFAULT_JSON_EXTENSION;
use crate::datasource::{
- file_format::{avro::AvroFormat, csv::CsvFormat},
+ file_format::{avro::AvroFormat, csv::CsvFormat, json::JsonFormat},
listing::ListingOptions,
};
@@ -178,3 +178,17 @@ impl<'a> Default for NdJsonReadOptions<'a> {
}
}
}
+
+impl<'a> NdJsonReadOptions<'a> {
+ /// Helper to convert these user facing options to `ListingTable` options
+ pub fn to_listing_options(&self, target_partitions: usize) -> ListingOptions {
+ let file_format = JsonFormat::default();
+ ListingOptions {
+ format: Arc::new(file_format),
+ collect_stat: false,
+ file_extension: self.file_extension.to_owned(),
+ target_partitions,
+ table_partition_cols: vec![],
+ }
+ }
+}
diff --git a/datafusion/src/logical_plan/builder.rs b/datafusion/src/logical_plan/builder.rs
index d0bfb5c..f2798b9 100644
--- a/datafusion/src/logical_plan/builder.rs
+++ b/datafusion/src/logical_plan/builder.rs
@@ -353,6 +353,55 @@ impl LogicalPlanBuilder {
Self::scan(table_name, Arc::new(provider), projection)
}
+ /// Scan an Json data source
+ pub async fn scan_json(
+ object_store: Arc<dyn ObjectStore>,
+ path: impl Into<String>,
+ options: NdJsonReadOptions<'_>,
+ projection: Option<Vec<usize>>,
+ target_partitions: usize,
+ ) -> Result<Self> {
+ let path = path.into();
+ Self::scan_json_with_name(
+ object_store,
+ path.clone(),
+ options,
+ projection,
+ path,
+ target_partitions,
+ )
+ .await
+ }
+
+ /// Scan an Json data source and register it with a given table name
+ pub async fn scan_json_with_name(
+ object_store: Arc<dyn ObjectStore>,
+ path: impl Into<String>,
+ options: NdJsonReadOptions<'_>,
+ projection: Option<Vec<usize>>,
+ table_name: impl Into<String>,
+ target_partitions: usize,
+ ) -> Result<Self> {
+ let listing_options = options.to_listing_options(target_partitions);
+
+ let path: String = path.into();
+
+ let resolved_schema = match options.schema {
+ Some(s) => s,
+ None => {
+ listing_options
+ .infer_schema(Arc::clone(&object_store), &path)
+ .await?
+ }
+ };
+ let config = ListingTableConfig::new(object_store, path)
+ .with_listing_options(listing_options)
+ .with_schema(resolved_schema);
+ let provider = ListingTable::try_new(config)?;
+
+ Self::scan(table_name, Arc::new(provider), projection)
+ }
+
/// Scan an empty data source, mainly used in tests
pub fn scan_empty(
name: Option<&str>,
diff --git a/datafusion/src/physical_plan/file_format/json.rs b/datafusion/src/physical_plan/file_format/json.rs
index 52a9e9f..14dfcf5 100644
--- a/datafusion/src/physical_plan/file_format/json.rs
+++ b/datafusion/src/physical_plan/file_format/json.rs
@@ -19,14 +19,19 @@
use async_trait::async_trait;
use crate::error::{DataFusionError, Result};
+use crate::execution::context::SessionState;
use crate::execution::context::TaskContext;
use crate::physical_plan::expressions::PhysicalSortExpr;
use crate::physical_plan::{
DisplayFormatType, ExecutionPlan, Partitioning, SendableRecordBatchStream, Statistics,
};
use arrow::{datatypes::SchemaRef, json};
+use futures::{StreamExt, TryStreamExt};
use std::any::Any;
+use std::fs;
+use std::path::Path;
use std::sync::Arc;
+use tokio::task::{self, JoinHandle};
use super::file_stream::{BatchIter, FileStream};
use super::FileScanConfig;
@@ -144,6 +149,44 @@ impl ExecutionPlan for NdJsonExec {
}
}
+pub async fn plan_to_json(
+ state: &SessionState,
+ plan: Arc<dyn ExecutionPlan>,
+ path: impl AsRef<str>,
+) -> Result<()> {
+ let path = path.as_ref();
+ // create directory to contain the CSV files (one per partition)
+ let fs_path = Path::new(path);
+ match fs::create_dir(fs_path) {
+ Ok(()) => {
+ let mut tasks = vec![];
+ for i in 0..plan.output_partitioning().partition_count() {
+ let plan = plan.clone();
+ let filename = format!("part-{}.json", i);
+ let path = fs_path.join(&filename);
+ let file = fs::File::create(path)?;
+ let mut writer = json::LineDelimitedWriter::new(file);
+ let task_ctx = Arc::new(TaskContext::from(state));
+ let stream = plan.execute(i, task_ctx).await?;
+ let handle: JoinHandle<Result<()>> = task::spawn(async move {
+ stream
+ .map(|batch| writer.write(batch?))
+ .try_collect()
+ .await
+ .map_err(DataFusionError::from)
+ });
+ tasks.push(handle);
+ }
+ futures::future::join_all(tasks).await;
+ Ok(())
+ }
+ Err(e) => Err(DataFusionError::Execution(format!(
+ "Could not create directory {}: {:?}",
+ path, e
+ ))),
+ }
+}
+
#[cfg(test)]
mod tests {
use arrow::array::Array;
@@ -156,7 +199,9 @@ mod tests {
local_object_reader_stream, local_unpartitioned_file, LocalFileSystem,
},
};
- use crate::prelude::SessionContext;
+ use crate::prelude::NdJsonReadOptions;
+ use crate::prelude::*;
+ use tempfile::TempDir;
use super::*;
@@ -302,4 +347,52 @@ mod tests {
assert_eq!(values.value(2), 2);
Ok(())
}
+
+ #[tokio::test]
+ async fn write_json_results() -> Result<()> {
+ // create partitioned input file and context
+ let tmp_dir = TempDir::new()?;
+ let mut ctx =
+ SessionContext::with_config(SessionConfig::new().with_target_partitions(8));
+
+ let path = format!("{}/1.json", TEST_DATA_BASE);
+
+ // register json file with the execution context
+ ctx.register_json("test", path.as_str(), NdJsonReadOptions::default())
+ .await?;
+
+ // execute a simple query and write the results to CSV
+ let out_dir = tmp_dir.as_ref().to_str().unwrap().to_string() + "/out";
+ let df = ctx.sql("SELECT a, b FROM test").await?;
+ df.write_json(&out_dir).await?;
+
+ // create a new context and verify that the results were saved to a partitioned csv file
+ let mut ctx = SessionContext::new();
+
+ // register each partition as well as the top level dir
+ let json_read_option = NdJsonReadOptions::default();
+ ctx.register_json(
+ "part0",
+ &format!("{}/part-0.json", out_dir),
+ json_read_option.clone(),
+ )
+ .await?;
+ ctx.register_json("allparts", &out_dir, json_read_option)
+ .await?;
+
+ let part0 = ctx.sql("SELECT a, b FROM part0").await?.collect().await?;
+ let allparts = ctx
+ .sql("SELECT a, b FROM allparts")
+ .await?
+ .collect()
+ .await?;
+
+ let allparts_count: usize = allparts.iter().map(|batch| batch.num_rows()).sum();
+
+ assert_eq!(part0[0].schema(), allparts[0].schema());
+
+ assert_eq!(allparts_count, 4);
+
+ Ok(())
+ }
}
diff --git a/datafusion/src/physical_plan/file_format/mod.rs b/datafusion/src/physical_plan/file_format/mod.rs
index 0e1e859..1e359ab 100644
--- a/datafusion/src/physical_plan/file_format/mod.rs
+++ b/datafusion/src/physical_plan/file_format/mod.rs
@@ -35,6 +35,7 @@ use arrow::{
pub use avro::AvroExec;
pub(crate) use csv::plan_to_csv;
pub use csv::CsvExec;
+pub(crate) use json::plan_to_json;
pub use json::NdJsonExec;
use crate::error::DataFusionError;
diff --git a/datafusion/tests/jsons/3.json b/datafusion/tests/jsons/3.json
new file mode 100644
index 0000000..f4ea5c6
--- /dev/null
+++ b/datafusion/tests/jsons/3.json
@@ -0,0 +1 @@
+{"mycol": null}
\ No newline at end of file
diff --git a/datafusion/tests/sql/json.rs b/datafusion/tests/sql/json.rs
new file mode 100644
index 0000000..c87bcfd
--- /dev/null
+++ b/datafusion/tests/sql/json.rs
@@ -0,0 +1,103 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use super::*;
+
+const TEST_DATA_BASE: &str = "tests/jsons";
+
+#[tokio::test]
+async fn json_query() {
+ let mut ctx = SessionContext::new();
+ let path = format!("{}/2.json", TEST_DATA_BASE);
+ ctx.register_json("t1", &path, NdJsonReadOptions::default())
+ .await
+ .unwrap();
+
+ let sql = "SELECT a, b FROM t1";
+ let actual = execute_to_batches(&ctx, sql).await;
+ let expected = vec![
+ "+-----------------+------+",
+ "| a | b |",
+ "+-----------------+------+",
+ "| 1 | 2 |",
+ "| -10 | -3.5 |",
+ "| 2 | 0.6 |",
+ "| 1 | 2 |",
+ "| 7 | -3.5 |",
+ "| 1 | 0.6 |",
+ "| 1 | 2 |",
+ "| 5 | -3.5 |",
+ "| 1 | 0.6 |",
+ "| 1 | 2 |",
+ "| 1 | -3.5 |",
+ "| 100000000000000 | 0.6 |",
+ "+-----------------+------+",
+ ];
+
+ assert_batches_eq!(expected, &actual);
+}
+
+#[tokio::test]
+#[should_panic]
+async fn json_single_nan_schema() {
+ let mut ctx = SessionContext::new();
+ let path = format!("{}/3.json", TEST_DATA_BASE);
+ ctx.register_json("single_nan", &path, NdJsonReadOptions::default())
+ .await
+ .unwrap();
+ let sql = "SELECT mycol FROM single_nan";
+ let plan = ctx.create_logical_plan(sql).unwrap();
+ let plan = ctx.optimize(&plan).unwrap();
+ let plan = ctx.create_physical_plan(&plan).await.unwrap();
+ let task_ctx = ctx.task_ctx();
+ let results = collect(plan, task_ctx).await.unwrap();
+ for batch in results {
+ assert_eq!(1, batch.num_rows());
+ assert_eq!(1, batch.num_columns());
+ }
+}
+
+#[tokio::test]
+async fn json_explain() {
+ let mut ctx = SessionContext::new();
+ let path = format!("{}/2.json", TEST_DATA_BASE);
+ ctx.register_json("t1", &path, NdJsonReadOptions::default())
+ .await
+ .unwrap();
+
+ let sql = "EXPLAIN SELECT count(*) from t1";
+ let actual = execute(&mut ctx, sql).await;
+ let actual = normalize_vec_for_explain(actual);
+ let expected = vec![
+ vec![
+ "logical_plan",
+ "Projection: #COUNT(UInt8(1))\
+ \n Aggregate: groupBy=[[]], aggr=[[COUNT(UInt8(1))]]\
+ \n TableScan: t1 projection=Some([0])",
+ ],
+ vec![
+ "physical_plan",
+ "ProjectionExec: expr=[COUNT(UInt8(1))@0 as COUNT(UInt8(1))]\
+ \n HashAggregateExec: mode=Final, gby=[], aggr=[COUNT(UInt8(1))]\
+ \n CoalescePartitionsExec\
+ \n HashAggregateExec: mode=Partial, gby=[], aggr=[COUNT(UInt8(1))]\
+ \n RepartitionExec: partitioning=RoundRobinBatch(NUM_CORES)\
+ \n JsonExec: limit=None, files=[tests/jsons/2.json]\n",
+ ],
+ ];
+ assert_eq!(expected, actual);
+}
diff --git a/datafusion/tests/sql/mod.rs b/datafusion/tests/sql/mod.rs
index 05c5abd..40fe7e7 100644
--- a/datafusion/tests/sql/mod.rs
+++ b/datafusion/tests/sql/mod.rs
@@ -84,6 +84,7 @@ pub mod functions;
pub mod group_by;
pub mod intersection;
pub mod joins;
+pub mod json;
pub mod limit;
pub mod order;
pub mod parquet;