Skip to content

Commit b410fce

Browse files
authored
Handle trailing tbl column in TPCH benchmarks (#4821)
* Handle trailing tbl column in TPCH benchmarks * Clippy * Fix benchmarks
1 parent 702e117 commit b410fce

2 files changed

Lines changed: 28 additions & 8 deletions

File tree

‎benchmarks/src/bin/tpch.rs‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -402,7 +402,6 @@ async fn get_table(
402402
unimplemented!("Invalid file format '{}'", other);
403403
}
404404
};
405-
let schema = Arc::new(get_tpch_table_schema(table));
406405

407406
let options = ListingOptions::new(format)
408407
.with_file_extension(extension)
@@ -412,10 +411,11 @@ async fn get_table(
412411
let table_path = ListingTableUrl::parse(path)?;
413412
let config = ListingTableConfig::new(table_path).with_listing_options(options);
414413

415-
let config = if table_format == "parquet" {
416-
config.infer_schema(&state).await?
417-
} else {
418-
config.with_schema(schema)
414+
let config = match table_format {
415+
"parquet" => config.infer_schema(&state).await?,
416+
"tbl" => config.with_schema(Arc::new(get_tbl_tpch_table_schema(table))),
417+
"csv" => config.with_schema(Arc::new(get_tpch_table_schema(table))),
418+
_ => unreachable!(),
419419
};
420420

421421
Ok(Arc::new(ListingTable::try_new(config)?))
@@ -827,6 +827,7 @@ mod tests {
827827
#[cfg(feature = "ci")]
828828
mod ci {
829829
use super::*;
830+
use arrow::datatypes::{DataType, Field};
830831
use datafusion_proto::bytes::{logical_plan_from_bytes, logical_plan_to_bytes};
831832

832833
async fn serde_round_trip(query: usize) -> Result<()> {
@@ -1086,7 +1087,6 @@ mod ci {
10861087
/// * the correct number of rows are returned
10871088
/// * the content of the rows is correct
10881089
async fn verify_query(n: usize) -> Result<()> {
1089-
use datafusion::arrow::datatypes::{DataType, Field};
10901090
use datafusion::common::ScalarValue;
10911091
use datafusion::logical_expr::expr::Cast;
10921092

@@ -1214,7 +1214,8 @@ mod ci {
12141214
}
12151215

12161216
fn get_tpch_data_path() -> Result<String> {
1217-
let path = std::env::var("TPCH_DATA").unwrap_or("benchmarks/data".to_string());
1217+
let path =
1218+
std::env::var("TPCH_DATA").unwrap_or_else(|_| "benchmarks/data".to_string());
12181219
if !Path::new(&path).exists() {
12191220
return Err(DataFusionError::Execution(format!(
12201221
"Benchmark data not found (set TPCH_DATA env var to override): {}",

‎benchmarks/src/tpch.rs‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,15 @@ pub const TPCH_TABLES: &[&str] = &[
4343
"part", "supplier", "partsupp", "customer", "orders", "lineitem", "nation", "region",
4444
];
4545

46+
/// The `.tbl` file contains a trailing column
47+
pub fn get_tbl_tpch_table_schema(table: &str) -> Schema {
48+
let mut schema = get_tpch_table_schema(table);
49+
schema
50+
.fields
51+
.push(Field::new("__placeholder", DataType::Utf8, false));
52+
schema
53+
}
54+
4655
/// Get the schema for the benchmarks derived from TPC-H
4756
pub fn get_tpch_table_schema(table: &str) -> Schema {
4857
// note that the schema intentionally uses signed integers so that any generated Parquet
@@ -331,7 +340,7 @@ pub async fn convert_tbl(
331340
let output_root_path = Path::new(output_path);
332341
for table in TPCH_TABLES {
333342
let start = Instant::now();
334-
let schema = get_tpch_table_schema(table);
343+
let schema = get_tbl_tpch_table_schema(table);
335344

336345
let input_path = format!("{input_path}/{table}.tbl");
337346
let options = CsvReadOptions::new()
@@ -346,6 +355,16 @@ pub async fn convert_tbl(
346355
// build plan to read the TBL file
347356
let mut csv = ctx.read_csv(&input_path, options).await?;
348357

358+
// Select all apart from the padding column
359+
let selection = csv
360+
.schema()
361+
.fields()
362+
.iter()
363+
.take(schema.fields.len() - 1)
364+
.map(|d| Expr::Column(d.qualified_column()))
365+
.collect();
366+
367+
csv = csv.select(selection)?;
349368
// optionally, repartition the file
350369
if partitions > 1 {
351370
csv = csv.repartition(Partitioning::RoundRobinBatch(partitions))?

0 commit comments

Comments
 (0)