This is an automated email from the ASF dual-hosted git repository.
agrove pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/master by this push:
new ee09cb6 ARROW-8839: [Rust] [DataFusion] support CSV schema inference
in logical plan
ee09cb6 is described below
commit ee09cb6edce31899841afa0c01fb4ef1dd971ef9
Author: Qingping Hou <[email protected]>
AuthorDate: Tue May 19 18:47:52 2020 -0600
ARROW-8839: [Rust] [DataFusion] support CSV schema inference in logical plan
This PR changes schema argument for scan_csv method into `Option<&Schema>`.
Other related changes are needed to make this happen including:
* added delimiter argument to all csv related structs and functions
* fixed a bug in schema field inference function
* made `arrow::csv::reader::infer_file_schema` public so it can be used by
data fusion
Known limitations:
* when provided with a directory of csv files, schema inference code only
reads rows from the first file.
* to avoid adding yet another argument to all csv related functions, i hard
coded number of rows to read for schema inference to 1000
Open questions:
* Should we rename `datasource::csv::CsvFile` struct to `CsvTable` to keep
it consistent with ParquetTable and MemoryTable? The implementation of CsvFile
also supports reading from a directory of files, so `CsvFile` is not an
accurate name.
* csv related function arguments are getting a bit long, should we
introduce a csv option struct to capture the following configs with sensible
defaults?
- schema
- has_header
- delimiter
- infer_max_read_records
Closes #7210 from houqp/csv_schema_infer
Authored-by: Qingping Hou <[email protected]>
Signed-off-by: Andy Grove <[email protected]>
---
rust/arrow/examples/read_csv.rs | 2 +-
rust/arrow/examples/read_csv_infer_schema.rs | 2 +-
rust/arrow/src/csv/mod.rs | 1 +
rust/arrow/src/csv/reader.rs | 88 ++++++++++++++--------
rust/datafusion/benches/aggregate_query_sql.rs | 1 +
rust/datafusion/examples/csv_sql.rs | 1 +
rust/datafusion/src/datasource/csv.rs | 44 ++++++++++-
rust/datafusion/src/execution/context.rs | 54 ++++++++++---
rust/datafusion/src/execution/physical_plan/csv.rs | 35 ++++++++-
.../src/execution/physical_plan/hash_aggregate.rs | 2 +-
.../src/execution/physical_plan/limit.rs | 2 +-
.../src/execution/physical_plan/merge.rs | 2 +-
.../src/execution/physical_plan/projection.rs | 2 +-
.../src/execution/physical_plan/selection.rs | 2 +-
rust/datafusion/src/execution/table_impl.rs | 1 +
rust/datafusion/src/logicalplan.rs | 20 ++++-
.../src/optimizer/projection_push_down.rs | 2 +
rust/datafusion/tests/sql.rs | 4 +-
18 files changed, 209 insertions(+), 56 deletions(-)
diff --git a/rust/arrow/examples/read_csv.rs b/rust/arrow/examples/read_csv.rs
index cde59d7..fb14b4e 100644
--- a/rust/arrow/examples/read_csv.rs
+++ b/rust/arrow/examples/read_csv.rs
@@ -34,7 +34,7 @@ fn main() -> Result<()> {
let file = File::open("test/data/uk_cities.csv").unwrap();
- let mut csv = csv::Reader::new(file, Arc::new(schema), false, 1024, None);
+ let mut csv = csv::Reader::new(file, Arc::new(schema), false, None, 1024,
None);
let batch = csv.next().unwrap().unwrap();
print_batches(&vec![batch])
}
diff --git a/rust/arrow/examples/read_csv_infer_schema.rs
b/rust/arrow/examples/read_csv_infer_schema.rs
index 07c28c7..eff9dc7 100644
--- a/rust/arrow/examples/read_csv_infer_schema.rs
+++ b/rust/arrow/examples/read_csv_infer_schema.rs
@@ -25,7 +25,7 @@ use std::fs::File;
fn main() -> Result<()> {
let file = File::open("test/data/uk_cities_with_headers.csv").unwrap();
let builder = csv::ReaderBuilder::new()
- .has_headers(true)
+ .has_header(true)
.infer_schema(Some(100));
let mut csv = builder.build(file).unwrap();
let batch = csv.next().unwrap().unwrap();
diff --git a/rust/arrow/src/csv/mod.rs b/rust/arrow/src/csv/mod.rs
index 55de4b2..b17b787 100644
--- a/rust/arrow/src/csv/mod.rs
+++ b/rust/arrow/src/csv/mod.rs
@@ -20,6 +20,7 @@
pub mod reader;
pub mod writer;
+pub use self::reader::infer_file_schema;
pub use self::reader::Reader;
pub use self::reader::ReaderBuilder;
pub use self::writer::Writer;
diff --git a/rust/arrow/src/csv/reader.rs b/rust/arrow/src/csv/reader.rs
index 6b99710..0d15570 100644
--- a/rust/arrow/src/csv/reader.rs
+++ b/rust/arrow/src/csv/reader.rs
@@ -36,7 +36,7 @@
//!
//! let file = File::open("test/data/uk_cities.csv").unwrap();
//!
-//! let mut csv = csv::Reader::new(file, Arc::new(schema), false, 1024, None);
+//! let mut csv = csv::Reader::new(file, Arc::new(schema), false, None, 1024,
None);
//! let batch = csv.next().unwrap().unwrap();
//! ```
@@ -57,7 +57,7 @@ use self::csv_crate::{StringRecord, StringRecordsIntoIter};
lazy_static! {
static ref DECIMAL_RE: Regex = Regex::new(r"^-?(\d+\.\d+)$").unwrap();
- static ref INTEGER_RE: Regex = Regex::new(r"^-?(\d*.)$").unwrap();
+ static ref INTEGER_RE: Regex = Regex::new(r"^-?(\d+)$").unwrap();
static ref BOOLEAN_RE: Regex = RegexBuilder::new(r"^(true)$|^(false)$")
.case_insensitive(true)
.build()
@@ -87,19 +87,19 @@ fn infer_field_schema(string: &str) -> DataType {
/// with `max_read_records` controlling the maximum number of records to read.
///
/// If `max_read_records` is not set, the whole file is read to infer its
schema.
-fn infer_file_schema<R: Read + Seek>(
+pub fn infer_file_schema<R: Read + Seek>(
reader: &mut BufReader<R>,
delimiter: u8,
max_read_records: Option<usize>,
- has_headers: bool,
+ has_header: bool,
) -> Result<Schema> {
let mut csv_reader = csv::ReaderBuilder::new()
.delimiter(delimiter)
.from_reader(reader);
// get or create header names
- // when has_headers is false, creates default column names with column_
prefix
- let headers: Vec<String> = if has_headers {
+ // when has_header is false, creates default column names with column_
prefix
+ let headers: Vec<String> = if has_header {
let headers = &csv_reader.headers()?.clone();
headers.iter().map(|s| s.to_string()).collect()
} else {
@@ -198,14 +198,16 @@ impl<R: Read> Reader<R> {
pub fn new(
reader: R,
schema: Arc<Schema>,
- has_headers: bool,
+ has_header: bool,
+ delimiter: Option<u8>,
batch_size: usize,
projection: Option<Vec<usize>>,
) -> Self {
Self::from_buf_reader(
BufReader::new(reader),
schema,
- has_headers,
+ has_header,
+ delimiter,
batch_size,
projection,
)
@@ -233,20 +235,29 @@ impl<R: Read> Reader<R> {
pub fn from_buf_reader(
buf_reader: BufReader<R>,
schema: Arc<Schema>,
- has_headers: bool,
+ has_header: bool,
+ delimiter: Option<u8>,
batch_size: usize,
projection: Option<Vec<usize>>,
) -> Self {
- let csv_reader = csv::ReaderBuilder::new()
- .has_headers(has_headers)
- .from_reader(buf_reader);
+ let mut reader_builder = csv_crate::ReaderBuilder::new();
+ reader_builder.has_headers(has_header);
+
+ match delimiter {
+ Some(c) => {
+ reader_builder.delimiter(c);
+ }
+ _ => (),
+ }
+
+ let csv_reader = reader_builder.from_reader(buf_reader);
let record_iter = csv_reader.into_records();
Self {
schema,
projection,
record_iter,
batch_size,
- line_number: if has_headers { 1 } else { 0 },
+ line_number: if has_header { 1 } else { 0 },
}
}
@@ -369,8 +380,9 @@ impl<R: Read> Reader<R> {
Err(_) => {
// TODO: we should surface the underlying error
here.
return Err(ArrowError::ParseError(format!(
- "Error while parsing value {} at line {}",
+ "Error while parsing value {} for column {} at
line {}",
s,
+ col_idx,
self.line_number + row_index
)));
}
@@ -394,7 +406,7 @@ pub struct ReaderBuilder {
///
/// If schema inference is run on a file with no headers, default column
names
/// are created.
- has_headers: bool,
+ has_header: bool,
/// An optional column delimiter. Defaults to `b','`
delimiter: Option<u8>,
/// Optional maximum number of records to read during schema inference
@@ -413,7 +425,7 @@ impl Default for ReaderBuilder {
fn default() -> ReaderBuilder {
ReaderBuilder {
schema: None,
- has_headers: false,
+ has_header: false,
delimiter: None,
max_records: None,
batch_size: 1024,
@@ -457,8 +469,8 @@ impl ReaderBuilder {
}
/// Set whether the CSV file has headers
- pub fn has_headers(mut self, has_headers: bool) -> Self {
- self.has_headers = has_headers;
+ pub fn has_header(mut self, has_header: bool) -> Self {
+ self.has_header = has_header;
self
}
@@ -492,22 +504,23 @@ impl ReaderBuilder {
pub fn build<R: Read + Seek>(self, reader: R) -> Result<Reader<R>> {
// check if schema should be inferred
let mut buf_reader = BufReader::new(reader);
+ let delimiter = self.delimiter.unwrap_or(b',');
let schema = match self.schema {
Some(schema) => schema,
None => {
let inferred_schema = infer_file_schema(
&mut buf_reader,
- self.delimiter.unwrap_or(b','),
+ delimiter,
self.max_records,
- self.has_headers,
+ self.has_header,
)?;
Arc::new(inferred_schema)
}
};
- let csv_reader = csv::ReaderBuilder::new()
- .delimiter(self.delimiter.unwrap_or(b','))
- .has_headers(self.has_headers)
+ let csv_reader = csv_crate::ReaderBuilder::new()
+ .delimiter(delimiter)
+ .has_headers(self.has_header)
.from_reader(buf_reader);
let record_iter = csv_reader.into_records();
Ok(Reader {
@@ -515,7 +528,7 @@ impl ReaderBuilder {
projection: self.projection.clone(),
record_iter,
batch_size: self.batch_size,
- line_number: if self.has_headers { 1 } else { 0 },
+ line_number: if self.has_header { 1 } else { 0 },
})
}
}
@@ -540,7 +553,8 @@ mod tests {
let file = File::open("test/data/uk_cities.csv").unwrap();
- let mut csv = Reader::new(file, Arc::new(schema.clone()), false, 1024,
None);
+ let mut csv =
+ Reader::new(file, Arc::new(schema.clone()), false, None, 1024,
None);
assert_eq!(Arc::new(schema), csv.schema());
let batch = csv.next().unwrap().unwrap();
assert_eq!(37, batch.num_rows());
@@ -582,6 +596,7 @@ mod tests {
BufReader::new(both_files),
Arc::new(schema),
true,
+ None,
1024,
None,
);
@@ -594,7 +609,7 @@ mod tests {
fn test_csv_with_schema_inference() {
let file = File::open("test/data/uk_cities_with_headers.csv").unwrap();
- let builder =
ReaderBuilder::new().has_headers(true).infer_schema(None);
+ let builder = ReaderBuilder::new().has_header(true).infer_schema(None);
let mut csv = builder.build(file).unwrap();
let expected_schema = Schema::new(vec![
@@ -673,7 +688,8 @@ mod tests {
let file = File::open("test/data/uk_cities.csv").unwrap();
- let mut csv = Reader::new(file, Arc::new(schema), false, 1024,
Some(vec![0, 1]));
+ let mut csv =
+ Reader::new(file, Arc::new(schema), false, None, 1024,
Some(vec![0, 1]));
let projected_schema = Arc::new(Schema::new(vec![
Field::new("city", DataType::Utf8, false),
Field::new("lat", DataType::Float64, false),
@@ -695,7 +711,7 @@ mod tests {
let file = File::open("test/data/null_test.csv").unwrap();
- let mut csv = Reader::new(file, Arc::new(schema), true, 1024, None);
+ let mut csv = Reader::new(file, Arc::new(schema), true, None, 1024,
None);
let batch = csv.next().unwrap().unwrap();
assert_eq!(false, batch.column(1).is_null(0));
@@ -711,7 +727,7 @@ mod tests {
let builder = ReaderBuilder::new()
.infer_schema(None)
- .has_headers(true)
+ .has_header(true)
.with_delimiter(b'|')
.with_batch_size(512)
.with_projection(vec![0, 1, 2, 3]);
@@ -754,7 +770,7 @@ mod tests {
let builder = ReaderBuilder::new()
.with_schema(Arc::new(schema))
- .has_headers(true)
+ .has_header(true)
.with_delimiter(b'|')
.with_batch_size(512)
.with_projection(vec![0, 1, 2, 3]);
@@ -762,10 +778,20 @@ mod tests {
let mut csv = builder.build(file).unwrap();
match csv.next() {
Err(e) => assert_eq!(
- "ParseError(\"Error while parsing value 4.x4 at line 4\")",
+ "ParseError(\"Error while parsing value 4.x4 for column 1 at
line 4\")",
format!("{:?}", e)
),
Ok(_) => panic!("should have failed"),
}
}
+
+ #[test]
+ fn test_infer_field_schema() {
+ assert_eq!(infer_field_schema("A"), DataType::Utf8);
+ assert_eq!(infer_field_schema("\"123\""), DataType::Utf8);
+ assert_eq!(infer_field_schema("10"), DataType::Int64);
+ assert_eq!(infer_field_schema("10.2"), DataType::Float64);
+ assert_eq!(infer_field_schema("true"), DataType::Boolean);
+ assert_eq!(infer_field_schema("false"), DataType::Boolean);
+ }
}
diff --git a/rust/datafusion/benches/aggregate_query_sql.rs
b/rust/datafusion/benches/aggregate_query_sql.rs
index e3b3d56..88b534b 100644
--- a/rust/datafusion/benches/aggregate_query_sql.rs
+++ b/rust/datafusion/benches/aggregate_query_sql.rs
@@ -63,6 +63,7 @@ fn create_context() -> ExecutionContext {
&format!("{}/csv/aggregate_test_100.csv", testdata),
&schema,
true,
+ None,
);
let mem_table = MemTable::load(&csv).unwrap();
diff --git a/rust/datafusion/examples/csv_sql.rs
b/rust/datafusion/examples/csv_sql.rs
index 17106d2..dc538bf 100644
--- a/rust/datafusion/examples/csv_sql.rs
+++ b/rust/datafusion/examples/csv_sql.rs
@@ -52,6 +52,7 @@ fn main() -> Result<()> {
&format!("{}/csv/aggregate_test_100.csv", testdata),
&schema,
true,
+ None,
);
let sql = "SELECT c1, MIN(c12), MAX(c12) FROM aggregate_test_100 WHERE c11
> 0.1 AND c11 < 0.9 GROUP BY c1";
diff --git a/rust/datafusion/src/datasource/csv.rs
b/rust/datafusion/src/datasource/csv.rs
index fb3014e..a6a957d 100644
--- a/rust/datafusion/src/datasource/csv.rs
+++ b/rust/datafusion/src/datasource/csv.rs
@@ -36,17 +36,54 @@ pub struct CsvFile {
filename: String,
schema: Arc<Schema>,
has_header: bool,
+ delimiter: Option<u8>,
}
impl CsvFile {
#[allow(missing_docs)]
- pub fn new(filename: &str, schema: &Schema, has_header: bool) -> Self {
+ pub fn new(
+ filename: &str,
+ schema: &Schema,
+ has_header: bool,
+ delimiter: Option<u8>,
+ ) -> Self {
Self {
filename: String::from(filename),
schema: Arc::new(schema.clone()),
has_header,
+ delimiter,
}
}
+
+ /// Attempt to initialize a new `CsvFile` from a file path
+ pub fn try_new(
+ filename: &str,
+ schema: Option<&Schema>,
+ has_header: bool,
+ delimiter: Option<u8>,
+ ) -> Result<Self> {
+ let schema = match schema {
+ Some(s) => Arc::new(s.clone()),
+ None => {
+ let schema_infer_batch_size = 1024;
+ let csv_exec = CsvExec::try_new(
+ filename,
+ None,
+ has_header,
+ delimiter,
+ None,
+ schema_infer_batch_size,
+ )?;
+ csv_exec.schema()
+ }
+ };
+ Ok(Self {
+ filename: String::from(filename),
+ schema: schema,
+ has_header,
+ delimiter,
+ })
+ }
}
impl TableProvider for CsvFile {
@@ -61,8 +98,9 @@ impl TableProvider for CsvFile {
) -> Result<Vec<ScanResult>> {
let exec = CsvExec::try_new(
&self.filename,
- self.schema.clone(),
+ Some(self.schema.clone()),
self.has_header,
+ self.delimiter,
projection.clone(),
batch_size,
)?;
@@ -88,6 +126,7 @@ impl CsvBatchIterator {
filename: &str,
schema: Arc<Schema>,
has_header: bool,
+ delimiter: Option<u8>,
projection: &Option<Vec<usize>>,
batch_size: usize,
) -> Result<Self> {
@@ -96,6 +135,7 @@ impl CsvBatchIterator {
file,
schema.clone(),
has_header,
+ delimiter,
batch_size,
projection.clone(),
);
diff --git a/rust/datafusion/src/execution/context.rs
b/rust/datafusion/src/execution/context.rs
index 9b8d068..d804ae4 100644
--- a/rust/datafusion/src/execution/context.rs
+++ b/rust/datafusion/src/execution/context.rs
@@ -100,7 +100,7 @@ impl ExecutionContext {
ref header_row,
} => match file_type {
FileType::CSV => {
- self.register_csv(name, location, schema, *header_row);
+ self.register_csv(name, location, schema, *header_row,
None);
Ok(vec![])
}
FileType::Parquet => {
@@ -216,8 +216,12 @@ impl ExecutionContext {
filename: &str,
schema: &Schema,
has_header: bool,
+ delimiter: Option<u8>,
) {
- self.register_table(name, Box::new(CsvFile::new(filename, schema,
has_header)));
+ self.register_table(
+ name,
+ Box::new(CsvFile::new(filename, schema, has_header, delimiter)),
+ );
}
/// Register a Parquet file as a table so that it can be queried from SQL
@@ -314,12 +318,14 @@ impl ExecutionContext {
path,
schema,
has_header,
+ delimiter,
projection,
..
} => Ok(Arc::new(CsvExec::try_new(
path,
- Arc::new(schema.as_ref().to_owned()),
+ Some(Arc::new(schema.as_ref().to_owned())),
*has_header,
+ *delimiter,
projection.to_owned(),
batch_size,
)?)),
@@ -864,11 +870,35 @@ mod tests {
]));
// register each partition as well as the top level dir
- ctx.register_csv("part0", &format!("{}/part-0.csv", out_dir), &schema,
true);
- ctx.register_csv("part1", &format!("{}/part-1.csv", out_dir), &schema,
true);
- ctx.register_csv("part2", &format!("{}/part-2.csv", out_dir), &schema,
true);
- ctx.register_csv("part3", &format!("{}/part-3.csv", out_dir), &schema,
true);
- ctx.register_csv("allparts", &out_dir, &schema, true);
+ ctx.register_csv(
+ "part0",
+ &format!("{}/part-0.csv", out_dir),
+ &schema,
+ true,
+ None,
+ );
+ ctx.register_csv(
+ "part1",
+ &format!("{}/part-1.csv", out_dir),
+ &schema,
+ true,
+ None,
+ );
+ ctx.register_csv(
+ "part2",
+ &format!("{}/part-2.csv", out_dir),
+ &schema,
+ true,
+ None,
+ );
+ ctx.register_csv(
+ "part3",
+ &format!("{}/part-3.csv", out_dir),
+ &schema,
+ true,
+ None,
+ );
+ ctx.register_csv("allparts", &out_dir, &schema, true, None);
let part0 = collect(&mut ctx, "SELECT c1, c2 FROM part0")?;
let part1 = collect(&mut ctx, "SELECT c1, c2 FROM part1")?;
@@ -1031,7 +1061,13 @@ mod tests {
}
// register csv file with the execution context
- ctx.register_csv("test", tmp_dir.path().to_str().unwrap(), &schema,
true);
+ ctx.register_csv(
+ "test",
+ tmp_dir.path().to_str().unwrap(),
+ &schema,
+ true,
+ None,
+ );
Ok(ctx)
}
diff --git a/rust/datafusion/src/execution/physical_plan/csv.rs
b/rust/datafusion/src/execution/physical_plan/csv.rs
index a07417c..14df0d5 100644
--- a/rust/datafusion/src/execution/physical_plan/csv.rs
+++ b/rust/datafusion/src/execution/physical_plan/csv.rs
@@ -18,9 +18,10 @@
//! Execution plan for reading CSV files
use std::fs::File;
+use std::io::BufReader;
use std::sync::{Arc, Mutex};
-use crate::error::Result;
+use crate::error::{ExecutionError, Result};
use crate::execution::physical_plan::common;
use crate::execution::physical_plan::{BatchIterator, ExecutionPlan, Partition};
use arrow::csv;
@@ -35,6 +36,8 @@ pub struct CsvExec {
schema: Arc<Schema>,
/// Does the CSV file have a header?
has_header: bool,
+ /// An optional column delimiter. Defaults to `b','`
+ delimiter: Option<u8>,
/// Optional projection for which columns to load
projection: Option<Vec<usize>>,
/// Batch size
@@ -58,6 +61,7 @@ impl ExecutionPlan for CsvExec {
&filename,
self.schema.clone(),
self.has_header,
+ self.delimiter,
self.projection.clone(),
self.batch_size,
)) as Arc<dyn Partition>
@@ -71,15 +75,35 @@ impl CsvExec {
/// Create a new execution plan for reading a set of CSV files
pub fn try_new(
path: &str,
- schema: Arc<Schema>,
+ schema: Option<Arc<Schema>>,
has_header: bool,
+ delimiter: Option<u8>,
projection: Option<Vec<usize>>,
batch_size: usize,
) -> Result<Self> {
+ let schema = match schema {
+ Some(s) => s,
+ None => {
+ let mut filenames: Vec<String> = vec![];
+ common::build_file_list(path, &mut filenames, ".csv")?;
+ if filenames.is_empty() {
+ return Err(ExecutionError::General("No files
found".to_string()));
+ }
+
+ let f = File::open(&filenames[0])?;
+ Arc::new(csv::infer_file_schema(
+ &mut BufReader::new(f),
+ delimiter.unwrap_or(b','),
+ Some(1000),
+ has_header,
+ )?)
+ }
+ };
Ok(Self {
path: path.to_string(),
schema,
has_header,
+ delimiter,
projection,
batch_size,
})
@@ -94,6 +118,8 @@ struct CsvPartition {
schema: Arc<Schema>,
/// Does the CSV file have a header?
has_header: bool,
+ /// An optional column delimiter. Defaults to `b','`
+ delimiter: Option<u8>,
/// Optional projection for which columns to load
projection: Option<Vec<usize>>,
/// Batch size
@@ -105,6 +131,7 @@ impl CsvPartition {
path: &str,
schema: Arc<Schema>,
has_header: bool,
+ delimiter: Option<u8>,
projection: Option<Vec<usize>>,
batch_size: usize,
) -> Self {
@@ -112,6 +139,7 @@ impl CsvPartition {
path: path.to_string(),
schema,
has_header,
+ delimiter,
projection,
batch_size,
}
@@ -125,6 +153,7 @@ impl Partition for CsvPartition {
&self.path,
self.schema.clone(),
self.has_header,
+ self.delimiter,
&self.projection,
self.batch_size,
)?)))
@@ -143,6 +172,7 @@ impl CsvIterator {
filename: &str,
schema: Arc<Schema>,
has_header: bool,
+ delimiter: Option<u8>,
projection: &Option<Vec<usize>>,
batch_size: usize,
) -> Result<Self> {
@@ -151,6 +181,7 @@ impl CsvIterator {
file,
schema.clone(),
has_header,
+ delimiter,
batch_size,
projection.clone(),
);
diff --git a/rust/datafusion/src/execution/physical_plan/hash_aggregate.rs
b/rust/datafusion/src/execution/physical_plan/hash_aggregate.rs
index 0d1a376..83ffa73 100644
--- a/rust/datafusion/src/execution/physical_plan/hash_aggregate.rs
+++ b/rust/datafusion/src/execution/physical_plan/hash_aggregate.rs
@@ -736,7 +736,7 @@ mod tests {
let partitions = 4;
let path = test::create_partitioned_csv("aggregate_test_100.csv",
partitions)?;
- let csv = CsvExec::try_new(&path, schema.clone(), true, None, 1024)?;
+ let csv = CsvExec::try_new(&path, Some(schema.clone()), true, None,
None, 1024)?;
let group_expr: Vec<Arc<dyn PhysicalExpr>> = vec![col(1,
schema.as_ref())];
diff --git a/rust/datafusion/src/execution/physical_plan/limit.rs
b/rust/datafusion/src/execution/physical_plan/limit.rs
index bdc42df..4aa25f8 100644
--- a/rust/datafusion/src/execution/physical_plan/limit.rs
+++ b/rust/datafusion/src/execution/physical_plan/limit.rs
@@ -183,7 +183,7 @@ mod tests {
let path =
test::create_partitioned_csv("aggregate_test_100.csv",
num_partitions)?;
- let csv = CsvExec::try_new(&path, schema.clone(), true, None, 1024)?;
+ let csv = CsvExec::try_new(&path, Some(schema.clone()), true, None,
None, 1024)?;
// input should have 4 partitions
let input = csv.partitions()?;
diff --git a/rust/datafusion/src/execution/physical_plan/merge.rs
b/rust/datafusion/src/execution/physical_plan/merge.rs
index 0ef8a39..18283a9 100644
--- a/rust/datafusion/src/execution/physical_plan/merge.rs
+++ b/rust/datafusion/src/execution/physical_plan/merge.rs
@@ -111,7 +111,7 @@ mod tests {
let path =
test::create_partitioned_csv("aggregate_test_100.csv",
num_partitions)?;
- let csv = CsvExec::try_new(&path, schema.clone(), true, None, 1024)?;
+ let csv = CsvExec::try_new(&path, Some(schema.clone()), true, None,
None, 1024)?;
// input should have 4 partitions
let input = csv.partitions()?;
diff --git a/rust/datafusion/src/execution/physical_plan/projection.rs
b/rust/datafusion/src/execution/physical_plan/projection.rs
index 3ae8ba0..d90050c 100644
--- a/rust/datafusion/src/execution/physical_plan/projection.rs
+++ b/rust/datafusion/src/execution/physical_plan/projection.rs
@@ -150,7 +150,7 @@ mod tests {
let partitions = 4;
let path = test::create_partitioned_csv("aggregate_test_100.csv",
partitions)?;
- let csv = CsvExec::try_new(&path, schema.clone(), true, None, 1024)?;
+ let csv = CsvExec::try_new(&path, Some(schema.clone()), true, None,
None, 1024)?;
let projection = ProjectionExec::try_new(
vec![Arc::new(Column::new(0, &schema.as_ref().field(0).name()))],
diff --git a/rust/datafusion/src/execution/physical_plan/selection.rs
b/rust/datafusion/src/execution/physical_plan/selection.rs
index ea74a89..e8956e7 100644
--- a/rust/datafusion/src/execution/physical_plan/selection.rs
+++ b/rust/datafusion/src/execution/physical_plan/selection.rs
@@ -159,7 +159,7 @@ mod tests {
let partitions = 4;
let path = test::create_partitioned_csv("aggregate_test_100.csv",
partitions)?;
- let csv = CsvExec::try_new(&path, schema.clone(), true, None, 1024)?;
+ let csv = CsvExec::try_new(&path, Some(schema.clone()), true, None,
None, 1024)?;
let predicate: Arc<dyn PhysicalExpr> = binary(
binary(
diff --git a/rust/datafusion/src/execution/table_impl.rs
b/rust/datafusion/src/execution/table_impl.rs
index 10d65a4..de3bcb6 100644
--- a/rust/datafusion/src/execution/table_impl.rs
+++ b/rust/datafusion/src/execution/table_impl.rs
@@ -267,6 +267,7 @@ mod tests {
&format!("{}/csv/aggregate_test_100.csv", testdata),
&schema,
true,
+ None,
);
}
}
diff --git a/rust/datafusion/src/logicalplan.rs
b/rust/datafusion/src/logicalplan.rs
index fddd60d..ecae04e 100644
--- a/rust/datafusion/src/logicalplan.rs
+++ b/rust/datafusion/src/logicalplan.rs
@@ -25,6 +25,7 @@ use std::fmt;
use arrow::datatypes::{DataType, Field, Schema};
+use crate::datasource::csv::CsvFile;
use crate::datasource::parquet::ParquetTable;
use crate::datasource::TableProvider;
use crate::error::{ExecutionError, Result};
@@ -560,6 +561,8 @@ pub enum LogicalPlan {
schema: Box<Schema>,
/// Whether the CSV file(s) have a header containing column names
has_header: bool,
+ /// An optional column delimiter. Defaults to `b','`
+ delimiter: Option<u8>,
/// Optional column indices to use as a projection
projection: Option<Vec<usize>>,
/// The projected schema
@@ -794,16 +797,27 @@ impl LogicalPlanBuilder {
pub fn scan_csv(
path: &str,
has_header: bool,
- schema: &Schema,
+ schema: Option<&Schema>,
+ delimiter: Option<u8>,
projection: Option<Vec<usize>>,
) -> Result<Self> {
+ let schema: Schema = match schema {
+ Some(s) => s.to_owned(),
+ None => CsvFile::try_new(path, None, has_header, delimiter)?
+ .schema()
+ .as_ref()
+ .to_owned(),
+ };
+
let projected_schema = projection
.clone()
.map(|p| Schema::new(p.iter().map(|i|
schema.field(*i).clone()).collect()));
+
Ok(Self::from(&LogicalPlan::CsvScan {
path: path.to_owned(),
schema: Box::new(schema.to_owned()),
has_header,
+ delimiter,
projection,
projected_schema: Box::new(
projected_schema.or(Some(schema.clone())).unwrap(),
@@ -815,7 +829,6 @@ impl LogicalPlanBuilder {
pub fn scan_parquet(path: &str, projection: Option<Vec<usize>>) ->
Result<Self> {
let p = ParquetTable::try_new(path)?;
let schema = p.schema().as_ref().to_owned();
- println!("{:?}", schema);
let projected_schema = projection
.clone()
.map(|p| Schema::new(p.iter().map(|i|
schema.field(*i).clone()).collect()));
@@ -957,7 +970,8 @@ mod tests {
let plan = LogicalPlanBuilder::scan_csv(
"employee.csv",
true,
- &employee_schema(),
+ Some(&employee_schema()),
+ None,
Some(vec![0, 3]),
)?
.filter(col("state").eq(&lit_str("CO")))?
diff --git a/rust/datafusion/src/optimizer/projection_push_down.rs
b/rust/datafusion/src/optimizer/projection_push_down.rs
index 8f614fc..ea47215 100644
--- a/rust/datafusion/src/optimizer/projection_push_down.rs
+++ b/rust/datafusion/src/optimizer/projection_push_down.rs
@@ -133,6 +133,7 @@ impl ProjectionPushDown {
LogicalPlan::CsvScan {
path,
has_header,
+ delimiter,
schema,
projection,
..
@@ -144,6 +145,7 @@ impl ProjectionPushDown {
path: path.to_owned(),
has_header: *has_header,
schema: schema.clone(),
+ delimiter: *delimiter,
projection: Some(projection),
projected_schema: Box::new(projected_schema),
})
diff --git a/rust/datafusion/tests/sql.rs b/rust/datafusion/tests/sql.rs
index 08aa656..a4bf16e 100644
--- a/rust/datafusion/tests/sql.rs
+++ b/rust/datafusion/tests/sql.rs
@@ -56,7 +56,7 @@ fn nyc() -> Result<()> {
]);
let mut ctx = ExecutionContext::new();
- ctx.register_csv("tripdata", "file.csv", &schema, true);
+ ctx.register_csv("tripdata", "file.csv", &schema, true, None);
let logical_plan = ctx.create_logical_plan(
"SELECT passenger_count, MIN(fare_amount), MAX(fare_amount) \
@@ -442,7 +442,7 @@ fn register_csv(
filename: &str,
schema: &Arc<Schema>,
) {
- ctx.register_csv(name, filename, &schema, true);
+ ctx.register_csv(name, filename, &schema, true, None);
}
fn register_alltypes_parquet(ctx: &mut ExecutionContext) {
