Xuanwo commented on code in PR #23833: URL: https://github.com/apache/datafusion/pull/23833#discussion_r3999440157
########## benchmarks/src/asof.rs: ########## @@ -0,0 +1,226 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::util::{BenchmarkRun, CommonOpt, QueryResult}; +use clap::Args; +use datafusion::physical_plan::execute_stream; +use datafusion::{error::Result, prelude::SessionContext}; +use datafusion_common::instant::Instant; +use datafusion_common::{DataFusionError, exec_datafusion_err, exec_err}; +use futures::StreamExt; + +/// Run end-to-end ASOF join benchmarks. +/// +/// The cases cover broadcast-side size asymmetry, equality-key cardinality and +/// skew, left-side parallelism, optimizer-inserted ordering, wide payload +/// materialization, and descending successor matching. +#[derive(Debug, Args, Clone)] +#[command(verbatim_doc_comment)] +pub struct RunOpt { + /// Query number (between 1 and 6). If not specified, runs all queries + #[arg(short, long)] + query: Option<usize>, + + /// Common options + #[command(flatten)] + common: CommonOpt, + + /// If present, write results json here + #[arg(short = 'o', long = "output")] + output_path: Option<std::path::PathBuf>, +} + +const ASOF_QUERIES: &[&str] = &[ + // Q1: small broadcast input and a large, equality-free probe input + r#" + WITH left_input AS ( + SELECT value AS ts, value AS payload FROM range(1000000) + ), + right_input AS ( + SELECT value AS ts, value AS payload FROM range(10000) + ) + SELECT l.ts, l.payload, r.payload AS right_payload + FROM left_input l + ASOF JOIN right_input r MATCH_CONDITION (l.ts >= r.ts) + "#, + // Q2: grouped predecessor, optimizer partitions left and coalesces right Review Comment: Good catch. Q07 now covers the pre-sorted keyed case. Thanks! ########## datafusion/physical-plan/benches/asof_join.rs: ########## @@ -0,0 +1,219 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Criterion benchmarks for the pre-sorted broadcast ASOF join kernel. + +use std::sync::Arc; + +use arrow::array::{ + ArrayRef, Int64Array, RecordBatch, StringArray, StringDictionaryBuilder, +}; +use arrow::datatypes::{DataType, Field, Int32Type, Schema, SchemaRef}; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_execution::{TaskContext, config::SessionConfig}; +use datafusion_expr::Operator; +use datafusion_physical_expr::expressions::col; +use datafusion_physical_plan::joins::{AsOfJoinExec, AsOfMatchExpr, utils::JoinOn}; +use datafusion_physical_plan::test::TestMemoryExec; +use datafusion_physical_plan::{ExecutionPlan, collect}; +use tokio::runtime::Runtime; + +#[derive(Clone, Copy)] +enum Payload { + Int64, + WideUtf8, + Dictionary, +} + +impl Payload { + fn name(self) -> &'static str { + match self { + Self::Int64 => "int64", + Self::WideUtf8 => "wide_utf8", + Self::Dictionary => "dictionary", + } + } + + fn data_type(self) -> DataType { + match self { + Self::Int64 => DataType::Int64, + Self::WideUtf8 => DataType::Utf8, + Self::Dictionary => { + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)) + } + } + } + + fn array(self, rows: &[(i64, i64, usize)]) -> ArrayRef { + match self { + Self::Int64 => Arc::new(Int64Array::from_iter_values( + rows.iter().map(|(_, _, row)| *row as i64), + )), + Self::WideUtf8 => Arc::new(StringArray::from_iter_values( + rows.iter() + .map(|(_, _, row)| format!("row_{row:08}_{}", "x".repeat(244))), + )), + Self::Dictionary => { + let mut builder = StringDictionaryBuilder::<Int32Type>::new(); + for (_, _, row) in rows { + builder.append_value(format!("category_{}", row % 64)); + } + Arc::new(builder.finish()) + } + } + } +} + +fn schema(payload: Payload) -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int64, false), + Field::new("ts", DataType::Int64, false), + Field::new("payload", payload.data_type(), false), + ])) +} + +fn build_sorted_batches( + num_rows: usize, + num_groups: usize, + time_offset: i64, + payload: Payload, + schema: &SchemaRef, +) -> Vec<RecordBatch> { + let mut rows = (0..num_rows) + .map(|row| { + ( + (row % num_groups) as i64, + (row / num_groups) as i64 + time_offset, + row, + ) + }) + .collect::<Vec<_>>(); + rows.sort_unstable_by_key(|(key, ts, _)| (*key, *ts)); + + let batch = RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(Int64Array::from_iter_values( + rows.iter().map(|(key, _, _)| *key), + )), + Arc::new(Int64Array::from_iter_values( + rows.iter().map(|(_, ts, _)| *ts), + )), + payload.array(&rows), + ], + ) + .unwrap(); + + let mut batches = Vec::new(); + let mut offset = 0; + while offset < batch.num_rows() { + let len = (batch.num_rows() - offset).min(8192); + batches.push(batch.slice(offset, len)); + offset += len; + } + batches +} + +fn partition_batches( + batches: &[RecordBatch], + partition_count: usize, +) -> Vec<Vec<RecordBatch>> { + let mut partitions = vec![Vec::new(); partition_count]; + for (index, batch) in batches.iter().enumerate() { + partitions[index % partition_count].push(batch.clone()); + } + partitions +} + +fn make_exec( + partitions: &[Vec<RecordBatch>], + schema: &SchemaRef, +) -> Arc<dyn ExecutionPlan> { + TestMemoryExec::try_new_exec(partitions, Arc::clone(schema), None).unwrap() +} + +fn do_join( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + rt: &Runtime, +) -> usize { + let on: JoinOn = vec![( + col("key", &left.schema()).unwrap(), + col("key", &right.schema()).unwrap(), + )]; + let left_match = col("ts", &left.schema()).unwrap(); + let right_match = col("ts", &right.schema()).unwrap(); + let join = AsOfJoinExec::try_new( + left, + right, + on, + AsOfMatchExpr::new(left_match, Operator::GtEq, right_match), Review Comment: Covered in the SQL suite now. Thanks! ########## benchmarks/src/asof.rs: ########## @@ -0,0 +1,226 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::util::{BenchmarkRun, CommonOpt, QueryResult}; +use clap::Args; +use datafusion::physical_plan::execute_stream; +use datafusion::{error::Result, prelude::SessionContext}; +use datafusion_common::instant::Instant; +use datafusion_common::{DataFusionError, exec_datafusion_err, exec_err}; +use futures::StreamExt; + +/// Run end-to-end ASOF join benchmarks. +/// +/// The cases cover broadcast-side size asymmetry, equality-key cardinality and +/// skew, left-side parallelism, optimizer-inserted ordering, wide payload +/// materialization, and descending successor matching. +#[derive(Debug, Args, Clone)] +#[command(verbatim_doc_comment)] +pub struct RunOpt { + /// Query number (between 1 and 6). If not specified, runs all queries + #[arg(short, long)] + query: Option<usize>, + + /// Common options + #[command(flatten)] + common: CommonOpt, + + /// If present, write results json here + #[arg(short = 'o', long = "output")] + output_path: Option<std::path::PathBuf>, +} + +const ASOF_QUERIES: &[&str] = &[ + // Q1: small broadcast input and a large, equality-free probe input + r#" + WITH left_input AS ( + SELECT value AS ts, value AS payload FROM range(1000000) + ), + right_input AS ( + SELECT value AS ts, value AS payload FROM range(10000) + ) + SELECT l.ts, l.payload, r.payload AS right_payload + FROM left_input l + ASOF JOIN right_input r MATCH_CONDITION (l.ts >= r.ts) + "#, + // Q2: grouped predecessor, optimizer partitions left and coalesces right Review Comment: Agreed, updated! ########## benchmarks/src/asof.rs: ########## @@ -0,0 +1,226 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::util::{BenchmarkRun, CommonOpt, QueryResult}; +use clap::Args; +use datafusion::physical_plan::execute_stream; +use datafusion::{error::Result, prelude::SessionContext}; +use datafusion_common::instant::Instant; +use datafusion_common::{DataFusionError, exec_datafusion_err, exec_err}; +use futures::StreamExt; + +/// Run end-to-end ASOF join benchmarks. +/// +/// The cases cover broadcast-side size asymmetry, equality-key cardinality and +/// skew, left-side parallelism, optimizer-inserted ordering, wide payload +/// materialization, and descending successor matching. +#[derive(Debug, Args, Clone)] +#[command(verbatim_doc_comment)] +pub struct RunOpt { + /// Query number (between 1 and 6). If not specified, runs all queries + #[arg(short, long)] + query: Option<usize>, + + /// Common options + #[command(flatten)] + common: CommonOpt, + + /// If present, write results json here + #[arg(short = 'o', long = "output")] + output_path: Option<std::path::PathBuf>, +} + +const ASOF_QUERIES: &[&str] = &[ + // Q1: small broadcast input and a large, equality-free probe input + r#" + WITH left_input AS ( + SELECT value AS ts, value AS payload FROM range(1000000) + ), + right_input AS ( + SELECT value AS ts, value AS payload FROM range(10000) + ) + SELECT l.ts, l.payload, r.payload AS right_payload + FROM left_input l + ASOF JOIN right_input r MATCH_CONDITION (l.ts >= r.ts) + "#, + // Q2: grouped predecessor, optimizer partitions left and coalesces right + r#" + WITH left_input AS ( + SELECT value % 10000 AS key, + value / 10000 + 1 AS ts, + value AS payload + FROM range(1000000) + ), + right_input AS ( + SELECT value % 10000 AS key, + value / 10000 AS ts, + value AS payload + FROM range(1000000) + ) + SELECT l.key, l.ts, l.payload, r.payload AS right_payload + FROM left_input l + ASOF JOIN right_input r MATCH_CONDITION (l.ts >= r.ts) + ON l.key = r.key + "#, + // Q3: grouped predecessor with a wide payload + r#" + WITH left_input AS ( + SELECT value % 10000 AS key, + value / 10000 + 1 AS ts, + repeat('x', 256) AS payload + FROM range(250000) + ), + right_input AS ( + SELECT value % 10000 AS key, + value / 10000 AS ts, + repeat('y', 256) AS payload + FROM range(250000) + ) + SELECT l.key, l.ts, l.payload, r.payload AS right_payload + FROM left_input l + ASOF JOIN right_input r MATCH_CONDITION (l.ts >= r.ts) + ON l.key = r.key + "#, + // Q4: successor matching requires descending input order Review Comment: Yep, agreed. The semantics are explicit now. ########## datafusion/physical-plan/benches/asof_join.rs: ########## @@ -0,0 +1,219 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Criterion benchmarks for the pre-sorted broadcast ASOF join kernel. Review Comment: Sounds good. Keeping only the SQL benchmark now. -- 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]
