Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(source): add ndjson source for streaming load #4561

Merged
merged 4 commits into from
Mar 24, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

8 changes: 8 additions & 0 deletions common/datavalues/src/types/deserializations/boolean.rs
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,14 @@ impl TypeDeserializer for BooleanDeserializer {
Ok(())
}

fn de_json(&mut self, value: &serde_json::Value) -> Result<()> {
match value {
serde_json::Value::Bool(v) => self.builder.append_value(*v),
_ => return Err(ErrorCode::BadBytes("Incorrect boolean value")),
}
Ok(())
}

fn finish_to_column(&mut self) -> ColumnRef {
self.builder.to_column()
}
Expand Down
7 changes: 7 additions & 0 deletions common/datavalues/src/types/deserializations/date.rs
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,13 @@ where
Ok(())
}

fn de_json(&mut self, value: &serde_json::Value) -> Result<()> {
match value {
serde_json::Value::String(v) => self.de_text(v.as_bytes()),
_ => Err(ErrorCode::BadBytes("Incorrect boolean value")),
}
}

fn de_text(&mut self, reader: &[u8]) -> Result<()> {
match lexical_core::parse::<T>(reader) {
Ok(v) => {
Expand Down
7 changes: 7 additions & 0 deletions common/datavalues/src/types/deserializations/date_time.rs
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,13 @@ where
Ok(())
}

fn de_json(&mut self, value: &serde_json::Value) -> Result<()> {
match value {
serde_json::Value::String(v) => self.de_text(v.as_bytes()),
_ => Err(ErrorCode::BadBytes("Incorrect boolean value")),
}
}

fn de_text(&mut self, reader: &[u8]) -> Result<()> {
let v = std::str::from_utf8(reader)
.map_err_to_code(ErrorCode::BadBytes, || "Cannot convert value to utf8")?;
Expand Down
3 changes: 3 additions & 0 deletions common/datavalues/src/types/deserializations/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
// limitations under the License.

use common_exception::Result;
use serde_json::Value;

use crate::prelude::*;

Expand Down Expand Up @@ -42,6 +43,8 @@ pub trait TypeDeserializer: Send + Sync {
/// If error occurrs, append a null by default
fn de_text(&mut self, reader: &[u8]) -> Result<()>;

fn de_json(&mut self, reader: &Value) -> Result<()>;

fn de_null(&mut self) -> bool {
false
}
Expand Down
5 changes: 5 additions & 0 deletions common/datavalues/src/types/deserializations/null.rs
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,11 @@ impl TypeDeserializer for NullDeserializer {
Ok(())
}

fn de_json(&mut self, _value: &serde_json::Value) -> Result<()> {
self.builder.append_default();
Ok(())
}

fn de_text(&mut self, _reader: &[u8]) -> Result<()> {
self.builder.append_default();
Ok(())
Expand Down
13 changes: 13 additions & 0 deletions common/datavalues/src/types/deserializations/nullable.rs
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,19 @@ impl TypeDeserializer for NullableDeserializer {
unreachable!()
}

fn de_json(&mut self, value: &serde_json::Value) -> Result<()> {
match value {
serde_json::Value::Null => {
self.de_null();
Ok(())
}
other => {
self.bitmap.push(true);
self.inner.de_json(other)
}
}
}

// TODO: support null text setting
fn de_text(&mut self, reader: &[u8]) -> Result<()> {
self.inner.de_text(reader)?;
Expand Down
11 changes: 11 additions & 0 deletions common/datavalues/src/types/deserializations/number.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.

use common_exception::ErrorCode;
use common_exception::Result;
use common_io::prelude::*;
use lexical_core::FromLexical;
Expand Down Expand Up @@ -46,6 +47,16 @@ where
Ok(())
}

fn de_json(&mut self, value: &serde_json::Value) -> Result<()> {
match value {
serde_json::Value::Number(v) => {
let str = v.to_string();
self.de_text(str.as_bytes())
}
_ => Err(ErrorCode::BadBytes("Incorrect json value, must be number")),
}
}

fn de_text(&mut self, reader: &[u8]) -> Result<()> {
let value = lexical_core::parse_partial::<T>(reader)
.unwrap_or((T::default(), 0))
Expand Down
11 changes: 11 additions & 0 deletions common/datavalues/src/types/deserializations/string.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@

use std::io::Read;

use common_exception::ErrorCode;
use common_exception::Result;
use common_io::prelude::BinaryRead;

Expand Down Expand Up @@ -62,6 +63,16 @@ impl TypeDeserializer for StringDeserializer {
Ok(())
}

fn de_json(&mut self, value: &serde_json::Value) -> Result<()> {
match value {
serde_json::Value::String(s) => {
self.builder.append_value(s);
Ok(())
}
_ => Err(ErrorCode::BadBytes("Incorrect json value, must be string")),
}
}

fn de_text(&mut self, reader: &[u8]) -> Result<()> {
self.builder.append_value(reader);
Ok(())
Expand Down
5 changes: 5 additions & 0 deletions common/datavalues/src/types/deserializations/variant.rs
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,11 @@ impl TypeDeserializer for VariantDeserializer {
Ok(())
}

fn de_json(&mut self, value: &serde_json::Value) -> Result<()> {
self.builder.append_value(value.clone());
Ok(())
}

fn de_text(&mut self, reader: &[u8]) -> Result<()> {
let val = serde_json::from_slice(reader)?;
self.builder.append_value(val);
Expand Down
1 change: 1 addition & 0 deletions common/streams/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ async-trait = "0.1.52"
csv-async = "1.2.4"
futures = "0.3.21"
pin-project-lite = "0.2.8"
serde_json = { version = "1.0.79", default-features = false, features = ["preserve_order"] }
tempfile = "3.3.0"

[dev-dependencies]
Expand Down
2 changes: 2 additions & 0 deletions common/streams/src/sources/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,10 @@

mod source;
mod source_csv;
mod source_ndjson;
mod source_parquet;

pub use source::Source;
pub use source_csv::CsvSourceBuilder;
pub use source_ndjson::NDJsonSourceBuilder;
pub use source_parquet::ParquetSourceBuilder;
6 changes: 3 additions & 3 deletions common/streams/src/sources/source_csv.rs
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ impl CsvSourceBuilder {
field_delimiter: b',',
record_delimiter: Terminator::CRLF,
block_size: 10000,
size_limit: 0,
size_limit: usize::MAX,
}
}

Expand Down Expand Up @@ -128,7 +128,7 @@ where R: AsyncRead + Unpin + Send
{
async fn read(&mut self) -> Result<Option<DataBlock>> {
// Check size_limit.
if self.builder.size_limit > 0 && self.rows >= self.builder.size_limit {
if self.rows >= self.builder.size_limit {
return Ok(None);
}

Expand Down Expand Up @@ -161,7 +161,7 @@ where R: AsyncRead + Unpin + Send
self.rows += 1;

// Check size_limit.
if self.builder.size_limit > 0 && self.rows >= self.builder.size_limit {
if self.rows >= self.builder.size_limit {
break;
}

Expand Down
156 changes: 156 additions & 0 deletions common/streams/src/sources/source_ndjson.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,156 @@
// Copyright 2021 Datafuse Labs.
//
// Licensed 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 std::io::BufRead;

use async_trait::async_trait;
use common_datablocks::DataBlock;
use common_datavalues::DataSchemaRef;
use common_exception::ErrorCode;
use common_exception::Result;
use common_exception::ToErrorCode;

use crate::Source;

#[derive(Debug, Clone)]
pub struct NDJsonSourceBuilder {
schema: DataSchemaRef,
block_size: usize,
size_limit: usize,
}

impl NDJsonSourceBuilder {
pub fn create(schema: DataSchemaRef) -> Self {
NDJsonSourceBuilder {
schema,
block_size: 10000,
size_limit: usize::MAX,
}
}

pub fn block_size(&mut self, block_size: usize) -> &mut Self {
self.block_size = block_size;
self
}

pub fn size_limit(&mut self, size_limit: usize) -> &mut Self {
self.size_limit = size_limit;
self
}

pub fn build<R>(&self, reader: R) -> Result<NDJsonSource<R>>
where R: BufRead + Unpin + Send {
NDJsonSource::try_create(self.clone(), reader)
}
}

pub struct NDJsonSource<R> {
builder: NDJsonSourceBuilder,
reader: R,
rows: usize,
buffer: String,
}

impl<R> NDJsonSource<R>
where R: BufRead + Unpin + Send
{
fn try_create(builder: NDJsonSourceBuilder, reader: R) -> Result<Self> {
Ok(Self {
builder,
reader,
rows: 0,
buffer: String::new(),
})
}
}

#[async_trait]
impl<R> Source for NDJsonSource<R>
where R: BufRead + Unpin + Send
{
async fn read(&mut self) -> Result<Option<DataBlock>> {
// Check size_limit.
if self.rows >= self.builder.size_limit {
return Ok(None);
}

let mut packs = self
.builder
.schema
.fields()
.iter()
.map(|f| f.data_type().create_deserializer(self.builder.block_size))
.collect::<Vec<_>>();

let names = self
.builder
.schema
.fields()
.iter()
.map(|f| f.name())
.collect::<Vec<_>>();

let mut rows = 0;

loop {
self.buffer.clear();
if self
.reader
.read_line(&mut self.buffer)
Copy link
Contributor

@DCjanus DCjanus Mar 26, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here is a block IO in an async function, might block runtime thread.

In this PR, it would not cause serious problem, because we are using std::io::Curse, but in the future, peopel may trying to use this in other situnations, and shot their foot.

Maybe we should replace R: std::io::BufRead with R: tokio::io::BufRead or R: futures::io::AsyncBufRead

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, good catch, I will create a fix for that.

.map_err_to_code(ErrorCode::BadBytes, || {
format!("Parse NDJson error at line {}", self.rows)
})?
== 0
{
break;
}

if self.buffer.trim().is_empty() {
continue;
}

let json: serde_json::Value = serde_json::from_reader(self.buffer.as_bytes())?;
// Deserialize each field.

for (name, deser) in names.iter().zip(packs.iter_mut()) {
let value = &json[name];
deser.de_json(value)?;
}

rows += 1;
self.rows += 1;

// Check size_limit.
if self.rows >= self.builder.size_limit {
break;
}

// Check block_size.
if rows >= self.builder.block_size {
break;
}
}

if rows == 0 {
return Ok(None);
}

let series = packs
.iter_mut()
.map(|deser| deser.finish_to_column())
.collect::<Vec<_>>();

Ok(Some(DataBlock::create(self.builder.schema.clone(), series)))
}
}
Loading