jayzhan211 commented on code in PR #25434: URL: https://github.com/apache/datafusion/pull/25434#discussion_r4047001845
########## datafusion/physical-plan/src/joins/hash_join/compact_hash_map/tests.rs: ########## @@ -0,0 +1,449 @@ +// 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 super::*; + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use arrow::array::{ArrayRef, DictionaryArray, Int32Array, StringArray, StructArray}; +use arrow::compute::concat_batches; +use arrow::datatypes::{Field, Int32Type, Schema}; +use datafusion_common::cast::as_int32_array; +use datafusion_common::{DataFusionError, JoinType}; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryConsumer, MemoryPool}; +use datafusion_execution::runtime_env::RuntimeEnvBuilder; +use datafusion_expr::{Volatility, create_udf}; +use datafusion_physical_expr::ScalarFunctionExpr; + +use super::super::exec::HASH_JOIN_SEED; +use crate::ExecutionPlan; +use crate::common; +use crate::joins::PartitionMode; +use crate::joins::hash_join::HashJoinExec; +use crate::joins::join_hash_map::{JoinHashMapType, JoinHashMapU32}; +use crate::joins::utils::update_hash; +use crate::test::TestMemoryExec; + +#[tokio::test] +async fn compact_hash_build_with_duplicates_and_nulls() -> Result<()> { Review Comment: No case with `HASH_BUILD_CHUNK_ROWS < distinct << rows`, so nothing pins table size in the mid-cardinality range. Add one that bounds the final allocation. ```rs let rows = 1_000_000; let distinct = 10_000; let batches: Vec<_> = (0..rows) .step_by(HASH_BUILD_CHUNK_ROWS) .map(|start| { let end = (start + HASH_BUILD_CHUNK_ROWS).min(rows); RecordBatch::try_from_iter([( "key", Arc::new(Int32Array::from_iter_values( (start..end).map(|i| (i % distinct) as i32), )) as ArrayRef, )]) }) .collect::<Result<_, _>>()?; let on = vec![Arc::new(Column::new("key", 0)) as PhysicalExprRef]; let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(1 << 30)); let reservation = MemoryConsumer::new("mid cardinality").register(&pool); let (table, _next) = build_compact_hash_map::<u32>( &batches, &on, rows, HASH_JOIN_SEED.random_state(), NullEquality::NullEqualsNothing, &reservation, &mut 0, )?; assert!( table.allocation_size() <= estimate_memory_size::<(u64, u32)>(4 * distinct, size_of::<JoinHashMapU32>())? ); ``` ########## datafusion/physical-plan/src/joins/hash_join/compact_hash_map.rs: ########## @@ -0,0 +1,195 @@ +// 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. + +//! Adapt hash lookup capacity to the build keys while retaining every row index. + +use std::fmt; +use std::mem::size_of; + +use arrow::datatypes::DataType; +use arrow::record_batch::RecordBatch; +use datafusion_common::hash_utils::{RandomState, create_hashes}; +use datafusion_common::utils::memory::estimate_memory_size; +use datafusion_common::{NullEquality, Result}; +use datafusion_execution::memory_pool::MemoryReservation; +use datafusion_physical_expr::PhysicalExprRef; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; +use hashbrown::HashTable; + +use crate::joins::join_hash_map::update_from_iter; +use crate::joins::utils::matchable_join_keys; + +/// Bound hash scratch and the number of potentially new hashes per insertion. +const HASH_BUILD_CHUNK_ROWS: usize = 8192; + +// Only flat column arrays can be sliced without repeatedly hashing unused +// dictionary values or changing the input batch seen by computed expressions. +fn is_flat_join_type(data_type: &DataType) -> bool { + data_type.is_primitive() + || matches!( + data_type, + DataType::Null + | DataType::Boolean + | DataType::FixedSizeBinary(_) + | DataType::Utf8 + | DataType::LargeUtf8 + | DataType::Binary + | DataType::LargeBinary + | DataType::Utf8View + | DataType::BinaryView + ) +} + +// Start with one chunk of buckets so low-cardinality builds stay small. On +// growth, try the row-count capacity once to avoid repeated rehashing of unique +// keys. If it cannot fit, continue growing by observed hashes plus the next chunk. +#[expect(clippy::type_complexity)] +pub(super) fn build_compact_hash_map<T>( + batches: &[RecordBatch], + on: &[PhysicalExprRef], + num_rows: usize, + random_state: &RandomState, + null_equality: NullEquality, + reservation: &MemoryReservation, + peak: &mut usize, +) -> Result<(HashTable<(u64, T)>, Vec<T>)> +where + T: Copy + Default + TryFrom<usize> + PartialOrd, + <T as TryFrom<usize>>::Error: fmt::Debug, +{ + let initial_reserved = reservation.size(); + let fixed_bytes = size_of::<HashTable<(u64, T)>>() + size_of::<Vec<T>>(); + let chain_bytes = num_rows.checked_mul(size_of::<T>()).ok_or_else(|| { + datafusion_common::exec_datafusion_err!("Hash join row-index size overflow") + })?; + reservation.try_grow(fixed_bytes + chain_bytes)?; + *peak = (*peak).max(reservation.size() - initial_reserved); + let mut next = vec![T::default(); num_rows]; + let mut table = HashTable::new(); + let scratch = reservation.new_empty(); + let max_batch_rows = batches.iter().map(RecordBatch::num_rows).max().unwrap_or(0); + let chunk_rows = max_batch_rows.min(HASH_BUILD_CHUNK_ROWS); + let mut hash_chunk_rows = chunk_rows; + if let Some(batch) = batches.first() { + for expr in on { + if !expr.is::<Column>() + || !is_flat_join_type(&expr.data_type(&batch.schema())?) + { + // Slicing a dictionary retains all its values. Hash complex keys + // and evaluate computed keys once per original batch. + hash_chunk_rows = max_batch_rows; + break; + } + } + } + scratch.try_grow(hash_chunk_rows * size_of::<u64>())?; + let mut hashes = vec![0; hash_chunk_rows]; + *peak = (*peak).max(reservation.size() - initial_reserved + scratch.size()); + let mut tried_preallocation = false; + 'build: loop { + let mut offset = 0; + for batch in batches.iter().rev() { + for hash_start in (0..batch.num_rows()).step_by(hash_chunk_rows.max(1)).rev() + { + let hash_rows = (batch.num_rows() - hash_start).min(hash_chunk_rows); + let chunk = (hash_rows < batch.num_rows()) + .then(|| batch.slice(hash_start, hash_rows)); + let keys = + evaluate_expressions_to_arrays(on, chunk.as_ref().unwrap_or(batch))?; + hashes[..hash_rows].fill(0); + let hashes = + create_hashes(&keys, random_state, &mut hashes[..hash_rows])?; + let valid = matchable_join_keys(&keys, null_equality); + for start in (0..hash_rows).step_by(HASH_BUILD_CHUNK_ROWS).rev() { + let rows = (hash_rows - start).min(HASH_BUILD_CHUNK_ROWS); + let hashes = &hashes[start..start + rows]; + let valid = valid.as_ref().map(|valid| valid.slice(start, rows)); + let additional = rows - valid.as_ref().map_or(0, |n| n.null_count()); + if additional > table.capacity() - table.len() { + // Keep the old allocation charged during rehash. The fixed-size + // allowance covers control-group padding; at least eight elements + // also covers hashbrown's minimum bucket sizes. + let minimum = (table.len() + additional).max(chunk_rows); Review Comment: Second growth always jumps to `num_rows`, and it triggers once `len + 8192 > 14336`, so any build with more than ~6,144 distinct hashes gets the same table as `main`. Measured with 1M Utf8 rows: 6,000 distinct → 278,536 B table; 6,200 distinct → 35,651,592 B; 10K and 100K distinct → 35,651,592 B. Cheap fix: shrink after the build. With this, 6,200 → 139 KB, 10K → 278 KB, 100K → 2.2 MB, ≥500K unchanged, and all 6 new tests still pass. Peak is still row-count-sized when the pool admits it ```diff break; } + // The one-shot row-count preallocation overshoots when only a fraction + // of the rows carry distinct hashes; return the unused buckets. + if table.capacity() / 4 > table.len() { + let old_bytes = table.allocation_size(); + table.shrink_to(table.len(), |&(hash, _)| hash); + reservation.shrink(old_bytes - table.allocation_size()); + } drop(hashes); ``` -- 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]
