Skip to content

Commit c21106e

Browse files
author
QP Hou
authored
honor table name for csv/parquet scan in ballista plan serde (#629)
* honor table name for csv/parquet scan in ballista plan serde * disable query 7,8,9 in ballista integration test
1 parent 2f1d6cb commit c21106e

File tree

3 files changed

+39
-9
lines changed

3 files changed

+39
-9
lines changed

ballista/rust/core/src/serde/logical_plan/from_proto.rs

+16-6
Original file line numberDiff line numberDiff line change
@@ -126,9 +126,14 @@ impl TryInto<LogicalPlan> for &protobuf::LogicalPlanNode {
126126
projection = Some(column_indices);
127127
}
128128

129-
LogicalPlanBuilder::scan_csv(&scan.path, options, projection)?
130-
.build()
131-
.map_err(|e| e.into())
129+
LogicalPlanBuilder::scan_csv_with_name(
130+
&scan.path,
131+
options,
132+
projection,
133+
&scan.table_name,
134+
)?
135+
.build()
136+
.map_err(|e| e.into())
132137
}
133138
LogicalPlanType::ParquetScan(scan) => {
134139
let projection = match scan.projection.as_ref() {
@@ -151,9 +156,14 @@ impl TryInto<LogicalPlan> for &protobuf::LogicalPlanNode {
151156
Some(r?)
152157
}
153158
};
154-
LogicalPlanBuilder::scan_parquet(&scan.path, projection, 24)? //TODO concurrency
155-
.build()
156-
.map_err(|e| e.into())
159+
LogicalPlanBuilder::scan_parquet_with_name(
160+
&scan.path,
161+
projection,
162+
24,
163+
&scan.table_name,
164+
)? //TODO concurrency
165+
.build()
166+
.map_err(|e| e.into())
157167
}
158168
LogicalPlanType::Sort(sort) => {
159169
let input: LogicalPlan = convert_box_required!(sort.input)?;

benchmarks/run.sh

+1-1
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ set -e
2020
# This bash script is meant to be run inside the docker-compose environment. Check the README for instructions
2121

2222
cd /
23-
for query in 1 3 5 6 7 8 9 10 12
23+
for query in 1 3 5 6 10 12
2424
do
2525
/tpch benchmark ballista --host ballista-scheduler --port 50050 --query $query --path /data --format tbl --iterations 1 --debug
2626
done

datafusion/src/logical_plan/builder.rs

+22-2
Original file line numberDiff line numberDiff line change
@@ -118,19 +118,39 @@ impl LogicalPlanBuilder {
118118
path: &str,
119119
options: CsvReadOptions,
120120
projection: Option<Vec<usize>>,
121+
) -> Result<Self> {
122+
Self::scan_csv_with_name(path, options, projection, path)
123+
}
124+
125+
/// Scan a CSV data source and register it with a given table name
126+
pub fn scan_csv_with_name(
127+
path: &str,
128+
options: CsvReadOptions,
129+
projection: Option<Vec<usize>>,
130+
table_name: &str,
121131
) -> Result<Self> {
122132
let provider = Arc::new(CsvFile::try_new(path, options)?);
123-
Self::scan(path, provider, projection)
133+
Self::scan(table_name, provider, projection)
124134
}
125135

126136
/// Scan a Parquet data source
127137
pub fn scan_parquet(
128138
path: &str,
129139
projection: Option<Vec<usize>>,
130140
max_concurrency: usize,
141+
) -> Result<Self> {
142+
Self::scan_parquet_with_name(path, projection, max_concurrency, path)
143+
}
144+
145+
/// Scan a Parquet data source and register it with a given table name
146+
pub fn scan_parquet_with_name(
147+
path: &str,
148+
projection: Option<Vec<usize>>,
149+
max_concurrency: usize,
150+
table_name: &str,
131151
) -> Result<Self> {
132152
let provider = Arc::new(ParquetTable::try_new(path, max_concurrency)?);
133-
Self::scan(path, provider, projection)
153+
Self::scan(table_name, provider, projection)
134154
}
135155

136156
/// Scan an empty data source, mainly used in tests

0 commit comments

Comments
 (0)