villebro commented on code in PR #2791:
URL:
https://github.com/apache/datafusion-ballista/pull/2791#discussion_r4502722546
##########
ballista/client/tests/common/mod.rs:
##########
@@ -258,42 +255,250 @@ async fn connect_to_scheduler(
}
}
-#[allow(dead_code)]
pub async fn standalone_context() -> SessionContext {
SessionContext::standalone().await.unwrap()
}
-#[allow(dead_code)]
+/// A standalone context with `alltypes_plain.parquet` registered as `test`.
+pub async fn alltypes_context() -> SessionContext {
+ let ctx = standalone_context().await;
+ register_alltypes(&ctx, "test").await;
+ ctx
+}
+
pub async fn remote_context() -> SessionContext {
let (host, port) = setup_test_cluster().await;
SessionContext::remote(&format!("df://{host}:{port}"))
.await
.unwrap()
}
-#[allow(dead_code)]
pub async fn standalone_context_with_state() -> SessionContext {
- let config = SessionConfig::new_with_ballista();
- let state = SessionStateBuilder::new()
- .with_config(config)
- .with_default_features()
- .build();
+ let state = state_with_config(SessionConfig::new_with_ballista());
SessionContext::standalone_with_state(state).await.unwrap()
}
-#[allow(dead_code)]
pub async fn remote_context_with_state() -> SessionContext {
- let config = SessionConfig::new_with_ballista();
- let state = SessionStateBuilder::new()
- .with_config(config)
- .with_default_features()
- .build();
+ let state = state_with_config(SessionConfig::new_with_ballista());
let (host, port) = setup_test_cluster_with_state(state.clone()).await;
SessionContext::remote_with_state(&format!("df://{host}:{port}"), state)
.await
.unwrap()
}
+/// A session state with DataFusion's default features and `config`.
+pub fn state_with_config(config: SessionConfig) -> SessionState {
+ SessionStateBuilder::new()
+ .with_config(config)
+ .with_default_features()
+ .build()
+}
+
+/// A Ballista session state with four target partitions and four standalone
+/// executor vcores, so a stage of four partitions runs at once whatever the
+/// machine's core count. `options` are set on top, in order.
+pub fn ballista_state(options: &[(&str, &str)]) -> SessionState {
+ let mut config = SessionConfig::new_with_ballista()
+ .with_target_partitions(4)
+ .set_str(BALLISTA_STANDALONE_PARALLELISM, "4");
+ for (key, value) in options {
+ config = config.set_str(key, value);
+ }
+ state_with_config(config)
+}
+
+/// A DataFusion session with four target partitions, used as the reference
+/// for results from a Ballista session.
+pub fn datafusion_context() -> SessionContext {
+
SessionContext::new_with_config(SessionConfig::new().with_target_partitions(4))
+}
+
+/// Writes the tables of [`orders_customers_dir`] to `dir`. `orders` is split
+/// over four files, so its scan has four partitions.
+fn write_orders_customers(dir: &Path) {
+ let orders = dir.join("orders");
+ std::fs::create_dir(&orders).unwrap();
+ for file in 0..4 {
+ let mut out =
+
std::fs::File::create(orders.join(format!("part-{file}.csv"))).unwrap();
+ writeln!(out, "o_id,o_cust,o_amount").unwrap();
+ for i in 0..250 {
+ let id = file * 250 + i;
+ writeln!(out, "{id},{},{}", id % 97, (id * 7) % 1000).unwrap();
+ }
+ }
+ let customers = dir.join("customers");
+ std::fs::create_dir(&customers).unwrap();
+ let mut out = std::fs::File::create(customers.join("part-0.csv")).unwrap();
+ writeln!(out, "c_id,c_name,c_region").unwrap();
+ for id in 0..100 {
+ writeln!(out, "{id},customer-{id},region-{}", id % 5).unwrap();
+ }
+}
+
+/// A temporary directory holding an `orders` table of four files and a small
+/// `customers` table, one directory each.
+pub fn orders_customers_dir() -> TempDir {
+ let dir = TempDir::new().unwrap();
+ write_orders_customers(dir.path());
+ dir
+}
+
+/// Registers the tables of [`orders_customers_dir`].
+pub async fn register_orders_customers(ctx: &SessionContext, dir: &Path) {
+ register_csv_tables(ctx, dir, &["orders", "customers"]).await;
+}
+
+/// Queries over the tables of [`orders_customers_dir`], each with whether its
+/// result has no defined row order, so its rows are compared sorted. The first
+/// three are a scan, an aggregate and a join.
+pub const ORDERS_QUERIES: &[(&str, bool)] = &[
+ // Scan and a collapse to one partition.
+ ("SELECT count(*), sum(o_amount) FROM orders", false),
+ // Hash repartition, aggregate, sort.
+ (
+ "SELECT o_cust, count(*) AS n, sum(o_amount) AS total \
+ FROM orders GROUP BY o_cust ORDER BY o_cust",
+ false,
+ ),
+ // Join (broadcast, since customers is small), aggregate, sort.
+ (
+ "SELECT c_region, count(*) AS n, sum(o_amount) AS total \
+ FROM orders JOIN customers ON o_cust = c_id \
+ GROUP BY c_region ORDER BY c_region",
+ false,
+ ),
+ // Top-k over an aggregate.
+ (
+ "SELECT o_cust, sum(o_amount) AS total FROM orders \
+ GROUP BY o_cust ORDER BY total DESC, o_cust LIMIT 5",
+ false,
+ ),
+ // Filter with no exchange but the final merge.
+ (
+ "SELECT o_id, o_amount FROM orders WHERE o_amount > 990 ORDER BY o_id",
+ false,
+ ),
+ // Union of two scans.
+ (
+ "SELECT o_id AS id FROM orders WHERE o_id < 10 \
+ UNION ALL SELECT c_id AS id FROM customers WHERE c_id < 5",
+ true,
+ ),
+ // Window function over a hash-partitioned input.
+ (
+ "SELECT o_cust, o_id, \
+ row_number() OVER (PARTITION BY o_cust ORDER BY o_id) AS rn \
+ FROM orders WHERE o_cust < 3 ORDER BY o_cust, o_id",
+ false,
+ ),
+ // A scalar subquery, joined to every row: each task of the join's
+ // stage reads the subquery's whole result.
+ (
+ "SELECT o_id, o_amount FROM orders \
+ WHERE o_amount > (SELECT avg(o_amount) + 400 FROM orders) ORDER BY
o_id",
+ false,
+ ),
+];
+
+/// A standalone cluster over [`ballista_state`] with `options`, with the
+/// tables of [`orders_customers_dir`] registered.
+pub async fn orders_customers_context(
+ dir: &Path,
+ options: &[(&str, &str)],
+) -> SessionContext {
+ let ctx = SessionContext::standalone_with_state(ballista_state(options))
+ .await
+ .unwrap();
+ register_orders_customers(&ctx, dir).await;
+ ctx
+}
+
+/// Registers each directory `dir/<name>` of header-bearing `.csv` files as
+/// the table `<name>`.
+pub async fn register_csv_tables(ctx: &SessionContext, dir: &Path, names:
&[&str]) {
+ for name in names {
+ ctx.register_csv(
+ *name,
+ dir.join(name).to_str().unwrap(),
+ CsvReadOptions::new()
+ .has_header(true)
+ .file_extension(".csv"),
+ )
+ .await
+ .unwrap();
+ }
+}
+
+/// Writes `batch` to the Parquet file `path`.
+pub fn write_parquet_file(path: &Path, batch: &RecordBatch) {
+ let file = std::fs::File::create(path).expect("create parquet file");
+ let mut writer = ArrowWriter::try_new(file, batch.schema(), None).unwrap();
+ writer.write(batch).unwrap();
+ writer.close().unwrap();
+}
+
+/// Runs `SET key = 'value'`.
+pub async fn set_option(ctx: &SessionContext, key: &str, value: &str) {
+ ctx.sql(&format!("SET {key} = '{value}'"))
Review Comment:
Good idea, I added escaping
https://github.com/apache/datafusion-ballista/pull/2791/changes/d998710f2bb3a77dcee7e9e6e05eedeb87b7f6c9
so we're future proof if any `SET` operation needs to contain single quotes.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
