diff --git a/datafusion/core/src/physical_plan/common.rs b/datafusion/core/src/physical_plan/common.rs index 18b13444d65ce..e766c225b51cd 100644 --- a/datafusion/core/src/physical_plan/common.rs +++ b/datafusion/core/src/physical_plan/common.rs @@ -21,6 +21,7 @@ use super::SendableRecordBatchStream; use crate::error::{DataFusionError, Result}; use crate::execution::context::TaskContext; use crate::execution::memory_pool::MemoryReservation; +use crate::physical_plan::stream::RecordBatchReceiverStream; use crate::physical_plan::{displayable, ColumnStatistics, ExecutionPlan, Statistics}; use arrow::datatypes::Schema; use arrow::ipc::writer::{FileWriter, IpcWriteOptions}; @@ -131,6 +132,31 @@ pub(crate) fn spawn_execution( }) } +/// If running in a tokio context spawns the execution of `stream` to a separate task +/// allowing it to execute in parallel with an intermediate buffer of size `buffer` +pub(crate) fn spawn_buffered( + mut input: SendableRecordBatchStream, + buffer: usize, +) -> SendableRecordBatchStream { + // Use tokio only if running from a tokio context (#2201) + let handle = match tokio::runtime::Handle::try_current() { + Ok(handle) => handle, + Err(_) => return input, + }; + + let schema = input.schema(); + let (sender, receiver) = mpsc::channel(buffer); + let join = handle.spawn(async move { + while let Some(item) = input.next().await { + if sender.send(item).await.is_err() { + return; + } + } + }); + + RecordBatchReceiverStream::create(&schema, receiver, join) +} + /// Computes the statistics for an in-memory RecordBatch /// /// Only computes statistics that are in arrows metadata (num rows, byte size and nulls) diff --git a/datafusion/core/src/physical_plan/file_format/parquet.rs b/datafusion/core/src/physical_plan/file_format/parquet.rs index 641be002b6adc..9843ecc22e4ce 100644 --- a/datafusion/core/src/physical_plan/file_format/parquet.rs +++ b/datafusion/core/src/physical_plan/file_format/parquet.rs @@ -482,7 +482,10 @@ impl FileOpener for ParquetOpener { let table_schema = self.table_schema.clone(); let reorder_predicates = self.reorder_filters; let pushdown_filters = self.pushdown_filters; - let enable_page_index = self.enable_page_index; + let enable_page_index = should_enable_page_index( + self.enable_page_index, + &self.page_pruning_predicate, + ); let limit = self.limit; Ok(Box::pin(async move { @@ -572,6 +575,18 @@ impl FileOpener for ParquetOpener { } } +fn should_enable_page_index( + enable_page_index: bool, + page_pruning_predicate: &Option>, +) -> bool { + enable_page_index + && page_pruning_predicate.is_some() + && page_pruning_predicate + .as_ref() + .map(|p| p.filter_number() > 0) + .unwrap_or(false) +} + /// Factory of parquet file readers. /// /// Provides means to implement custom data access interface. diff --git a/datafusion/core/src/physical_plan/file_format/parquet/page_filter.rs b/datafusion/core/src/physical_plan/file_format/parquet/page_filter.rs index 410c43c57d71c..00e55c41ad09b 100644 --- a/datafusion/core/src/physical_plan/file_format/parquet/page_filter.rs +++ b/datafusion/core/src/physical_plan/file_format/parquet/page_filter.rs @@ -222,6 +222,11 @@ impl PagePruningPredicate { file_metrics.page_index_rows_filtered.add(total_skip); Ok(Some(final_selection)) } + + /// Returns the number of filters in the [`PagePruningPredicate`] + pub fn filter_number(&self) -> usize { + self.predicates.len() + } } /// Returns the column index in the row group metadata for the single diff --git a/datafusion/core/src/physical_plan/sorts/sort.rs b/datafusion/core/src/physical_plan/sorts/sort.rs index 12018944d4cd8..544abb24740dd 100644 --- a/datafusion/core/src/physical_plan/sorts/sort.rs +++ b/datafusion/core/src/physical_plan/sorts/sort.rs @@ -25,7 +25,7 @@ use crate::execution::memory_pool::{ human_readable_size, MemoryConsumer, MemoryReservation, }; use crate::execution::runtime_env::RuntimeEnv; -use crate::physical_plan::common::{batch_byte_size, IPCWriter}; +use crate::physical_plan::common::{batch_byte_size, spawn_buffered, IPCWriter}; use crate::physical_plan::expressions::PhysicalSortExpr; use crate::physical_plan::metrics::{ BaselineMetrics, CompositeMetricsSet, MemTrackingMetrics, MetricsSet, @@ -284,11 +284,13 @@ impl ExternalSorter { self.partition_id, &self.runtime.memory_pool, ); - sort_batch_stream(batch, self.expr.clone(), self.fetch, metrics) + Ok(spawn_buffered( + sort_batch_stream(batch, self.expr.clone(), self.fetch, metrics)?, + 1, + )) }) .collect::>()?; - // TODO: Run batch sorts concurrently (#6162) // TODO: Pushdown fetch to streaming merge (#6000) streaming_merge( diff --git a/datafusion/core/src/physical_plan/sorts/sort_preserving_merge.rs b/datafusion/core/src/physical_plan/sorts/sort_preserving_merge.rs index b78058360184d..88916c93f0f45 100644 --- a/datafusion/core/src/physical_plan/sorts/sort_preserving_merge.rs +++ b/datafusion/core/src/physical_plan/sorts/sort_preserving_merge.rs @@ -22,18 +22,17 @@ use std::sync::Arc; use arrow::datatypes::SchemaRef; use log::{debug, trace}; -use tokio::sync::mpsc; use crate::error::{DataFusionError, Result}; use crate::execution::context::TaskContext; +use crate::physical_plan::common::spawn_buffered; use crate::physical_plan::metrics::{ ExecutionPlanMetricsSet, MemTrackingMetrics, MetricsSet, }; use crate::physical_plan::sorts::streaming_merge; -use crate::physical_plan::stream::RecordBatchReceiverStream; use crate::physical_plan::{ - common::spawn_execution, expressions::PhysicalSortExpr, DisplayFormatType, - Distribution, ExecutionPlan, Partitioning, SendableRecordBatchStream, Statistics, + expressions::PhysicalSortExpr, DisplayFormatType, Distribution, ExecutionPlan, + Partitioning, SendableRecordBatchStream, Statistics, }; use datafusion_physical_expr::{EquivalenceProperties, PhysicalSortRequirement}; @@ -181,29 +180,12 @@ impl ExecutionPlan for SortPreservingMergeExec { result } _ => { - // Use tokio only if running from a tokio context (#2201) - let receivers = match tokio::runtime::Handle::try_current() { - Ok(_) => (0..input_partitions) - .map(|part_i| { - let (sender, receiver) = mpsc::channel(1); - let join_handle = spawn_execution( - self.input.clone(), - sender, - part_i, - context.clone(), - ); - - RecordBatchReceiverStream::create( - &schema, - receiver, - join_handle, - ) - }) - .collect(), - Err(_) => (0..input_partitions) - .map(|partition| self.input.execute(partition, context.clone())) - .collect::>()?, - }; + let receivers = (0..input_partitions) + .map(|partition| { + let stream = self.input.execute(partition, context.clone())?; + Ok(spawn_buffered(stream, 1)) + }) + .collect::>()?; debug!("Done setting up sender-receiver for SortPreservingMergeExec::execute"); @@ -262,6 +244,7 @@ mod tests { use crate::physical_plan::memory::MemoryExec; use crate::physical_plan::metrics::MetricValue; use crate::physical_plan::sorts::sort::SortExec; + use crate::physical_plan::stream::RecordBatchReceiverStream; use crate::physical_plan::{collect, common}; use crate::prelude::{SessionConfig, SessionContext}; use crate::test::exec::{assert_strong_count_converges_to_zero, BlockingExec}; @@ -812,7 +795,7 @@ mod tests { let mut streams = Vec::with_capacity(partition_count); for partition in 0..partition_count { - let (sender, receiver) = mpsc::channel(1); + let (sender, receiver) = tokio::sync::mpsc::channel(1); let mut stream = batches.execute(partition, task_ctx.clone()).unwrap(); let join_handle = tokio::spawn(async move { while let Some(batch) = stream.next().await { diff --git a/datafusion/core/tests/parquet/page_pruning.rs b/datafusion/core/tests/parquet/page_pruning.rs index baf9d2d36a17c..1d444326bbc43 100644 --- a/datafusion/core/tests/parquet/page_pruning.rs +++ b/datafusion/core/tests/parquet/page_pruning.rs @@ -23,6 +23,7 @@ use datafusion::datasource::listing::PartitionedFile; use datafusion::datasource::object_store::ObjectStoreUrl; use datafusion::execution::context::SessionState; use datafusion::physical_plan::file_format::{FileScanConfig, ParquetExec}; +use datafusion::physical_plan::metrics::MetricValue; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionContext; use datafusion_common::{ScalarValue, Statistics, ToDFSchema}; @@ -714,3 +715,41 @@ async fn prune_decimal_in_list() { ) .await; } + +#[tokio::test] +async fn without_pushdown_filter() { + let mut context = ContextWithParquet::new(Scenario::Timestamps, Page).await; + + let output1 = context.query("SELECT * FROM t").await; + + let mut context = ContextWithParquet::new(Scenario::Timestamps, Page).await; + + let output2 = context + .query("SELECT * FROM t where nanos < to_timestamp('2023-01-02 01:01:11Z')") + .await; + + let bytes_scanned_without_filter = cast_count_metric( + output1 + .parquet_metrics + .sum_by_name("bytes_scanned") + .unwrap(), + ) + .unwrap(); + let bytes_scanned_with_filter = cast_count_metric( + output2 + .parquet_metrics + .sum_by_name("bytes_scanned") + .unwrap(), + ) + .unwrap(); + + // Without filter will not read pageIndex. + assert!(bytes_scanned_with_filter > bytes_scanned_without_filter); +} + +fn cast_count_metric(metric: MetricValue) -> Option { + match metric { + MetricValue::Count { count, .. } => Some(count.value()), + _ => None, + } +} diff --git a/datafusion/core/tests/sql/joins.rs b/datafusion/core/tests/sql/joins.rs index 3213ae3e0595e..54ecfe4951dc7 100644 --- a/datafusion/core/tests/sql/joins.rs +++ b/datafusion/core/tests/sql/joins.rs @@ -18,392 +18,6 @@ use super::*; use datafusion::from_slice::FromSlice; -#[tokio::test] -async fn equijoin() -> Result<()> { - let ctx = create_join_context_qualified("t1", "t2")?; - let equivalent_sql = [ - "SELECT t1.a, t2.b FROM t1 INNER JOIN t2 ON t1.a = t2.a ORDER BY t1.a", - "SELECT t1.a, t2.b FROM t1 INNER JOIN t2 ON t2.a = t1.a ORDER BY t1.a", - ]; - let expected = vec![ - "+---+-----+", - "| a | b |", - "+---+-----+", - "| 1 | 100 |", - "| 2 | 200 |", - "| 4 | 400 |", - "+---+-----+", - ]; - for sql in equivalent_sql.iter() { - let actual = execute_to_batches(&ctx, sql).await; - assert_batches_eq!(expected, &actual); - } - - Ok(()) -} - -#[tokio::test] -async fn left_join_unbalanced() -> Result<()> { - // the t1_id is larger than t2_id so the join_selection optimizer should kick in - let ctx = create_join_context_unbalanced("t1_id", "t2_id")?; - let equivalent_sql = [ - "SELECT t1_id, t1_name, t2_name FROM t1 LEFT JOIN t2 ON t1_id = t2_id ORDER BY t1_id", - "SELECT t1_id, t1_name, t2_name FROM t1 LEFT JOIN t2 ON t2_id = t1_id ORDER BY t1_id", - ]; - let expected = vec![ - "+-------+---------+---------+", - "| t1_id | t1_name | t2_name |", - "+-------+---------+---------+", - "| 11 | a | z |", - "| 22 | b | y |", - "| 33 | c | |", - "| 44 | d | x |", - "| 77 | e | |", - "+-------+---------+---------+", - ]; - for sql in equivalent_sql.iter() { - let actual = execute_to_batches(&ctx, sql).await; - assert_batches_eq!(expected, &actual); - } - Ok(()) -} - -#[tokio::test] -async fn left_join_null_filter() -> Result<()> { - // Since t2 is the non-preserved side of the join, we cannot push down a NULL filter. - // Note that this is only true because IS NULL does not remove nulls. For filters that - // remove nulls, we can rewrite the join as an inner join and then push down the filter. - let ctx = create_join_context_with_nulls()?; - let sql = "SELECT t1_id, t2_id, t2_name FROM t1 LEFT JOIN t2 ON t1_id = t2_id WHERE t2_name IS NULL ORDER BY t1_id"; - let expected = vec![ - "+-------+-------+---------+", - "| t1_id | t2_id | t2_name |", - "+-------+-------+---------+", - "| 22 | 22 | |", - "| 33 | | |", - "| 77 | | |", - "| 88 | | |", - "+-------+-------+---------+", - ]; - - let actual = execute_to_batches(&ctx, sql).await; - assert_batches_eq!(expected, &actual); - Ok(()) -} - -#[tokio::test] -async fn left_join_null_filter_on_join_column() -> Result<()> { - // Again, since t2 is the non-preserved side of the join, we cannot push down a NULL filter. - let ctx = create_join_context_with_nulls()?; - let sql = "SELECT t1_id, t2_id, t2_name FROM t1 LEFT JOIN t2 ON t1_id = t2_id WHERE t2_id IS NULL ORDER BY t1_id"; - let expected = vec![ - "+-------+-------+---------+", - "| t1_id | t2_id | t2_name |", - "+-------+-------+---------+", - "| 33 | | |", - "| 77 | | |", - "| 88 | | |", - "+-------+-------+---------+", - ]; - - let actual = execute_to_batches(&ctx, sql).await; - assert_batches_eq!(expected, &actual); - Ok(()) -} - -#[tokio::test] -async fn left_join_not_null_filter() -> Result<()> { - let ctx = create_join_context_with_nulls()?; - let sql = "SELECT t1_id, t2_id, t2_name FROM t1 LEFT JOIN t2 ON t1_id = t2_id WHERE t2_name IS NOT NULL ORDER BY t1_id"; - let expected = vec![ - "+-------+-------+---------+", - "| t1_id | t2_id | t2_name |", - "+-------+-------+---------+", - "| 11 | 11 | z |", - "| 44 | 44 | x |", - "| 99 | 99 | u |", - "+-------+-------+---------+", - ]; - - let actual = execute_to_batches(&ctx, sql).await; - assert_batches_eq!(expected, &actual); - Ok(()) -} - -#[tokio::test] -async fn left_join_not_null_filter_on_join_column() -> Result<()> { - let ctx = create_join_context_with_nulls()?; - let sql = "SELECT t1_id, t2_id, t2_name FROM t1 LEFT JOIN t2 ON t1_id = t2_id WHERE t2_id IS NOT NULL ORDER BY t1_id"; - let expected = vec![ - "+-------+-------+---------+", - "| t1_id | t2_id | t2_name |", - "+-------+-------+---------+", - "| 11 | 11 | z |", - "| 22 | 22 | |", - "| 44 | 44 | x |", - "| 99 | 99 | u |", - "+-------+-------+---------+", - ]; - - let actual = execute_to_batches(&ctx, sql).await; - assert_batches_eq!(expected, &actual); - Ok(()) -} - -#[tokio::test] -async fn self_join_non_equijoin() -> Result<()> { - let ctx = create_join_context_with_nulls()?; - let sql = - "SELECT x.t1_id, y.t1_id FROM t1 x JOIN t1 y ON x.t1_id = 11 AND y.t1_id = 44"; - let expected = vec![ - "+-------+-------+", - "| t1_id | t1_id |", - "+-------+-------+", - "| 11 | 44 |", - "+-------+-------+", - ]; - - let actual = execute_to_batches(&ctx, sql).await; - assert_batches_eq!(expected, &actual); - Ok(()) -} - -#[tokio::test] -async fn right_join_null_filter() -> Result<()> { - let ctx = create_join_context_with_nulls()?; - let sql = "SELECT t1_id, t1_name, t2_id FROM t1 RIGHT JOIN t2 ON t1_id = t2_id WHERE t1_name IS NULL ORDER BY t2_id"; - let expected = vec![ - "+-------+---------+-------+", - "| t1_id | t1_name | t2_id |", - "+-------+---------+-------+", - "| | | 55 |", - "| 99 | | 99 |", - "+-------+---------+-------+", - ]; - - let actual = execute_to_batches(&ctx, sql).await; - assert_batches_eq!(expected, &actual); - Ok(()) -} - -#[tokio::test] -async fn right_join_null_filter_on_join_column() -> Result<()> { - let ctx = create_join_context_with_nulls()?; - let sql = "SELECT t1_id, t1_name, t2_id FROM t1 RIGHT JOIN t2 ON t1_id = t2_id WHERE t1_id IS NULL ORDER BY t2_id"; - let expected = vec![ - "+-------+---------+-------+", - "| t1_id | t1_name | t2_id |", - "+-------+---------+-------+", - "| | | 55 |", - "+-------+---------+-------+", - ]; - - let actual = execute_to_batches(&ctx, sql).await; - assert_batches_eq!(expected, &actual); - Ok(()) -} - -#[tokio::test] -async fn right_join_not_null_filter() -> Result<()> { - let ctx = create_join_context_with_nulls()?; - let sql = "SELECT t1_id, t1_name, t2_id FROM t1 RIGHT JOIN t2 ON t1_id = t2_id WHERE t1_name IS NOT NULL ORDER BY t2_id"; - let expected = vec![ - "+-------+---------+-------+", - "| t1_id | t1_name | t2_id |", - "+-------+---------+-------+", - "| 11 | a | 11 |", - "| 22 | b | 22 |", - "| 44 | d | 44 |", - "+-------+---------+-------+", - ]; - - let actual = execute_to_batches(&ctx, sql).await; - assert_batches_eq!(expected, &actual); - Ok(()) -} - -#[tokio::test] -async fn right_join_not_null_filter_on_join_column() -> Result<()> { - let ctx = create_join_context_with_nulls()?; - let sql = "SELECT t1_id, t1_name, t2_id FROM t1 RIGHT JOIN t2 ON t1_id = t2_id WHERE t1_id IS NOT NULL ORDER BY t2_id"; - let expected = vec![ - "+-------+---------+-------+", - "| t1_id | t1_name | t2_id |", - "+-------+---------+-------+", - "| 11 | a | 11 |", - "| 22 | b | 22 |", - "| 44 | d | 44 |", - "| 99 | | 99 |", - "+-------+---------+-------+", - ]; - - let actual = execute_to_batches(&ctx, sql).await; - assert_batches_eq!(expected, &actual); - Ok(()) -} - -#[tokio::test] -async fn full_join_null_filter() -> Result<()> { - let ctx = create_join_context_with_nulls()?; - let sql = "SELECT t1_id, t1_name, t2_id FROM t1 FULL OUTER JOIN t2 ON t1_id = t2_id WHERE t1_name IS NULL ORDER BY t1_id"; - let expected = vec![ - "+-------+---------+-------+", - "| t1_id | t1_name | t2_id |", - "+-------+---------+-------+", - "| 88 | | |", - "| 99 | | 99 |", - "| | | 55 |", - "+-------+---------+-------+", - ]; - - let actual = execute_to_batches(&ctx, sql).await; - assert_batches_eq!(expected, &actual); - Ok(()) -} - -#[tokio::test] -async fn full_join_not_null_filter() -> Result<()> { - let ctx = create_join_context_with_nulls()?; - let sql = "SELECT t1_id, t1_name, t2_id FROM t1 FULL OUTER JOIN t2 ON t1_id = t2_id WHERE t1_name IS NOT NULL ORDER BY t1_id"; - let expected = vec![ - "+-------+---------+-------+", - "| t1_id | t1_name | t2_id |", - "+-------+---------+-------+", - "| 11 | a | 11 |", - "| 22 | b | 22 |", - "| 33 | c | |", - "| 44 | d | 44 |", - "| 77 | e | |", - "+-------+---------+-------+", - ]; - - let actual = execute_to_batches(&ctx, sql).await; - assert_batches_eq!(expected, &actual); - Ok(()) -} - -#[tokio::test] -async fn left_join_using() -> Result<()> { - let test_repartition_joins = vec![true, false]; - for repartition_joins in test_repartition_joins { - let ctx = create_join_context("id", "id", repartition_joins)?; - let sql = - "SELECT id, t1_name, t2_name FROM t1 LEFT JOIN t2 USING (id) ORDER BY id"; - let actual = execute_to_batches(&ctx, sql).await; - let expected = vec![ - "+----+---------+---------+", - "| id | t1_name | t2_name |", - "+----+---------+---------+", - "| 11 | a | z |", - "| 22 | b | y |", - "| 33 | c | |", - "| 44 | d | x |", - "+----+---------+---------+", - ]; - assert_batches_eq!(expected, &actual); - } - Ok(()) -} - -#[tokio::test] -async fn cross_join() { - let test_repartition_joins = vec![true, false]; - for repartition_joins in test_repartition_joins { - let ctx = create_join_context("t1_id", "t2_id", repartition_joins).unwrap(); - - let sql = "SELECT t1_id, t1_name, t2_name FROM t1, t2 ORDER BY t1_id"; - let actual = execute(&ctx, sql).await; - - assert_eq!(4 * 4, actual.len()); - - let sql = "SELECT t1_id, t1_name, t2_name FROM t1, t2 WHERE 1=1 ORDER BY t1_id"; - let actual = execute(&ctx, sql).await; - - assert_eq!(4 * 4, actual.len()); - - let sql = "SELECT t1_id, t1_name, t2_name FROM t1 CROSS JOIN t2 ORDER BY t1_id"; - - let actual = execute(&ctx, sql).await; - assert_eq!(4 * 4, actual.len()); - - let actual = execute_to_batches(&ctx, sql).await; - let expected = vec![ - "+-------+---------+---------+", - "| t1_id | t1_name | t2_name |", - "+-------+---------+---------+", - "| 11 | a | z |", - "| 11 | a | y |", - "| 11 | a | x |", - "| 11 | a | w |", - "| 22 | b | z |", - "| 22 | b | y |", - "| 22 | b | x |", - "| 22 | b | w |", - "| 33 | c | z |", - "| 33 | c | y |", - "| 33 | c | x |", - "| 33 | c | w |", - "| 44 | d | z |", - "| 44 | d | y |", - "| 44 | d | x |", - "| 44 | d | w |", - "+-------+---------+---------+", - ]; - - assert_batches_eq!(expected, &actual); - - // Two partitions (from UNION) on the left - let sql = "SELECT * FROM (SELECT t1_id, t1_name FROM t1 UNION ALL SELECT t1_id, t1_name FROM t1) AS t1 CROSS JOIN t2"; - let actual = execute(&ctx, sql).await; - - assert_eq!(4 * 4 * 2, actual.len()); - - // Two partitions (from UNION) on the right - let sql = "SELECT t1_id, t1_name, t2_name FROM t1 CROSS JOIN (SELECT t2_name FROM t2 UNION ALL SELECT t2_name FROM t2) AS t2"; - let actual = execute(&ctx, sql).await; - - assert_eq!(4 * 4 * 2, actual.len()); - } -} - -#[tokio::test] -async fn cross_join_unbalanced() { - // the t1_id is larger than t2_id so the join_selection optimizer should kick in - let ctx = create_join_context_unbalanced("t1_id", "t2_id").unwrap(); - - // the order of the values is not determinisitic, so we need to sort to check the values - let sql = - "SELECT t1_id, t1_name, t2_name FROM t1 CROSS JOIN t2 ORDER BY t1_id, t1_name, t2_name"; - let actual = execute_to_batches(&ctx, sql).await; - let expected = vec![ - "+-------+---------+---------+", - "| t1_id | t1_name | t2_name |", - "+-------+---------+---------+", - "| 11 | a | w |", - "| 11 | a | x |", - "| 11 | a | y |", - "| 11 | a | z |", - "| 22 | b | w |", - "| 22 | b | x |", - "| 22 | b | y |", - "| 22 | b | z |", - "| 33 | c | w |", - "| 33 | c | x |", - "| 33 | c | y |", - "| 33 | c | z |", - "| 44 | d | w |", - "| 44 | d | x |", - "| 44 | d | y |", - "| 44 | d | z |", - "| 77 | e | w |", - "| 77 | e | x |", - "| 77 | e | y |", - "| 77 | e | z |", - "+-------+---------+---------+", - ]; - assert_batches_eq!(expected, &actual); -} - #[tokio::test] async fn test_join_timestamp() -> Result<()> { let ctx = SessionContext::new(); @@ -595,35 +209,6 @@ async fn nestedjoin_without_alias() -> Result<()> { Ok(()) } -#[tokio::test] -async fn issue_3002() -> Result<()> { - // repro case for https://github.com/apache/arrow-datafusion/issues/3002 - let sql = "select a.a, b.b from a join b on a.a = b.b"; - let expected = vec!["++", "++"]; - let ctx = create_join_context_qualified("a", "b")?; - let actual = execute_to_batches(&ctx, sql).await; - assert_batches_eq!(expected, &actual); - Ok(()) -} - -#[tokio::test] -async fn inner_join_nulls() { - let sql = "SELECT * FROM (SELECT null AS id1) t1 - INNER JOIN (SELECT null AS id2) t2 ON id1 = id2"; - - #[rustfmt::skip] - let expected = vec![ - "++", - "++", - ]; - - let ctx = create_join_context_qualified("t1", "t2").unwrap(); - let actual = execute_to_batches(&ctx, sql).await; - - // left and right shouldn't match anything - assert_batches_eq!(expected, &actual); -} - #[tokio::test] async fn join_tables_with_duplicated_column_name_not_in_on_constraint() -> Result<()> { let ctx = SessionContext::new(); diff --git a/datafusion/core/tests/sql/mod.rs b/datafusion/core/tests/sql/mod.rs index bd9e213dfe216..7041cf7a1f5a7 100644 --- a/datafusion/core/tests/sql/mod.rs +++ b/datafusion/core/tests/sql/mod.rs @@ -520,100 +520,6 @@ fn create_hashjoin_datatype_context() -> Result { Ok(ctx) } -/// the table column_left has more rows than the table column_right -fn create_join_context_unbalanced( - column_left: &str, - column_right: &str, -) -> Result { - let ctx = SessionContext::new(); - - let t1_schema = Arc::new(Schema::new(vec![ - Field::new(column_left, DataType::UInt32, true), - Field::new("t1_name", DataType::Utf8, true), - ])); - let t1_data = RecordBatch::try_new( - t1_schema, - vec![ - Arc::new(UInt32Array::from_slice([11, 22, 33, 44, 77])), - Arc::new(StringArray::from(vec![ - Some("a"), - Some("b"), - Some("c"), - Some("d"), - Some("e"), - ])), - ], - )?; - ctx.register_batch("t1", t1_data)?; - - let t2_schema = Arc::new(Schema::new(vec![ - Field::new(column_right, DataType::UInt32, true), - Field::new("t2_name", DataType::Utf8, true), - ])); - let t2_data = RecordBatch::try_new( - t2_schema, - vec![ - Arc::new(UInt32Array::from_slice([11, 22, 44, 55])), - Arc::new(StringArray::from(vec![ - Some("z"), - Some("y"), - Some("x"), - Some("w"), - ])), - ], - )?; - ctx.register_batch("t2", t2_data)?; - - Ok(ctx) -} - -// Create memory tables with nulls -fn create_join_context_with_nulls() -> Result { - let ctx = SessionContext::new(); - - let t1_schema = Arc::new(Schema::new(vec![ - Field::new("t1_id", DataType::UInt32, true), - Field::new("t1_name", DataType::Utf8, true), - ])); - let t1_data = RecordBatch::try_new( - t1_schema, - vec![ - Arc::new(UInt32Array::from(vec![11, 22, 33, 44, 77, 88, 99])), - Arc::new(StringArray::from(vec![ - Some("a"), - Some("b"), - Some("c"), - Some("d"), - Some("e"), - None, - None, - ])), - ], - )?; - ctx.register_batch("t1", t1_data)?; - - let t2_schema = Arc::new(Schema::new(vec![ - Field::new("t2_id", DataType::UInt32, true), - Field::new("t2_name", DataType::Utf8, true), - ])); - let t2_data = RecordBatch::try_new( - t2_schema, - vec![ - Arc::new(UInt32Array::from(vec![11, 22, 44, 55, 99])), - Arc::new(StringArray::from(vec![ - Some("z"), - None, - Some("x"), - Some("w"), - Some("u"), - ])), - ], - )?; - ctx.register_batch("t2", t2_data)?; - - Ok(ctx) -} - fn create_sort_merge_join_context( column_left: &str, column_right: &str, diff --git a/datafusion/core/tests/sqllogictests/test_files/joins.slt b/datafusion/core/tests/sqllogictests/test_files/joins.slt new file mode 100644 index 0000000000000..eb8f72470c6cb --- /dev/null +++ b/datafusion/core/tests/sqllogictests/test_files/joins.slt @@ -0,0 +1,624 @@ +# 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. + +########## +## Joins Tests +########## + +# create table t1 +statement ok +CREATE TABLE t1(a INT, b INT, c INT) AS VALUES +(1, 10, 50), +(2, 20, 60), +(3, 30, 70), +(4, 40, 80) + +# create table t2 +statement ok +CREATE TABLE t2(a INT, b INT, c INT) AS VALUES +(1, 100, 500), +(2, 200, 600), +(9, 300, 700), +(4, 400, 800) + +# equijoin +query II nosort +SELECT t1.a, t2.b FROM t1 INNER JOIN t2 ON t1.a = t2.a ORDER BY t1.a +---- +1 100 +2 200 +4 400 + +query II nosort +SELECT t1.a, t2.b FROM t1 INNER JOIN t2 ON t2.a = t1.a ORDER BY t1.a +---- +1 100 +2 200 +4 400 + +# inner_join_nulls +query ?? +SELECT * FROM (SELECT null AS id1) t1 +INNER JOIN (SELECT null AS id2) t2 ON id1 = id2 +---- + +statement ok +DROP TABLE t1 + +statement ok +DROP TABLE t2 + + +# create table a +statement ok +CREATE TABLE a(a INT, b INT, c INT) AS VALUES +(1, 10, 50), +(2, 20, 60), +(3, 30, 70), +(4, 40, 80) + +# create table b +statement ok +CREATE TABLE b(a INT, b INT, c INT) AS VALUES +(1, 100, 500), +(2, 200, 600), +(9, 300, 700), +(4, 400, 800) + +# issue_3002 +# // repro case for https://github.com/apache/arrow-datafusion/issues/3002 + +query II +select a.a, b.b from a join b on a.a = b.b +---- + +statement ok +DROP TABLE a + +statement ok +DROP TABLE b + +# create table t1 +statement ok +CREATE TABLE t1(t1_id INT, t1_name VARCHAR) AS VALUES +(11, 'a'), +(22, 'b'), +(33, 'c'), +(44, 'd'), +(77, 'e') + +# create table t2 +statement ok +CREATE TABLE t2(t2_id INT, t2_name VARCHAR) AS VALUES +(11, 'z'), +(22, 'y'), +(44, 'x'), +(55, 'w') + +# left_join_unbalanced +# // the t1_id is larger than t2_id so the join_selection optimizer should kick in +query ITT nosort +SELECT t1_id, t1_name, t2_name FROM t1 LEFT JOIN t2 ON t1_id = t2_id ORDER BY t1_id +---- +11 a z +22 b y +33 c NULL +44 d x +77 e NULL + +query ITT nosort +SELECT t1_id, t1_name, t2_name FROM t1 LEFT JOIN t2 ON t2_id = t1_id ORDER BY t1_id +---- +11 a z +22 b y +33 c NULL +44 d x +77 e NULL + + +# cross_join_unbalanced +query ITT nosort +SELECT t1_id, t1_name, t2_name FROM t1 CROSS JOIN t2 ORDER BY t1_id, t1_name, t2_name +---- +11 a w +11 a x +11 a y +11 a z +22 b w +22 b x +22 b y +22 b z +33 c w +33 c x +33 c y +33 c z +44 d w +44 d x +44 d y +44 d z +77 e w +77 e x +77 e y +77 e z + +statement ok +DROP TABLE t1 + +statement ok +DROP TABLE t2 + +# create table t1 +statement ok +CREATE TABLE t1(t1_id INT, t1_name VARCHAR) AS VALUES +(11, 'a'), +(22, 'b'), +(33, 'c'), +(44, 'd'), +(77, 'e'), +(88, NULL), +(99, NULL) + +# create table t2 +statement ok +CREATE TABLE t2(t2_id INT, t2_name VARCHAR) AS VALUES +(11, 'z'), +(22, NULL), +(44, 'x'), +(55, 'w'), +(99, 'u') + +# left_join_null_filter +# // Since t2 is the non-preserved side of the join, we cannot push down a NULL filter. +# // Note that this is only true because IS NULL does not remove nulls. For filters that +# // remove nulls, we can rewrite the join as an inner join and then push down the filter. +query IIT nosort +SELECT t1_id, t2_id, t2_name FROM t1 LEFT JOIN t2 ON t1_id = t2_id WHERE t2_name IS NULL ORDER BY t1_id +---- +22 22 NULL +33 NULL NULL +77 NULL NULL +88 NULL NULL + +# left_join_null_filter_on_join_column +# // Again, since t2 is the non-preserved side of the join, we cannot push down a NULL filter. +query IIT nosort +SELECT t1_id, t2_id, t2_name FROM t1 LEFT JOIN t2 ON t1_id = t2_id WHERE t2_id IS NULL ORDER BY t1_id +---- +33 NULL NULL +77 NULL NULL +88 NULL NULL + +# left_join_not_null_filter +query IIT nosort +SELECT t1_id, t2_id, t2_name FROM t1 LEFT JOIN t2 ON t1_id = t2_id WHERE t2_name IS NOT NULL ORDER BY t1_id +---- +11 11 z +44 44 x +99 99 u + +# left_join_not_null_filter_on_join_column +query IIT nosort +SELECT t1_id, t2_id, t2_name FROM t1 LEFT JOIN t2 ON t1_id = t2_id WHERE t2_id IS NOT NULL ORDER BY t1_id +---- +11 11 z +22 22 NULL +44 44 x +99 99 u + +# self_join_non_equijoin +query II nosort +SELECT x.t1_id, y.t1_id FROM t1 x JOIN t1 y ON x.t1_id = 11 AND y.t1_id = 44 +---- +11 44 + +# right_join_null_filter +query ITI nosort +SELECT t1_id, t1_name, t2_id FROM t1 RIGHT JOIN t2 ON t1_id = t2_id WHERE t1_name IS NULL ORDER BY t2_id +---- +NULL NULL 55 +99 NULL 99 + +# right_join_null_filter_on_join_column +query ITI nosort +SELECT t1_id, t1_name, t2_id FROM t1 RIGHT JOIN t2 ON t1_id = t2_id WHERE t1_id IS NULL ORDER BY t2_id +---- +NULL NULL 55 + +# right_join_not_null_filter +query ITI nosort +SELECT t1_id, t1_name, t2_id FROM t1 RIGHT JOIN t2 ON t1_id = t2_id WHERE t1_name IS NOT NULL ORDER BY t2_id +---- +11 a 11 +22 b 22 +44 d 44 + +# right_join_not_null_filter_on_join_column +query ITI nosort +SELECT t1_id, t1_name, t2_id FROM t1 RIGHT JOIN t2 ON t1_id = t2_id WHERE t1_id IS NOT NULL ORDER BY t2_id +---- +11 a 11 +22 b 22 +44 d 44 +99 NULL 99 + +# full_join_null_filter +query ITI nosort +SELECT t1_id, t1_name, t2_id FROM t1 FULL OUTER JOIN t2 ON t1_id = t2_id WHERE t1_name IS NULL ORDER BY t1_id +---- +88 NULL NULL +99 NULL 99 +NULL NULL 55 + +# full_join_not_null_filter +query ITI nosort +SELECT t1_id, t1_name, t2_id FROM t1 FULL OUTER JOIN t2 ON t1_id = t2_id WHERE t1_name IS NOT NULL ORDER BY t1_id +---- +11 a 11 +22 b 22 +33 c NULL +44 d 44 +77 e NULL + +statement ok +DROP TABLE t1 + +statement ok +DROP TABLE t2 + +# create table t1 +statement ok +CREATE TABLE t1(id INT, t1_name VARCHAR, t1_int INT) AS VALUES +(11, 'a', 1), +(22, 'b', 2), +(33, 'c', 3), +(44, 'd', 4) + +# create table t2 +statement ok +CREATE TABLE t2(id INT, t2_name VARCHAR, t2_int INT) AS VALUES +(11, 'z', 3), +(22, 'y', 1), +(44, 'x', 3), +(55, 'w', 3) + +# left_join_using + +# set repartition_joins to true +statement ok +set datafusion.optimizer.repartition_joins = true + +query ITT nosort +SELECT id, t1_name, t2_name FROM t1 LEFT JOIN t2 USING (id) ORDER BY id +---- +11 a z +22 b y +33 c NULL +44 d x + +# set repartition_joins to false +statement ok +set datafusion.optimizer.repartition_joins = false + +query ITT nosort +SELECT id, t1_name, t2_name FROM t1 LEFT JOIN t2 USING (id) ORDER BY id +---- +11 a z +22 b y +33 c NULL +44 d x + +statement ok +DROP TABLE t1 + +statement ok +DROP TABLE t2 + +# create table t1 +statement ok +CREATE TABLE t1(t1_id INT, t1_name VARCHAR, t1_int INT) AS VALUES +(11, 'a', 1), +(22, 'b', 2), +(33, 'c', 3), +(44, 'd', 4) + +# create table t2 +statement ok +CREATE TABLE t2(t2_id INT, t2_name VARCHAR, t2_int INT) AS VALUES +(11, 'z', 3), +(22, 'y', 1), +(44, 'x', 3), +(55, 'w', 3) + +# cross_join + +# set repartition_joins to true +statement ok +set datafusion.optimizer.repartition_joins = true + +query ITT nosort +SELECT t1_id, t1_name, t2_name FROM t1, t2 ORDER BY t1_id +---- +11 a z +11 a y +11 a x +11 a w +22 b z +22 b y +22 b x +22 b w +33 c z +33 c y +33 c x +33 c w +44 d z +44 d y +44 d x +44 d w + +query ITT nosort +SELECT t1_id, t1_name, t2_name FROM t1, t2 WHERE 1=1 ORDER BY t1_id +---- +11 a z +11 a y +11 a x +11 a w +22 b z +22 b y +22 b x +22 b w +33 c z +33 c y +33 c x +33 c w +44 d z +44 d y +44 d x +44 d w + +query ITT nosort +SELECT t1_id, t1_name, t2_name FROM t1 CROSS JOIN t2 ORDER BY t1_id +---- +11 a z +11 a y +11 a x +11 a w +22 b z +22 b y +22 b x +22 b w +33 c z +33 c y +33 c x +33 c w +44 d z +44 d y +44 d x +44 d w + +query ITITI rowsort +SELECT * FROM (SELECT t1_id, t1_name FROM t1 UNION ALL SELECT t1_id, t1_name FROM t1) AS t1 CROSS JOIN t2 +---- +11 a 11 z 3 +11 a 11 z 3 +11 a 22 y 1 +11 a 22 y 1 +11 a 44 x 3 +11 a 44 x 3 +11 a 55 w 3 +11 a 55 w 3 +22 b 11 z 3 +22 b 11 z 3 +22 b 22 y 1 +22 b 22 y 1 +22 b 44 x 3 +22 b 44 x 3 +22 b 55 w 3 +22 b 55 w 3 +33 c 11 z 3 +33 c 11 z 3 +33 c 22 y 1 +33 c 22 y 1 +33 c 44 x 3 +33 c 44 x 3 +33 c 55 w 3 +33 c 55 w 3 +44 d 11 z 3 +44 d 11 z 3 +44 d 22 y 1 +44 d 22 y 1 +44 d 44 x 3 +44 d 44 x 3 +44 d 55 w 3 +44 d 55 w 3 + +query ITT rowsort +SELECT t1_id, t1_name, t2_name FROM t1 CROSS JOIN (SELECT t2_name FROM t2 UNION ALL SELECT t2_name FROM t2) AS t2_data +---- +11 a w +11 a w +11 a x +11 a x +11 a y +11 a y +11 a z +11 a z +22 b w +22 b w +22 b x +22 b x +22 b y +22 b y +22 b z +22 b z +33 c w +33 c w +33 c x +33 c x +33 c y +33 c y +33 c z +33 c z +44 d w +44 d w +44 d x +44 d x +44 d y +44 d y +44 d z +44 d z + +# set repartition_joins to true +statement ok +set datafusion.optimizer.repartition_joins = false + +query ITT nosort +SELECT t1_id, t1_name, t2_name FROM t1, t2 ORDER BY t1_id +---- +11 a z +11 a y +11 a x +11 a w +22 b z +22 b y +22 b x +22 b w +33 c z +33 c y +33 c x +33 c w +44 d z +44 d y +44 d x +44 d w + +query ITT nosort +SELECT t1_id, t1_name, t2_name FROM t1, t2 WHERE 1=1 ORDER BY t1_id +---- +11 a z +11 a y +11 a x +11 a w +22 b z +22 b y +22 b x +22 b w +33 c z +33 c y +33 c x +33 c w +44 d z +44 d y +44 d x +44 d w + +query ITT nosort +SELECT t1_id, t1_name, t2_name FROM t1 CROSS JOIN t2 ORDER BY t1_id +---- +11 a z +11 a y +11 a x +11 a w +22 b z +22 b y +22 b x +22 b w +33 c z +33 c y +33 c x +33 c w +44 d z +44 d y +44 d x +44 d w + +query ITITI rowsort +SELECT * FROM (SELECT t1_id, t1_name FROM t1 UNION ALL SELECT t1_id, t1_name FROM t1) AS t1 CROSS JOIN t2 +---- +11 a 11 z 3 +11 a 11 z 3 +11 a 22 y 1 +11 a 22 y 1 +11 a 44 x 3 +11 a 44 x 3 +11 a 55 w 3 +11 a 55 w 3 +22 b 11 z 3 +22 b 11 z 3 +22 b 22 y 1 +22 b 22 y 1 +22 b 44 x 3 +22 b 44 x 3 +22 b 55 w 3 +22 b 55 w 3 +33 c 11 z 3 +33 c 11 z 3 +33 c 22 y 1 +33 c 22 y 1 +33 c 44 x 3 +33 c 44 x 3 +33 c 55 w 3 +33 c 55 w 3 +44 d 11 z 3 +44 d 11 z 3 +44 d 22 y 1 +44 d 22 y 1 +44 d 44 x 3 +44 d 44 x 3 +44 d 55 w 3 +44 d 55 w 3 + +query ITT rowsort +SELECT t1_id, t1_name, t2_name FROM t1 CROSS JOIN (SELECT t2_name FROM t2 UNION ALL SELECT t2_name FROM t2) AS t2_data +---- +11 a w +11 a w +11 a x +11 a x +11 a y +11 a y +11 a z +11 a z +22 b w +22 b w +22 b x +22 b x +22 b y +22 b y +22 b z +22 b z +33 c w +33 c w +33 c x +33 c x +33 c y +33 c y +33 c z +33 c z +44 d w +44 d w +44 d x +44 d x +44 d y +44 d y +44 d z +44 d z + +statement ok +DROP TABLE t1 + +statement ok +DROP TABLE t2 diff --git a/datafusion/physical-expr/src/expressions/case.rs b/datafusion/physical-expr/src/expressions/case.rs index 2d97d57324fee..72783145a04db 100644 --- a/datafusion/physical-expr/src/expressions/case.rs +++ b/datafusion/physical-expr/src/expressions/case.rs @@ -23,7 +23,7 @@ use crate::physical_expr::down_cast_any_ref; use crate::PhysicalExpr; use arrow::array::*; use arrow::compute::kernels::zip::zip; -use arrow::compute::{and, eq_dyn, is_null, not, or, or_kleene}; +use arrow::compute::{and, and_kleene, eq_dyn, is_not_null, is_null, not, or, or_kleene}; use arrow::datatypes::{DataType, Schema}; use arrow::record_batch::RecordBatch; use datafusion_common::{cast::as_boolean_array, DataFusionError, Result}; @@ -138,6 +138,8 @@ impl CaseExpr { let when_value = when_value.into_array(batch.num_rows()); // build boolean array representing which rows match the "when" value let when_match = eq_dyn(&when_value, base_value.as_ref())?; + // Treat nulls as false + let when_match = and_kleene(&when_match, &is_not_null(&when_match)?)?; let then_value = self.when_then_expr[i] .1 @@ -152,7 +154,7 @@ impl CaseExpr { current_value = zip(&when_match, then_value.as_ref(), current_value.as_ref())?; - remainder = and(&remainder, &or_kleene(¬(&when_match)?, &base_nulls)?)?; + remainder = and(&remainder, ¬(&when_match)?)?; } if let Some(e) = &self.else_expr { @@ -526,6 +528,35 @@ mod tests { Ok(()) } + #[test] + fn case_with_expr_when_null() -> Result<()> { + let batch = case_test_batch()?; + let schema = batch.schema(); + + // CASE a WHEN NULL THEN 0 WHEN a THEN 123 ELSE 999 END + let when1 = lit(ScalarValue::Utf8(None)); + let then1 = lit(0i32); + let when2 = col("a", &schema)?; + let then2 = lit(123i32); + let else_value = lit(999i32); + + let expr = generate_case_when_with_type_coercion( + Some(col("a", &schema)?), + vec![(when1, then1), (when2, then2)], + Some(else_value), + schema.as_ref(), + )?; + let result = expr.evaluate(&batch)?.into_array(batch.num_rows()); + let result = as_int32_array(&result)?; + + let expected = + &Int32Array::from(vec![Some(123), Some(123), Some(999), Some(123)]); + + assert_eq!(expected, result); + + Ok(()) + } + #[test] fn case_without_expr_divide_by_zero() -> Result<()> { let batch = case_test_batch1()?;