Skip to content

Commit f9696b6

Browse files
xudong963xudong963
authored
Extract useful methods from sqllogictest bin (#14267)
1 parent 144ebf9 commit f9696b6

File tree

3 files changed

+118
-88
lines changed

3 files changed

+118
-88
lines changed

‎datafusion/sqllogictest/bin/sqllogictests.rs

+7-88
Original file line numberDiff line numberDiff line change
@@ -18,16 +18,19 @@
1818
use clap::Parser;
1919
use datafusion_common::instant::Instant;
2020
use datafusion_common::utils::get_available_parallelism;
21-
use datafusion_common::{exec_datafusion_err, exec_err, DataFusionError, Result};
21+
use datafusion_common::{exec_err, DataFusionError, Result};
2222
use datafusion_common_runtime::SpawnedTask;
23-
use datafusion_sqllogictest::{DataFusion, TestContext};
23+
use datafusion_sqllogictest::{
24+
df_value_validator, read_dir_recursive, setup_scratch_dir, value_normalizer,
25+
DataFusion, TestContext,
26+
};
2427
use futures::stream::StreamExt;
2528
use indicatif::{
2629
HumanDuration, MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle,
2730
};
2831
use itertools::Itertools;
29-
use log::Level::{Info, Warn};
30-
use log::{info, log_enabled, warn};
32+
use log::Level::Info;
33+
use log::{info, log_enabled};
3134
use sqllogictest::{
3235
parse_file, strict_column_validator, AsyncDB, Condition, Normalizer, Record,
3336
Validator,
@@ -38,7 +41,6 @@ use crate::postgres_container::{
3841
initialize_postgres_container, terminate_postgres_container,
3942
};
4043
use std::ffi::OsStr;
41-
use std::fs;
4244
use std::path::{Path, PathBuf};
4345

4446
#[cfg(feature = "postgres")]
@@ -56,14 +58,6 @@ pub fn main() -> Result<()> {
5658
.block_on(run_tests())
5759
}
5860

59-
// Trailing whitespace from lines in SLT will typically be removed, but do not fail if it is not
60-
// If particular test wants to cover trailing whitespace on a value,
61-
// it should project additional non-whitespace column on the right.
62-
#[allow(clippy::ptr_arg)]
63-
fn value_normalizer(s: &String) -> String {
64-
s.trim_end().to_string()
65-
}
66-
6761
fn sqlite_value_validator(
6862
normalizer: Normalizer,
6963
actual: &[Vec<String>],
@@ -93,54 +87,6 @@ fn sqlite_value_validator(
9387
normalized_actual == normalized_expected
9488
}
9589

96-
fn df_value_validator(
97-
normalizer: Normalizer,
98-
actual: &[Vec<String>],
99-
expected: &[String],
100-
) -> bool {
101-
let normalized_expected = expected.iter().map(normalizer).collect::<Vec<_>>();
102-
let normalized_actual = actual
103-
.iter()
104-
.map(|strs| strs.iter().join(" "))
105-
.map(|str| str.trim_end().to_string())
106-
.collect_vec();
107-
108-
if log_enabled!(Warn) && normalized_actual != normalized_expected {
109-
warn!("df validation failed. actual vs expected:");
110-
for i in 0..normalized_actual.len() {
111-
warn!("[{i}] {}<eol>", normalized_actual[i]);
112-
warn!(
113-
"[{i}] {}<eol>",
114-
if normalized_expected.len() >= i {
115-
&normalized_expected[i]
116-
} else {
117-
"No more results"
118-
}
119-
);
120-
}
121-
}
122-
123-
normalized_actual == normalized_expected
124-
}
125-
126-
/// Sets up an empty directory at test_files/scratch/<name>
127-
/// creating it if needed and clearing any file contents if it exists
128-
/// This allows tests for inserting to external tables or copy to
129-
/// persist data to disk and have consistent state when running
130-
/// a new test
131-
fn setup_scratch_dir(name: &Path) -> Result<()> {
132-
// go from copy.slt --> copy
133-
let file_stem = name.file_stem().expect("File should have a stem");
134-
let path = PathBuf::from("test_files").join("scratch").join(file_stem);
135-
136-
info!("Creating scratch dir in {path:?}");
137-
if path.exists() {
138-
fs::remove_dir_all(&path)?;
139-
}
140-
fs::create_dir_all(&path)?;
141-
Ok(())
142-
}
143-
14490
async fn run_tests() -> Result<()> {
14591
// Enable logging (e.g. set RUST_LOG=debug to see debug logs)
14692
env_logger::init();
@@ -573,33 +519,6 @@ fn read_test_files<'a>(
573519
Ok(Box::new(paths.into_iter()))
574520
}
575521

576-
fn read_dir_recursive<P: AsRef<Path>>(path: P) -> Result<Vec<PathBuf>> {
577-
let mut dst = vec![];
578-
read_dir_recursive_impl(&mut dst, path.as_ref())?;
579-
Ok(dst)
580-
}
581-
582-
/// Append all paths recursively to dst
583-
fn read_dir_recursive_impl(dst: &mut Vec<PathBuf>, path: &Path) -> Result<()> {
584-
let entries = fs::read_dir(path)
585-
.map_err(|e| exec_datafusion_err!("Error reading directory {path:?}: {e}"))?;
586-
for entry in entries {
587-
let path = entry
588-
.map_err(|e| {
589-
exec_datafusion_err!("Error reading entry in directory {path:?}: {e}")
590-
})?
591-
.path();
592-
593-
if path.is_dir() {
594-
read_dir_recursive_impl(dst, &path)?;
595-
} else {
596-
dst.push(path);
597-
}
598-
}
599-
600-
Ok(())
601-
}
602-
603522
/// Parsed command line options
604523
///
605524
/// This structure attempts to mimic the command line options of the built-in rust test runner

‎datafusion/sqllogictest/src/lib.rs

+3
Original file line numberDiff line numberDiff line change
@@ -28,4 +28,7 @@ pub use engines::DataFusion;
2828
pub use engines::Postgres;
2929

3030
mod test_context;
31+
mod util;
32+
3133
pub use test_context::TestContext;
34+
pub use util::*;

‎datafusion/sqllogictest/src/util.rs

+108
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
use datafusion_common::{exec_datafusion_err, Result};
19+
use itertools::Itertools;
20+
use log::Level::Warn;
21+
use log::{info, log_enabled, warn};
22+
use sqllogictest::Normalizer;
23+
use std::fs;
24+
use std::path::{Path, PathBuf};
25+
26+
/// Sets up an empty directory at `test_files/scratch/<name>`
27+
/// creating it if needed and clearing any file contents if it exists
28+
/// This allows tests for inserting to external tables or copy to
29+
/// persist data to disk and have consistent state when running
30+
/// a new test
31+
pub fn setup_scratch_dir(name: &Path) -> Result<()> {
32+
// go from copy.slt --> copy
33+
let file_stem = name.file_stem().expect("File should have a stem");
34+
let path = PathBuf::from("test_files").join("scratch").join(file_stem);
35+
36+
info!("Creating scratch dir in {path:?}");
37+
if path.exists() {
38+
fs::remove_dir_all(&path)?;
39+
}
40+
fs::create_dir_all(&path)?;
41+
Ok(())
42+
}
43+
44+
/// Trailing whitespace from lines in SLT will typically be removed, but do not fail if it is not
45+
/// If particular test wants to cover trailing whitespace on a value,
46+
/// it should project additional non-whitespace column on the right.
47+
#[allow(clippy::ptr_arg)]
48+
pub fn value_normalizer(s: &String) -> String {
49+
s.trim_end().to_string()
50+
}
51+
52+
pub fn read_dir_recursive<P: AsRef<Path>>(path: P) -> Result<Vec<PathBuf>> {
53+
let mut dst = vec![];
54+
read_dir_recursive_impl(&mut dst, path.as_ref())?;
55+
Ok(dst)
56+
}
57+
58+
/// Append all paths recursively to dst
59+
fn read_dir_recursive_impl(dst: &mut Vec<PathBuf>, path: &Path) -> Result<()> {
60+
let entries = fs::read_dir(path)
61+
.map_err(|e| exec_datafusion_err!("Error reading directory {path:?}: {e}"))?;
62+
for entry in entries {
63+
let path = entry
64+
.map_err(|e| {
65+
exec_datafusion_err!("Error reading entry in directory {path:?}: {e}")
66+
})?
67+
.path();
68+
69+
if path.is_dir() {
70+
read_dir_recursive_impl(dst, &path)?;
71+
} else {
72+
dst.push(path);
73+
}
74+
}
75+
76+
Ok(())
77+
}
78+
79+
/// Validate the actual and expected values.
80+
pub fn df_value_validator(
81+
normalizer: Normalizer,
82+
actual: &[Vec<String>],
83+
expected: &[String],
84+
) -> bool {
85+
let normalized_expected = expected.iter().map(normalizer).collect::<Vec<_>>();
86+
let normalized_actual = actual
87+
.iter()
88+
.map(|strs| strs.iter().join(" "))
89+
.map(|str| str.trim_end().to_string())
90+
.collect_vec();
91+
92+
if log_enabled!(Warn) && normalized_actual != normalized_expected {
93+
warn!("df validation failed. actual vs expected:");
94+
for i in 0..normalized_actual.len() {
95+
warn!("[{i}] {}<eol>", normalized_actual[i]);
96+
warn!(
97+
"[{i}] {}<eol>",
98+
if normalized_expected.len() >= i {
99+
&normalized_expected[i]
100+
} else {
101+
"No more results"
102+
}
103+
);
104+
}
105+
}
106+
107+
normalized_actual == normalized_expected
108+
}

0 commit comments

Comments
 (0)