diff --git a/Cargo.lock b/Cargo.lock index baf18148735d8..87aefd6fdae9e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1223,6 +1223,7 @@ dependencies = [ "once_cell", "pprof", "rand 0.8.5", + "semver 1.0.14", "serde", "tikv-jemalloc-ctl", "tikv-jemalloc-sys", @@ -1521,10 +1522,14 @@ name = "common-hashtable" version = "0.1.0" dependencies = [ "ahash 0.8.0", + "bumpalo", + "cfg-if", "common-base", + "libc", "ordered-float 3.3.0", "primitive-types", "rand 0.8.5", + "semver 1.0.14", ] [[package]] diff --git a/src/common/base/Cargo.toml b/src/common/base/Cargo.toml index 6363f63b00e38..d0d91203b2b93 100644 --- a/src/common/base/Cargo.toml +++ b/src/common/base/Cargo.toml @@ -37,6 +37,7 @@ pprof = { version = "0.10.1", features = [ "protobuf-codec", "protobuf", ] } +semver = "1.0.10" serde = { workspace = true } tikv-jemalloc-ctl = { version = "0.4.2", optional = true } tikv-jemalloc-sys = "0.4.3" diff --git a/src/common/base/src/base/runtime_tracker.rs b/src/common/base/src/base/runtime_tracker.rs index a7d7ec25bacdc..7e607a4292d67 100644 --- a/src/common/base/src/base/runtime_tracker.rs +++ b/src/common/base/src/base/runtime_tracker.rs @@ -18,7 +18,7 @@ use std::sync::atomic::AtomicI64; use std::sync::atomic::Ordering; use std::sync::Arc; -use crate::mem_allocator::ALLOC; +use crate::mem_allocator::GlobalAllocator; #[thread_local] static mut TRACKER: *mut ThreadTracker = std::ptr::null_mut(); @@ -92,12 +92,15 @@ impl ThreadTracker { } #[inline] - pub fn realloc_memory(old_size: i64, new_size: i64) { - let addition = new_size - old_size; - match addition > 0 { - true => Self::alloc_memory(addition), - false => Self::dealloc_memory(-addition), - } + pub fn grow_memory(old_size: i64, new_size: i64) { + assert!(old_size <= new_size); + Self::alloc_memory(new_size - old_size) + } + + #[inline] + pub fn shrink_memory(old_size: i64, new_size: i64) { + assert!(old_size >= new_size); + Self::dealloc_memory(old_size - new_size) } } @@ -171,7 +174,7 @@ impl RuntimeTracker { let tracker = std::mem::replace(&mut TRACKER, std::ptr::null_mut()); std::ptr::drop_in_place(tracker as usize as *mut ThreadTracker); - ALLOC.dealloc(tracker as *mut u8, Layout::new::()) + GlobalAllocator.dealloc(tracker as *mut u8, Layout::new::()) } } diff --git a/src/common/base/src/lib.rs b/src/common/base/src/lib.rs index 6d6c33c8ca3ff..f2707ca7d7ece 100644 --- a/src/common/base/src/lib.rs +++ b/src/common/base/src/lib.rs @@ -12,7 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +#![feature(allocator_api)] #![feature(thread_local)] +#![feature(ptr_metadata)] #![feature(result_flattening)] #![feature(try_trait_v2)] #![allow(incomplete_features)] diff --git a/src/common/base/src/mem_allocator/global_allocator.rs b/src/common/base/src/mem_allocator/global_allocator.rs new file mode 100644 index 0000000000000..01eab7ce51c59 --- /dev/null +++ b/src/common/base/src/mem_allocator/global_allocator.rs @@ -0,0 +1,128 @@ +// Copyright 2022 Datafuse Labs. +// +// Licensed 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 std::alloc::AllocError; +use std::alloc::Allocator; +use std::alloc::GlobalAlloc; +use std::alloc::Layout; +use std::ptr::null_mut; +use std::ptr::NonNull; + +use super::je_allocator::JEAllocator; +use super::system_allocator::SystemAllocator; + +#[global_allocator] +pub static GLOBAL_ALLOCATOR: GlobalAllocator = GlobalAllocator; + +#[derive(Debug, Clone, Copy, Default)] +pub struct GlobalAllocator; + +type Fallback = JEAllocator; + +unsafe impl Allocator for GlobalAllocator { + #[inline(always)] + fn allocate(&self, layout: Layout) -> Result, AllocError> { + Fallback::new(Default::default()).allocate(layout) + } + + #[inline(always)] + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + Fallback::new(Default::default()).deallocate(ptr, layout) + } + + #[inline(always)] + fn allocate_zeroed(&self, layout: Layout) -> Result, AllocError> { + Fallback::new(Default::default()).allocate_zeroed(layout) + } + + #[inline(always)] + unsafe fn grow( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + Fallback::new(Default::default()).grow(ptr, old_layout, new_layout) + } + + #[inline(always)] + unsafe fn grow_zeroed( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + Fallback::new(Default::default()).grow_zeroed(ptr, old_layout, new_layout) + } + + #[inline(always)] + unsafe fn shrink( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + Fallback::new(Default::default()).shrink(ptr, old_layout, new_layout) + } +} + +unsafe impl GlobalAlloc for GlobalAllocator { + #[inline] + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + if let Ok(ptr) = GlobalAllocator.allocate(layout) { + ptr.as_ptr() as *mut u8 + } else { + null_mut() + } + } + + #[inline] + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + let ptr = NonNull::new(ptr).unwrap_unchecked(); + GlobalAllocator.deallocate(ptr, layout); + } + + #[inline] + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + if let Ok(ptr) = GlobalAllocator.allocate_zeroed(layout) { + ptr.as_ptr() as *mut u8 + } else { + null_mut() + } + } + + #[inline] + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + use std::cmp::Ordering::*; + let ptr = NonNull::new(ptr).unwrap_unchecked(); + let new_layout = Layout::from_size_align(new_size, layout.align()).unwrap(); + match layout.size().cmp(&new_size) { + Less => { + if let Ok(ptr) = GlobalAllocator.grow(ptr, layout, new_layout) { + ptr.as_ptr() as *mut u8 + } else { + null_mut() + } + } + Greater => { + if let Ok(ptr) = GlobalAllocator.shrink(ptr, layout, new_layout) { + ptr.as_ptr() as *mut u8 + } else { + null_mut() + } + } + Equal => ptr.as_ptr() as *mut u8, + } + } +} diff --git a/src/common/base/src/mem_allocator/je_allocator.rs b/src/common/base/src/mem_allocator/je_allocator.rs index 34af3d11f2e02..2b69812ebf9f8 100644 --- a/src/common/base/src/mem_allocator/je_allocator.rs +++ b/src/common/base/src/mem_allocator/je_allocator.rs @@ -16,41 +16,43 @@ // License, v. 2.0. If a copy of the MPL was not distributed with this // file, You can obtain one at https://mozilla.org/MPL/2.0/. -#[global_allocator] -pub static ALLOC: JEAllocator = JEAllocator; +#[derive(Debug, Clone, Copy, Default)] +pub struct JEAllocator { + #[allow(dead_code)] + allocator: T, +} -pub use platform::*; -pub use tikv_jemalloc_sys; +impl JEAllocator { + pub fn new(allocator: T) -> Self { + Self { allocator } + } +} -mod platform { - use std::alloc::GlobalAlloc; +#[cfg(any(target_os = "linux", target_os = "macos"))] +pub mod linux_or_macos { + use std::alloc::AllocError; + use std::alloc::Allocator; use std::alloc::Layout; - use std::os::raw::c_int; + use std::ptr::NonNull; + use libc::c_int; + use libc::c_void; use tikv_jemalloc_sys as ffi; + use super::JEAllocator; use crate::base::ThreadTracker; - use crate::mem_allocator::Allocator; - /// Memory allocation APIs compatible with libc - pub mod libc_compat { - pub use super::ffi::free; - pub use super::ffi::malloc; - pub use super::ffi::realloc; + impl JEAllocator { + pub const FALLBACK: bool = false; } - #[derive(Debug, Clone, Copy, Default)] - pub struct JEAllocator; - - // The minimum alignment guaranteed by the architecture. This value is used to - // add fast paths for low alignment values. #[cfg(all(any( target_arch = "arm", target_arch = "mips", target_arch = "mipsel", target_arch = "powerpc" )))] - const MIN_ALIGN: usize = 8; + const ALIGNOF_MAX_ALIGN_T: usize = 8; #[cfg(all(any( target_arch = "x86", target_arch = "x86_64", @@ -58,124 +60,204 @@ mod platform { target_arch = "powerpc64", target_arch = "powerpc64le", target_arch = "mips64", + target_arch = "riscv64", target_arch = "s390x", target_arch = "sparc64" )))] - const MIN_ALIGN: usize = 16; + const ALIGNOF_MAX_ALIGN_T: usize = 16; + /// If `align` is less than `_Alignof(max_align_t)`, and if the requested + /// allocation `size` is larger than the alignment, we are guaranteed to get a + /// suitably aligned allocation by default, without passing extra flags, and + /// this function returns `0`. + /// + /// Otherwise, it returns the alignment flag to pass to the jemalloc APIs. fn layout_to_flags(align: usize, size: usize) -> c_int { - // If our alignment is less than the minimum alignment they we may not - // have to pass special flags asking for a higher alignment. If the - // alignment is greater than the size, however, then this hits a sort of odd - // case where we still need to ask for a custom alignment. See #25 for more - // info. - if align <= MIN_ALIGN && align <= size { + if align <= ALIGNOF_MAX_ALIGN_T && align <= size { 0 } else { - // Equivalent to the MALLOCX_ALIGN(a) macro. - align.trailing_zeros() as _ + ffi::MALLOCX_ALIGN(align) } } - unsafe impl GlobalAlloc for JEAllocator { - #[inline] - unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + unsafe impl Allocator for JEAllocator { + #[inline(always)] + fn allocate(&self, layout: Layout) -> Result, AllocError> { + let data_address = if layout.size() == 0 { + unsafe { NonNull::new(layout.align() as *mut ()).unwrap_unchecked() } + } else { + let flags = layout_to_flags(layout.align(), layout.size()); + unsafe { + NonNull::new(ffi::mallocx(layout.size(), flags) as *mut ()).ok_or(AllocError)? + } + }; ThreadTracker::alloc_memory(layout.size() as i64); - let flags = layout_to_flags(layout.align(), layout.size()); - ffi::mallocx(layout.size(), flags) as *mut u8 + Ok(NonNull::<[u8]>::from_raw_parts(data_address, layout.size())) } - #[inline] - unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { - ThreadTracker::dealloc_memory(layout.size() as i64); - - let flags = layout_to_flags(layout.align(), layout.size()); - ffi::sdallocx(ptr as *mut _, layout.size(), flags) - } - - #[inline] - unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { - ThreadTracker::alloc_memory(layout.size() as i64); - - if layout.align() <= MIN_ALIGN && layout.align() <= layout.size() { - ffi::calloc(1, layout.size()) as *mut u8 + #[inline(always)] + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + if layout.size() == 0 { + debug_assert_eq!(ptr.as_ptr() as usize, layout.align()); } else { - let flags = layout_to_flags(layout.align(), layout.size()) | ffi::MALLOCX_ZERO; - ffi::mallocx(layout.size(), flags) as *mut u8 + let flags = layout_to_flags(layout.align(), layout.size()); + ffi::sdallocx(ptr.as_ptr() as *mut _, layout.size(), flags) } } - #[inline] - unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { - ThreadTracker::realloc_memory(layout.size() as i64, new_size as i64); - - let flags = layout_to_flags(layout.align(), new_size); - ffi::rallocx(ptr as *mut _, new_size, flags) as *mut u8 + #[inline(always)] + fn allocate_zeroed(&self, layout: Layout) -> Result, AllocError> { + let data_address = if layout.size() == 0 { + unsafe { NonNull::new(layout.align() as *mut ()).unwrap_unchecked() } + } else { + let flags = layout_to_flags(layout.align(), layout.size()) | ffi::MALLOCX_ZERO; + unsafe { + NonNull::new(ffi::mallocx(layout.size(), flags) as *mut ()).ok_or(AllocError)? + } + }; + Ok(NonNull::<[u8]>::from_raw_parts(data_address, layout.size())) } - } - unsafe impl Allocator for JEAllocator { - unsafe fn allocx(&mut self, layout: Layout, clear_mem: bool) -> *mut u8 { - if clear_mem { - self.alloc_zeroed(layout) + unsafe fn grow( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + debug_assert_eq!(old_layout.align(), new_layout.align()); + debug_assert!(old_layout.size() <= new_layout.size()); + let data_address = if new_layout.size() == 0 { + NonNull::new(new_layout.align() as *mut ()).unwrap_unchecked() + } else if old_layout.size() == 0 { + let flags = layout_to_flags(new_layout.align(), new_layout.size()); + NonNull::new(ffi::mallocx(new_layout.size(), flags) as *mut ()).ok_or(AllocError)? } else { - self.alloc(layout) - } + let flags = layout_to_flags(new_layout.align(), new_layout.size()); + NonNull::new(ffi::rallocx(ptr.cast().as_ptr(), new_layout.size(), flags) as *mut ()) + .unwrap() + }; + Ok(NonNull::<[u8]>::from_raw_parts( + data_address, + new_layout.size(), + )) } - unsafe fn deallocx(&mut self, ptr: *mut u8, layout: Layout) { - self.dealloc(ptr, layout) + unsafe fn grow_zeroed( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + debug_assert_eq!(old_layout.align(), new_layout.align()); + debug_assert!(old_layout.size() <= new_layout.size()); + let data_address = if new_layout.size() == 0 { + NonNull::new(new_layout.align() as *mut ()).unwrap_unchecked() + } else if old_layout.size() == 0 { + let flags = + layout_to_flags(new_layout.align(), new_layout.size()) | ffi::MALLOCX_ZERO; + NonNull::new(ffi::mallocx(new_layout.size(), flags) as *mut ()).ok_or(AllocError)? + } else { + let flags = layout_to_flags(new_layout.align(), new_layout.size()); + // Jemalloc doesn't support `grow_zeroed`, so it might be better to use + // mmap allocator for large frequent memory allocation and take jemalloc + // as fallback. + let raw = ffi::rallocx(ptr.cast().as_ptr(), new_layout.size(), flags) as *mut u8; + raw.add(old_layout.size()) + .write_bytes(0, new_layout.size() - old_layout.size()); + NonNull::new(raw as *mut ()).unwrap() + }; + Ok(NonNull::<[u8]>::from_raw_parts( + data_address, + new_layout.size(), + )) } - unsafe fn reallocx( - &mut self, - ptr: *mut u8, - layout: Layout, - new_size: usize, - clear_mem: bool, - ) -> *mut u8 { - ThreadTracker::realloc_memory(layout.size() as i64, new_size as i64); - - let mut flags = layout_to_flags(layout.align(), new_size); - if clear_mem { - flags |= ffi::MALLOCX_ZERO; + unsafe fn shrink( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + debug_assert_eq!(old_layout.align(), new_layout.align()); + debug_assert!(old_layout.size() >= new_layout.size()); + if old_layout.size() == 0 { + debug_assert_eq!(ptr.as_ptr() as usize, old_layout.align()); + let slice = std::slice::from_raw_parts_mut(ptr.as_ptr(), 0); + let ptr = NonNull::new(slice).unwrap_unchecked(); + return Ok(ptr); + } + let flags = layout_to_flags(new_layout.align(), new_layout.size()); + if new_layout.size() == 0 { + ffi::sdallocx(ptr.as_ptr() as *mut c_void, new_layout.size(), flags); + let slice = std::slice::from_raw_parts_mut(new_layout.align() as *mut u8, 0); + let ptr = NonNull::new(slice).unwrap_unchecked(); + Ok(ptr) + } else { + let data_address = + ffi::rallocx(ptr.cast().as_ptr(), new_layout.size(), flags) as *mut u8; + let metadata = new_layout.size(); + let slice = std::slice::from_raw_parts_mut(data_address, metadata); + let ptr = NonNull::new(slice).ok_or(AllocError)?; + Ok(ptr) } - ffi::rallocx(ptr as *mut _, new_size, flags) as *mut u8 } } } -#[cfg(test)] -mod test { +#[cfg(not(any(target_os = "linux", target_os = "macos")))] +pub mod fallback { + use std::alloc::AllocError; + use std::alloc::Allocator; use std::alloc::Layout; + use std::ptr::NonNull; - use crate::mem_allocator::Allocator; - use crate::mem_allocator::JEAllocator; + use super::JEAllocator; - #[test] - fn test_malloc() { - type T = i64; - let mut alloc = JEAllocator::default(); + impl JEAllocator { + pub const FALLBACK: bool = true; + } - let align = std::mem::align_of::(); - let size = std::mem::size_of::() * 100; - let new_size = std::mem::size_of::() * 1000000; + unsafe impl Allocator for JEAllocator { + #[inline(always)] + fn allocate(&self, layout: Layout) -> Result, AllocError> { + self.allocator.allocate(layout) + } + + #[inline(always)] + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + self.allocator.deallocate(ptr, layout) + } - unsafe { - let layout = Layout::from_size_align_unchecked(size, align); - let ptr = alloc.allocx(layout, true) as *mut T; - *ptr = 84; - assert_eq!(84, *ptr); - assert_eq!(0, *(ptr.offset(5))); + #[inline(always)] + fn allocate_zeroed(&self, layout: Layout) -> Result, AllocError> { + self.allocator.allocate_zeroed(layout) + } - *(ptr.offset(5)) = 1000; + unsafe fn grow( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + self.allocator.grow(ptr, old_layout, new_layout) + } - let new_ptr = alloc.reallocx(ptr as *mut u8, layout, new_size, true) as *mut T; - assert_eq!(84, *new_ptr); - assert_eq!(0, *(new_ptr.offset(4))); - assert_eq!(1000, *(new_ptr.offset(5))); + unsafe fn grow_zeroed( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + self.allocator.grow_zeroed(ptr, old_layout, new_layout) + } - alloc.deallocx(new_ptr as *mut u8, layout) + unsafe fn shrink( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + self.allocator.shrink(ptr, old_layout, new_layout) } } } diff --git a/src/common/base/src/mem_allocator/mmap_allocator.rs b/src/common/base/src/mem_allocator/mmap_allocator.rs index fd69799deb63c..75a735de21450 100644 --- a/src/common/base/src/mem_allocator/mmap_allocator.rs +++ b/src/common/base/src/mem_allocator/mmap_allocator.rs @@ -16,332 +16,320 @@ // License, v. 2.0. If a copy of the MPL was not distributed with this // file, You can obtain one at https://mozilla.org/MPL/2.0/. -use std::alloc::Layout; -use std::cell::Cell; -use std::os::raw::c_int; -use std::os::raw::c_long; -use std::os::raw::c_void; -use std::ptr; +#[derive(Debug, Clone, Copy, Default)] +pub struct MmapAllocator { + allocator: T, +} -use libc::off_t; -use libc::size_t; +impl MmapAllocator { + pub fn new(allocator: T) -> Self { + Self { allocator } + } +} -use super::JEAllocator; -use crate::base::ThreadTracker; -use crate::mem_allocator::Allocator as AllocatorTrait; +#[cfg(target_os = "linux")] +pub mod linux { + use std::alloc::AllocError; + use std::alloc::Allocator; + use std::alloc::Layout; + use std::ptr::null_mut; + use std::ptr::NonNull; -const MMAP_THRESHOLD: size_t = 64 * (1usize << 20); -const MALLOC_MIN_ALIGNMENT: size_t = 8; + use super::MmapAllocator; -/// Implementation of std::alloc::AllocatorTrait whose backend is mmap(2) -#[derive(Debug, Clone, Copy, Default)] -pub struct MmapAllocator { - allocator: JEAllocator, -} + // MADV_POPULATE_WRITE is supported since Linux 5.14. + const MADV_POPULATE_WRITE: i32 = 23; -/// # Portability -/// -/// allocx() calls mmap() with flag MAP_ANONYMOUS. -/// Many systems support the flag, however, it is not specified in POSIX. -/// -/// # Safety -/// -/// All functions are thread safe. -/// -/// # Error -/// -/// Each function don't cause panic but set OS errno on error. -/// -/// Note that it is not an error to deallocate pointer which is not allocated. -/// This is the spec of munmap(2). See `man 2 munmap` for details. -unsafe impl AllocatorTrait for MmapAllocator { - /// # Panics - /// - /// This method may panic if the align of `layout` is greater than the kernel page align. - /// (Basically, kernel page align is always greater than the align of `layout` that rust - /// generates unless the programer dares to build such a `layout` on purpose.) - #[inline] - unsafe fn allocx(&mut self, layout: Layout, clear_mem: bool) -> *mut u8 { - #[cfg(not(target_os = "linux"))] - let flags: c_int = libc::MAP_PRIVATE | libc::MAP_ANONYMOUS; - #[cfg(target_os = "linux")] - let flags: c_int = libc::MAP_PRIVATE - | libc::MAP_ANONYMOUS - | if MMAP_POPULATE { libc::MAP_POPULATE } else { 0 }; + const THRESHOLD: usize = 64 << 20; - const ADDR: *mut c_void = ptr::null_mut::(); - const PROT: c_int = libc::PROT_READ | libc::PROT_WRITE; - const FD: c_int = -1; // Should be -1 if flags includes MAP_ANONYMOUS. See `man 2 mmap` - const OFFSET: off_t = 0; // Should be 0 if flags includes MAP_ANONYMOUS. See `man 2 mmap` - let length = layout.size() as size_t; + impl MmapAllocator { + pub const FALLBACK: bool = false; + } - if layout.size() >= MMAP_THRESHOLD { - if layout.align() > page_size() { - // too large alignment, fallback to use allocator - return self.allocator.allocx(layout, clear_mem); + impl MmapAllocator { + #[inline(always)] + fn mmap_alloc(&self, layout: Layout) -> Result, AllocError> { + debug_assert!(layout.align() <= page_size()); + const PROT: i32 = libc::PROT_READ | libc::PROT_WRITE; + const FLAGS: i32 = libc::MAP_PRIVATE | libc::MAP_ANONYMOUS | libc::MAP_POPULATE; + let addr = unsafe { libc::mmap(null_mut(), layout.size(), PROT, FLAGS, -1, 0) }; + if addr == libc::MAP_FAILED { + return Err(AllocError); } + let addr = NonNull::new(addr as *mut ()).ok_or(AllocError)?; + Ok(NonNull::<[u8]>::from_raw_parts(addr, layout.size())) + } + + #[inline(always)] + unsafe fn mmap_dealloc(&self, ptr: NonNull, layout: Layout) { + debug_assert!(layout.align() <= page_size()); + let result = libc::munmap(ptr.cast().as_ptr(), layout.size()); + assert_eq!(result, 0, "Failed to deallocate."); + } - ThreadTracker::alloc_memory(layout.size() as i64); - match mmap(ADDR, length, PROT, flags, FD, OFFSET) { - libc::MAP_FAILED => ptr::null_mut::(), - ret => { - let ptr = ret as usize; - debug_assert_eq!(0, ptr % layout.align()); - ret as *mut u8 - } + #[inline(always)] + unsafe fn mmap_grow( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + debug_assert!(old_layout.align() <= page_size()); + debug_assert!(old_layout.align() == new_layout.align()); + const REMAP_FLAGS: i32 = libc::MREMAP_MAYMOVE; + let addr = libc::mremap( + ptr.cast().as_ptr(), + old_layout.size(), + new_layout.size(), + REMAP_FLAGS, + ); + if addr == libc::MAP_FAILED { + return Err(AllocError); } - } else { - self.allocator.allocx(layout, clear_mem) + let addr = NonNull::new(addr as *mut ()).ok_or(AllocError)?; + if linux_kernel_version() >= (5, 14, 0) { + libc::madvise(addr.cast().as_ptr(), new_layout.size(), MADV_POPULATE_WRITE); + } + Ok(NonNull::<[u8]>::from_raw_parts(addr, new_layout.size())) } - } - #[inline] - unsafe fn deallocx(&mut self, ptr: *mut u8, layout: Layout) { - let addr = ptr as *mut c_void; - if layout.size() >= MMAP_THRESHOLD { - munmap(addr, layout.size()); - } else { - self.allocator.deallocx(ptr, layout); + #[inline(always)] + unsafe fn mmap_shrink( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + debug_assert!(old_layout.align() <= page_size()); + debug_assert!(old_layout.align() == new_layout.align()); + const REMAP_FLAGS: i32 = libc::MREMAP_MAYMOVE; + let addr = libc::mremap( + ptr.cast().as_ptr(), + old_layout.size(), + new_layout.size(), + REMAP_FLAGS, + ); + if addr == libc::MAP_FAILED { + return Err(AllocError); + } + let addr = NonNull::new(addr as *mut ()).ok_or(AllocError)?; + Ok(NonNull::<[u8]>::from_raw_parts(addr, new_layout.size())) } } - #[inline] - unsafe fn reallocx( - &mut self, - ptr: *mut u8, - layout: Layout, - new_size: usize, - clear_mem: bool, - ) -> *mut u8 { - use libc::PROT_READ; - use libc::PROT_WRITE; - - if new_size == layout.size() { - return ptr; + unsafe impl Allocator for MmapAllocator { + #[inline(always)] + fn allocate(&self, layout: Layout) -> Result, AllocError> { + if layout.align() > page_size() { + return self.allocator.allocate(layout); + } + if layout.size() >= THRESHOLD { + self.mmap_alloc(layout) + } else { + self.allocator.allocate(layout) + } } - if layout.size() < MMAP_THRESHOLD - && new_size < MMAP_THRESHOLD - && layout.align() <= MALLOC_MIN_ALIGNMENT - { - self.allocator.reallocx(ptr, layout, new_size, clear_mem) - } else if layout.size() >= MMAP_THRESHOLD && new_size >= MMAP_THRESHOLD { + #[inline(always)] + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { if layout.align() > page_size() { - self.allocator.reallocx(ptr, layout, new_size, clear_mem) + return self.allocator.deallocate(ptr, layout); + } + if layout.size() >= THRESHOLD { + self.mmap_dealloc(ptr, layout); } else { - mremapx( - ptr as *mut c_void, - layout.size(), - new_size, - 0, - PROT_READ | PROT_WRITE, - ) as *mut u8 + self.allocator.deallocate(ptr, layout); } - } else { - let new_buf = self.allocx( - Layout::from_size_align_unchecked(new_size, layout.align()), - clear_mem, - ); + } - if new_buf.is_null() { - return ptr::null_mut(); + #[inline(always)] + fn allocate_zeroed(&self, layout: Layout) -> Result, AllocError> { + if layout.align() > page_size() { + return self.allocator.allocate_zeroed(layout); + } + if layout.size() >= THRESHOLD { + self.mmap_alloc(layout) + } else { + self.allocator.allocate_zeroed(layout) } - std::ptr::copy_nonoverlapping(ptr, new_buf, new_size.min(layout.size())); - self.deallocx(ptr, layout); - new_buf } - } -} - -#[inline] -unsafe fn mremapx( - old_address: *mut c_void, - old_size: size_t, - new_size: size_t, - flags: c_int, - mmap_prot: c_int, -) -> *mut c_void { - #[cfg(not(target_os = "linux"))] - { - const ADDR: *mut c_void = ptr::null_mut::(); - const FD: c_int = -1; // Should be -1 if flags includes MAP_ANONYMOUS. See `man 2 mmap` - const OFFSET: off_t = 0; // Should be 0 if flags includes MAP_ANONYMOUS. See `man 2 mmap` - match mmap(ADDR, new_size, mmap_prot, flags, FD, OFFSET) { - libc::MAP_FAILED => ptr::null_mut::(), - new_address => { + unsafe fn grow( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + if old_layout.align() > page_size() { + return self.allocator.grow(ptr, old_layout, new_layout); + } + if old_layout.size() >= THRESHOLD { + self.mmap_grow(ptr, old_layout, new_layout) + } else if new_layout.size() >= THRESHOLD { + let addr = self.mmap_alloc(new_layout)?; std::ptr::copy_nonoverlapping( - old_address as *const u8, - new_address as *mut u8, - old_size as usize, + ptr.as_ptr(), + addr.cast().as_ptr(), + old_layout.size(), ); - new_address + self.allocator.deallocate(ptr, old_layout); + Ok(addr) + } else { + self.allocator.grow(ptr, old_layout, new_layout) } } - } - #[cfg(target_os = "linux")] - mremap( - old_address, - old_size, - new_size, - flags | libc::MREMAP_MAYMOVE, - mmap_prot, - ) -} - -extern "C" { - fn mmap( - addr: *mut c_void, - length: size_t, - prot: c_int, - flags: c_int, - fd: c_int, - offset: off_t, - ) -> *mut c_void; - - fn munmap(addr: *mut c_void, length: size_t); - - #[cfg(target_os = "linux")] - fn mremap( - old_address: *mut c_void, - old_size: size_t, - new_size: size_t, - flags: c_int, - mmap_prot: c_int, - ) -> *mut c_void; -} - -#[cfg(test)] -mod tests { - use std::mem; - use std::ptr; - - use super::*; - use crate::mem_allocator::Allocator as AllocatorTrait; - type Aloc = MmapAllocator; - fn clear_errno() { - #[cfg(target_os = "linux")] - unsafe { - *libc::__errno_location() = 0 + unsafe fn grow_zeroed( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + if old_layout.align() > page_size() { + return self.allocator.grow_zeroed(ptr, old_layout, new_layout); + } + if old_layout.size() >= THRESHOLD { + self.mmap_grow(ptr, old_layout, new_layout) + } else if new_layout.size() >= THRESHOLD { + let addr = self.mmap_alloc(new_layout)?; + std::ptr::copy_nonoverlapping( + ptr.as_ptr(), + addr.cast().as_ptr(), + old_layout.size(), + ); + self.allocator.deallocate(ptr, old_layout); + Ok(addr) + } else { + self.allocator.grow_zeroed(ptr, old_layout, new_layout) + } } - } - - #[test] - fn default() { - let _alloc = Aloc::default(); - } - - #[test] - fn allocate() { - unsafe { - type T = i64; - let mut alloc = Aloc::default(); - - let layout = Layout::new::(); - let ptr = alloc.allocx(layout, false) as *mut T; - assert_ne!(std::ptr::null(), ptr); - - *ptr = 84; - assert_eq!(84, *ptr); - *ptr *= -2; - assert_eq!(-168, *ptr); - - alloc.deallocx(ptr as *mut u8, layout) + unsafe fn shrink( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + if old_layout.align() > page_size() { + return self.allocator.shrink(ptr, old_layout, new_layout); + } + if new_layout.size() >= THRESHOLD { + self.mmap_shrink(ptr, old_layout, new_layout) + } else if old_layout.size() >= THRESHOLD { + let addr = self.allocator.allocate(new_layout)?; + std::ptr::copy_nonoverlapping( + ptr.as_ptr(), + addr.cast().as_ptr(), + old_layout.size(), + ); + self.mmap_dealloc(ptr, old_layout); + Ok(addr) + } else { + self.allocator.shrink(ptr, old_layout, new_layout) + } } } - #[test] - fn allocate_too_large() { - unsafe { - clear_errno(); - - type T = String; - let mut alloc = Aloc::default(); - - let align = mem::align_of::(); - let size = std::usize::MAX - mem::size_of::(); - let layout = Layout::from_size_align_unchecked(size, align); - - assert_eq!(ptr::null(), alloc.allocx(layout, false)); + #[inline(always)] + fn page_size() -> usize { + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + const INVAILED: usize = 0; + static CACHE: AtomicUsize = AtomicUsize::new(INVAILED); + let fetch = CACHE.load(Ordering::Relaxed); + if fetch == INVAILED { + let result = unsafe { libc::sysconf(libc::_SC_PAGE_SIZE) as usize }; + debug_assert_eq!(result.count_ones(), 1); + CACHE.store(result, Ordering::Relaxed); + result + } else { + fetch } } - #[test] - fn alloc_zeroed() { - unsafe { - type T = [u8; 1025]; - let mut alloc = Aloc::default(); - - let layout = Layout::new::(); - let ptr = alloc.allocx(layout, false) as *mut T; - let s: &[u8] = &*ptr; - - for u in s { - assert_eq!(0, *u); + #[inline(always)] + fn linux_kernel_version() -> (u16, u8, u8) { + use std::sync::atomic::AtomicU32; + use std::sync::atomic::Ordering; + const INVAILED: u32 = 0; + static CACHE: AtomicU32 = AtomicU32::new(INVAILED); + let fetch = CACHE.load(Ordering::Relaxed); + let code = if fetch == INVAILED { + let mut uname = unsafe { std::mem::zeroed::() }; + assert_ne!(-1, unsafe { libc::uname(&mut uname) }); + let mut length = 0usize; + while length < uname.release.len() && uname.release[length] != 0 { + length += 1; } - - alloc.deallocx(ptr as *mut u8, layout); - } + let slice = unsafe { &*(&uname.release[..length] as *const _ as *const [u8]) }; + let ver = std::str::from_utf8(slice).unwrap(); + let semver = semver::Version::parse(ver).unwrap(); + let result = (semver.major.min(65535) as u32) << 16 + | (semver.minor.min(255) as u32) << 8 + | (semver.patch.min(255) as u32); + CACHE.store(result, Ordering::Relaxed); + result + } else { + fetch + }; + ((code >> 16) as u16, (code >> 8) as u8, code as u8) } +} - #[test] - fn reallocx() { - unsafe { - type T = [u8; 1025]; - let mut alloc = Aloc::default(); - - let layout = Layout::new::(); - let ptr = alloc.allocx(layout, false) as *mut T; - - let ts = &mut *ptr; - for t in ts.iter_mut() { - *t = 1; - } - - type U = (T, T); +#[cfg(not(target_os = "linux"))] +pub mod fallback { + use std::alloc::AllocError; + use std::alloc::Allocator; + use std::alloc::Layout; + use std::ptr::NonNull; - let new_size = mem::size_of::(); - let ptr = alloc.reallocx(ptr as *mut u8, layout, new_size, false) as *mut T; - let layout = Layout::from_size_align(new_size, layout.align()).unwrap(); + use super::MmapAllocator; - let ts = &mut *ptr; - for t in ts.iter_mut() { - assert_eq!(1, *t); - *t = 2; - } + impl MmapAllocator { + pub const FALLBACK: bool = true; + } - let new_size = mem::size_of::(); - let ptr = alloc.reallocx(ptr as *mut u8, layout, new_size, false); - let layout = Layout::from_size_align(new_size, layout.align()).unwrap(); + unsafe impl Allocator for MmapAllocator { + #[inline(always)] + fn allocate(&self, layout: Layout) -> Result, AllocError> { + self.allocator.allocate(layout) + } - assert_eq!(2, *ptr); + #[inline(always)] + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + self.allocator.deallocate(ptr, layout) + } - alloc.deallocx(ptr, layout); + #[inline(always)] + fn allocate_zeroed(&self, layout: Layout) -> Result, AllocError> { + self.allocator.allocate_zeroed(layout) } - } -} -thread_local! { - static PAGE_SIZE: Cell = Cell::new(0); -} + unsafe fn grow( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + self.allocator.grow(ptr, old_layout, new_layout) + } -/// Returns OS Page Size. -/// -/// See crate document for details. -#[inline] -pub fn page_size() -> usize { - PAGE_SIZE.with(|s| match s.get() { - 0 => { - let ret = unsafe { sysconf(libc::_SC_PAGE_SIZE) as usize }; - s.set(ret); - ret + unsafe fn grow_zeroed( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + self.allocator.grow_zeroed(ptr, old_layout, new_layout) } - ret => ret, - }) -} -extern "C" { - fn sysconf(name: c_int) -> c_long; + unsafe fn shrink( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + self.allocator.shrink(ptr, old_layout, new_layout) + } + } } diff --git a/src/common/base/src/mem_allocator/mod.rs b/src/common/base/src/mem_allocator/mod.rs index d18b78f9a8fa8..738b9aaa66f89 100644 --- a/src/common/base/src/mem_allocator/mod.rs +++ b/src/common/base/src/mem_allocator/mod.rs @@ -12,55 +12,18 @@ // See the License for the specific language governing permissions and // limitations under the License. +mod global_allocator; mod je_allocator; mod mmap_allocator; -mod stackful_allocator; -mod std_allocator; - -use std::alloc::Layout; +mod system_allocator; +pub use global_allocator::GlobalAllocator; pub use je_allocator::JEAllocator; -pub use je_allocator::ALLOC; pub use mmap_allocator::MmapAllocator; -pub use stackful_allocator::StackfulAllocator; -pub use std_allocator::StdAllocator; +pub use system_allocator::SystemAllocator; #[cfg(feature = "memory-profiling")] mod profiling; #[cfg(feature = "memory-profiling")] pub use profiling::dump_profile; - -// A new Allocator trait suits databend -/// # Safety -/// -/// All functions are thread safe. -pub unsafe trait Allocator { - /// # Safety - /// the caller must ensure that the layout is valid - unsafe fn allocx(&mut self, layout: Layout, clear_mem: bool) -> *mut u8; - /// # Safety - /// the caller must ensure that the ptr is valid - unsafe fn deallocx(&mut self, ptr: *mut u8, layout: Layout); - /// # Safety - /// the caller must ensure that the `new_size` does not overflow. - /// `layout.align()` comes from a `Layout` and is thus guaranteed to be valid. - unsafe fn reallocx( - &mut self, - ptr: *mut u8, - layout: Layout, - new_size: usize, - clear_mem: bool, - ) -> *mut u8 { - let new_layout = Layout::from_size_align_unchecked(new_size, layout.align()); - // SAFETY: the caller must ensure that `new_layout` is greater than zero. - let new_ptr = self.allocx(new_layout, clear_mem); - if !new_ptr.is_null() { - // SAFETY: the previously allocated block cannot overlap the newly allocated block. - // The safety contract for `dealloc` must be upheld by the caller. - std::ptr::copy_nonoverlapping(ptr, new_ptr, std::cmp::min(layout.size(), new_size)); - self.deallocx(ptr, layout); - } - new_ptr - } -} diff --git a/src/common/base/src/mem_allocator/stackful_allocator.rs b/src/common/base/src/mem_allocator/stackful_allocator.rs deleted file mode 100644 index e8656cfbaec6a..0000000000000 --- a/src/common/base/src/mem_allocator/stackful_allocator.rs +++ /dev/null @@ -1,174 +0,0 @@ -// Copyright 2022 Datafuse Labs. -// -// Licensed 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. - -// This Source Code Form is subject to the terms of the Mozilla Public -// License, v. 2.0. If a copy of the MPL was not distributed with this -// file, You can obtain one at https://mozilla.org/MPL/2.0/. - -use std::alloc::Layout; -use std::ptr; - -use super::Allocator as AllocatorTrait; - -/// Implementation of std::alloc::Allocator whose backend is mmap(2) -#[derive(Debug, Clone, Copy)] -pub struct StackfulAllocator { - allocator: Allocator, - stack_bytes: [u8; INIT_BYTES], -} - -impl Default - for StackfulAllocator -{ - fn default() -> Self { - Self { - allocator: Default::default(), - stack_bytes: [0; INIT_BYTES], - } - } -} - -/// # Safety -/// -/// We should guarantee that the stack is only used as inner memory, such as HashSet -unsafe impl AllocatorTrait - for StackfulAllocator -{ - unsafe fn allocx(&mut self, layout: Layout, clear_mem: bool) -> *mut u8 { - if layout.size() <= INIT_BYTES { - if clear_mem { - ptr::write_bytes(self.stack_bytes.as_mut_ptr(), 0, INIT_BYTES); - } - return self.stack_bytes.as_mut_ptr(); - } - self.allocator.allocx(layout, clear_mem) - } - - unsafe fn deallocx(&mut self, ptr: *mut u8, layout: Layout) { - if layout.size() > INIT_BYTES { - self.allocator.deallocx(ptr, layout); - } - } - - unsafe fn reallocx( - &mut self, - ptr: *mut u8, - layout: Layout, - new_size: usize, - clear_mem: bool, - ) -> *mut u8 { - if new_size <= INIT_BYTES { - return ptr; - } - if layout.size() > INIT_BYTES { - return self.allocator.reallocx(ptr, layout, new_size, clear_mem); - } - - let new_buf = self.allocator.allocx( - Layout::from_size_align_unchecked(new_size, layout.align()), - clear_mem, - ); - std::ptr::copy_nonoverlapping(ptr, new_buf, layout.size()); - new_buf - } -} - -#[cfg(test)] -mod tests { - use std::mem; - - use super::*; - use crate::mem_allocator::Allocator as AllocatorTrait; - use crate::mem_allocator::MmapAllocator; - type Aloc = StackfulAllocator<1024, MmapAllocator>; - - #[test] - fn default() { - let _alloc = Aloc::default(); - } - - #[test] - fn allocate() { - unsafe { - type T = i64; - let mut alloc = Aloc::default(); - - let layout = Layout::new::(); - let ptr = alloc.allocx(layout, false) as *mut T; - assert_ne!(std::ptr::null(), ptr); - - *ptr = 84; - assert_eq!(84, *ptr); - - *ptr *= -2; - assert_eq!(-168, *ptr); - - alloc.deallocx(ptr as *mut u8, layout) - } - } - - #[test] - fn alloc_zeroed() { - unsafe { - type T = [u8; 1025]; - let mut alloc = Aloc::default(); - - let layout = Layout::new::(); - let ptr = alloc.allocx(layout, true) as *mut T; - let s: &[u8] = &*ptr; - - for u in s { - assert_eq!(0, *u); - } - - alloc.deallocx(ptr as *mut u8, layout); - } - } - - #[test] - fn reallocx() { - unsafe { - type T = [u8; 1025]; - let mut alloc = Aloc::default(); - - let layout = Layout::new::(); - let ptr = alloc.allocx(layout, false) as *mut T; - - let ts = &mut *ptr; - for t in ts.iter_mut() { - *t = 1; - } - - type U = (T, T); - - let new_size = mem::size_of::(); - let ptr = alloc.reallocx(ptr as *mut u8, layout, new_size, false) as *mut T; - let layout = Layout::from_size_align(new_size, layout.align()).unwrap(); - - let ts = &mut *ptr; - for t in ts.iter_mut() { - assert_eq!(1, *t); - *t = 2; - } - - let new_size = mem::size_of::(); - let ptr = alloc.reallocx(ptr as *mut u8, layout, new_size, false); - let layout = Layout::from_size_align(new_size, layout.align()).unwrap(); - - assert_eq!(2, *ptr); - - alloc.deallocx(ptr, layout); - } - } -} diff --git a/src/common/base/src/mem_allocator/std_allocator.rs b/src/common/base/src/mem_allocator/std_allocator.rs deleted file mode 100644 index 124a48582802e..0000000000000 --- a/src/common/base/src/mem_allocator/std_allocator.rs +++ /dev/null @@ -1,87 +0,0 @@ -// Copyright 2022 Datafuse Labs. -// -// Licensed 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. - -// This Source Code Form is subject to the terms of the Mozilla Public -// License, v. 2.0. If a copy of the MPL was not distributed with this -// file, You can obtain one at https://mozilla.org/MPL/2.0/. - -use std::alloc::Layout; - -use super::Allocator; - -#[derive(Debug, Clone, Copy, Default)] -pub struct StdAllocator; - -unsafe impl Allocator for StdAllocator { - unsafe fn allocx(&mut self, layout: Layout, clear_mem: bool) -> *mut u8 { - if clear_mem { - std::alloc::alloc_zeroed(layout) - } else { - std::alloc::alloc(layout) - } - } - - unsafe fn deallocx(&mut self, ptr: *mut u8, layout: Layout) { - std::alloc::dealloc(ptr, layout) - } - - unsafe fn reallocx( - &mut self, - ptr: *mut u8, - layout: Layout, - new_size: usize, - clear_mem: bool, - ) -> *mut u8 { - let ptr = std::alloc::realloc(ptr, layout, new_size); - if clear_mem && new_size > layout.size() { - std::ptr::write_bytes(ptr.add(layout.size()), 0, new_size - layout.size()); - } - ptr - } -} - -#[cfg(test)] -mod test { - use std::alloc::Layout; - - use crate::mem_allocator::Allocator; - use crate::mem_allocator::StdAllocator; - - #[test] - fn test_malloc() { - type T = i64; - let mut alloc = StdAllocator::default(); - - let align = std::mem::align_of::(); - let size = std::mem::size_of::() * 100; - let new_size = std::mem::size_of::() * 1000000; - - unsafe { - let layout = Layout::from_size_align_unchecked(size, align); - let ptr = alloc.allocx(layout, true) as *mut T; - *ptr = 84; - assert_eq!(84, *ptr); - assert_eq!(0, *(ptr.offset(5))); - - *(ptr.offset(5)) = 1000; - - let new_ptr = alloc.reallocx(ptr as *mut u8, layout, new_size, true) as *mut T; - assert_eq!(84, *new_ptr); - assert_eq!(0, *(new_ptr.offset(4))); - assert_eq!(1000, *(new_ptr.offset(5))); - - alloc.deallocx(new_ptr as *mut u8, layout) - } - } -} diff --git a/src/common/base/src/mem_allocator/system_allocator.rs b/src/common/base/src/mem_allocator/system_allocator.rs new file mode 100644 index 0000000000000..c7b321d4b6abd --- /dev/null +++ b/src/common/base/src/mem_allocator/system_allocator.rs @@ -0,0 +1,77 @@ +// Copyright 2022 Datafuse Labs. +// +// Licensed 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 std::alloc::AllocError; +use std::alloc::Allocator; +use std::alloc::Layout; +use std::alloc::System; +use std::ptr::NonNull; + +use crate::base::ThreadTracker; + +#[derive(Debug, Clone, Copy, Default)] +pub struct SystemAllocator; + +unsafe impl Allocator for SystemAllocator { + #[inline(always)] + fn allocate(&self, layout: Layout) -> Result, AllocError> { + ThreadTracker::alloc_memory(layout.size() as i64); + System.allocate(layout) + } + + #[inline(always)] + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + ThreadTracker::dealloc_memory(layout.size() as i64); + System.deallocate(ptr, layout) + } + + #[inline(always)] + fn allocate_zeroed(&self, layout: Layout) -> Result, AllocError> { + ThreadTracker::alloc_memory(layout.size() as i64); + System.allocate_zeroed(layout) + } + + #[inline(always)] + unsafe fn grow( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + ThreadTracker::grow_memory(old_layout.size() as i64, new_layout.size() as i64); + System.grow(ptr, old_layout, new_layout) + } + + #[inline(always)] + unsafe fn grow_zeroed( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + ThreadTracker::grow_memory(old_layout.size() as i64, new_layout.size() as i64); + System.grow_zeroed(ptr, old_layout, new_layout) + } + + #[inline(always)] + unsafe fn shrink( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + ThreadTracker::shrink_memory(old_layout.size() as i64, new_layout.size() as i64); + System.shrink(ptr, old_layout, new_layout) + } +} diff --git a/src/common/hashtable/Cargo.toml b/src/common/hashtable/Cargo.toml index 8ba3adb3011f2..aacb0df567a36 100644 --- a/src/common/hashtable/Cargo.toml +++ b/src/common/hashtable/Cargo.toml @@ -16,8 +16,12 @@ common-base = { path = "../base" } # Crates.io dependencies ahash = "0.8.0" +bumpalo = "3.10.0" +cfg-if = "1.0.0" +libc = "0.2.126" ordered-float = { version = "3.1.0", features = ["serde"] } primitive-types = "0.12.0" +semver = "1.0.10" [dev-dependencies] rand = "0.8.5" diff --git a/src/common/hashtable/src/container.rs b/src/common/hashtable/src/container.rs new file mode 100644 index 0000000000000..da5326e82e26b --- /dev/null +++ b/src/common/hashtable/src/container.rs @@ -0,0 +1,205 @@ +// Copyright 2021 Datafuse Labs. +// +// Licensed 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 std::alloc::Allocator; +use std::alloc::Layout; +use std::mem::MaybeUninit; +use std::ops::Deref; +use std::ops::DerefMut; +use std::ptr::null_mut; +use std::ptr::NonNull; + +/// # Safety +/// +/// Any foreign type shouldn't implement this trait. +pub unsafe trait Container +where Self: Deref + DerefMut +{ + type T; + + type A: Allocator; + + fn len(&self) -> usize; + + unsafe fn new_zeroed(len: usize, allocator: Self::A) -> Self; + + unsafe fn grow_zeroed(&mut self, new_len: usize); +} + +#[derive(Debug)] +pub struct HeapContainer(Box<[T], A>); + +impl Deref for HeapContainer { + type Target = [T]; + + #[inline(always)] + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl DerefMut for HeapContainer { + #[inline(always)] + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.0 + } +} + +unsafe impl Container for HeapContainer { + type T = T; + + type A = A; + + #[inline(always)] + fn len(&self) -> usize { + self.as_ref().len() + } + + unsafe fn new_zeroed(len: usize, allocator: Self::A) -> Self { + Self(Box::new_zeroed_slice_in(len, allocator).assume_init()) + } + + unsafe fn grow_zeroed(&mut self, new_len: usize) { + debug_assert!(self.len() <= new_len); + let old_layout = Layout::array::(self.len()).unwrap(); + let new_layout = Layout::array::(new_len).unwrap(); + let old_box = std::ptr::read(&self.0); + let (old_raw, allocator) = Box::into_raw_with_allocator(old_box); + let old_ptr = NonNull::new(old_raw).unwrap().cast(); + let new_ptr = allocator + .grow_zeroed(old_ptr, old_layout, new_layout) + .unwrap(); + let new_raw = std::ptr::slice_from_raw_parts_mut(new_ptr.cast().as_ptr(), new_len); + let new_box = Box::from_raw_in(new_raw, allocator); + std::ptr::write(self, Self(new_box)); + } +} + +pub struct StackContainer { + allocator: A, + ptr: *mut T, + len: usize, + array: [MaybeUninit; N], +} + +impl Deref for StackContainer { + type Target = [T]; + + #[inline(always)] + fn deref(&self) -> &Self::Target { + unsafe { + std::slice::from_raw_parts( + if self.ptr.is_null() { + self.array.as_ptr() as *const _ + } else { + self.ptr + }, + self.len, + ) + } + } +} + +impl DerefMut for StackContainer { + #[inline(always)] + fn deref_mut(&mut self) -> &mut Self::Target { + unsafe { + std::slice::from_raw_parts_mut( + if self.ptr.is_null() { + self.array.as_mut_ptr() as *mut _ + } else { + self.ptr + }, + self.len, + ) + } + } +} + +unsafe impl Container for StackContainer { + type T = T; + + type A = A; + + #[inline(always)] + fn len(&self) -> usize { + self.len + } + + unsafe fn new_zeroed(len: usize, allocator: Self::A) -> Self { + if len <= N { + Self { + allocator, + ptr: null_mut(), + len, + array: std::array::from_fn(|_| MaybeUninit::zeroed()), + } + } else { + let ptr = allocator + .allocate_zeroed(Layout::array::(len).unwrap()) + .unwrap() + .cast() + .as_ptr(); + Self { + allocator, + ptr, + len, + array: std::array::from_fn(|_| MaybeUninit::uninit()), + } + } + } + + unsafe fn grow_zeroed(&mut self, new_len: usize) { + debug_assert!(self.len <= new_len); + if new_len <= N { + self.len = new_len; + } else if self.ptr.is_null() { + self.ptr = self + .allocator + .allocate_zeroed(Layout::array::(new_len).unwrap()) + .unwrap() + .cast() + .as_ptr(); + std::ptr::copy_nonoverlapping(self.array.as_ptr() as *mut _, self.ptr, self.len); + self.len = new_len; + } else { + let old_layout = Layout::array::(self.len).unwrap(); + let new_layout = Layout::array::(new_len).unwrap(); + self.ptr = self + .allocator + .grow_zeroed( + NonNull::new_unchecked(self.ptr).cast(), + old_layout, + new_layout, + ) + .unwrap() + .cast::() + .as_ptr(); + self.len = new_len; + } + } +} + +impl Drop for StackContainer { + fn drop(&mut self) { + if !self.ptr.is_null() { + unsafe { + self.allocator.deallocate( + NonNull::new(self.ptr).unwrap().cast(), + Layout::array::(self.len).unwrap(), + ); + } + } + } +} diff --git a/src/common/hashtable/src/hash_set.rs b/src/common/hashtable/src/hash_set.rs deleted file mode 100644 index 8707b326f7d16..0000000000000 --- a/src/common/hashtable/src/hash_set.rs +++ /dev/null @@ -1,37 +0,0 @@ -// Copyright 2021 Datafuse Labs. -// -// Licensed 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. -// -// Reference the ClickHouse HashTable to implement the Databend HashTable - -use common_base::mem_allocator::Allocator as AllocatorTrait; - -use crate::HashSet; -use crate::HashTableEntity; -use crate::HashTableGrower; -use crate::HashTableKeyable; - -impl - HashSet -{ - pub fn merge(&mut self, other: &Self) { - let mut inserted = false; - for value in other.iter() { - self.insert_key(value.get_key(), &mut inserted); - } - } - - pub fn contains(&self, key: &Key) -> bool { - self.find_key(key).is_some() - } -} diff --git a/src/common/hashtable/src/hash_table.rs b/src/common/hashtable/src/hash_table.rs deleted file mode 100644 index 12672c3754eaa..0000000000000 --- a/src/common/hashtable/src/hash_table.rs +++ /dev/null @@ -1,373 +0,0 @@ -// Copyright 2021 Datafuse Labs. -// -// Licensed 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. -// -// Reference the ClickHouse HashTable to implement the Databend HashTable - -use std::alloc::GlobalAlloc; -use std::alloc::Layout; -use std::marker::PhantomData; -use std::mem; - -use common_base::mem_allocator::Allocator as AllocatorTrait; -use common_base::mem_allocator::ALLOC; - -use crate::hash_table_grower::HashTableGrower; -use crate::HashTableEntity; -use crate::HashTableIter; -use crate::HashTableIteratorKind; -use crate::HashTableKeyable; - -pub struct HashTable< - Key: HashTableKeyable, - Entity: HashTableEntity, - Grower: HashTableGrower, - Allocator: AllocatorTrait, -> { - size: usize, - grower: Grower, - allocator: Allocator, - entities: *mut Entity, - entities_raw: *mut u8, - zero_entity: Option<*mut Entity>, - zero_entity_raw: Option<*mut u8>, - - // set to true if the table is converted to other hash table - pub(crate) entity_swapped: bool, - /// Generics hold - generics_hold: PhantomData, -} - -unsafe impl< - Key: HashTableKeyable + Send, - Entity: HashTableEntity + Send, - Grower: HashTableGrower, - Allocator: AllocatorTrait, -> Send for HashTable -{ -} - -unsafe impl< - Key: HashTableKeyable + Sync, - Entity: HashTableEntity + Sync, - Grower: HashTableGrower, - Allocator: AllocatorTrait, -> Sync for HashTable -{ -} - -impl< - Key: HashTableKeyable, - Entity: HashTableEntity, - Grower: HashTableGrower, - Allocator: AllocatorTrait, -> Drop for HashTable -{ - fn drop(&mut self) { - unsafe { - let item_size = self.grower.max_size() as usize; - - if std::mem::needs_drop::() && !self.entity_swapped { - for off in 0..item_size { - let entity = self.entities.add(off); - - if !entity.is_zero() { - std::ptr::drop_in_place(entity); - } - } - } - - let size = item_size * mem::size_of::(); - let layout = Layout::from_size_align_unchecked(size, std::mem::align_of::()); - - self.allocator.deallocx(self.entities_raw, layout); - if let Some(zero_entity) = self.zero_entity_raw { - if std::mem::needs_drop::() && !self.entity_swapped { - let entity = self.zero_entity.unwrap(); - std::ptr::drop_in_place(entity); - } - - let zero_layout = Layout::from_size_align_unchecked( - mem::size_of::(), - std::mem::align_of::(), - ); - ALLOC.dealloc(zero_entity, zero_layout); - } - } - } -} - -impl< - Key: HashTableKeyable, - Entity: HashTableEntity, - Grower: HashTableGrower, - Allocator: AllocatorTrait + Default, -> HashTable -{ - pub fn create() -> HashTable { - Self::with_capacity(1 << 8) - } - - pub fn with_capacity(capacity: usize) -> HashTable { - let mut grower = Grower::default(); - while (grower.max_size() as usize) < capacity { - grower.increase_size(); - } - - let size = grower.max_size() as usize * mem::size_of::(); - unsafe { - let layout = Layout::from_size_align_unchecked(size, mem::align_of::()); - let mut allocator = Allocator::default(); - let raw_ptr = allocator.allocx(layout, true); - - if raw_ptr.is_null() { - panic!( - "Failed to have enough memory to alloc {size} bytes to initial the hashtable" - ); - } - - let entities_ptr = raw_ptr as *mut Entity; - HashTable { - size: 0, - grower, - entities: entities_ptr, - entities_raw: raw_ptr, - zero_entity: None, - zero_entity_raw: None, - generics_hold: PhantomData::default(), - entity_swapped: false, - allocator, - } - } - } - - #[inline(always)] - pub fn len(&self) -> usize { - self.size - } - - #[inline(always)] - pub fn is_empty(&self) -> bool { - self.size == 0 - } - - #[inline(always)] - pub fn enum_iter(&self) -> HashTableIteratorKind { - HashTableIteratorKind::create_hash_table_iter( - self.grower.max_size(), - self.entities, - self.zero_entity, - ) - } - - #[inline(always)] - pub fn iter(&self) -> HashTableIter { - HashTableIter::::create( - self.grower.max_size(), - self.entities, - self.zero_entity, - ) - } - - #[inline(always)] - pub fn insert_key(&mut self, key: &Key, inserted: &mut bool) -> *mut Entity { - let hash = key.fast_hash(); - match self.insert_if_zero_key(key, hash, inserted) { - None => self.insert_non_zero_key(key, hash, inserted), - Some(zero_hash_table_entity) => zero_hash_table_entity, - } - } - - #[inline(always)] - pub fn insert_hash_key(&mut self, key: &Key, hash: u64, inserted: &mut bool) -> *mut Entity { - match self.insert_if_zero_key(key, hash, inserted) { - None => self.insert_non_zero_key(key, hash, inserted), - Some(zero_hash_table_entity) => zero_hash_table_entity, - } - } - - #[inline(always)] - pub fn find_key(&self, key: &Key) -> Option<*mut Entity> { - if !key.is_zero() { - let hash_value = key.fast_hash(); - let place_value = self.find_entity(key, hash_value); - unsafe { - let value = self.entities.offset(place_value); - return match value.is_zero() { - true => None, - false => Some(value), - }; - } - } - - self.zero_entity - } - - #[inline(always)] - fn find_entity(&self, key: &Key, hash_value: u64) -> isize { - unsafe { - let grower = &self.grower; - - let mut place_value = grower.place(hash_value); - - while !self.entities.offset(place_value).is_zero() - && !self - .entities - .offset(place_value) - .key_equals(key, hash_value) - { - place_value = grower.next_place(place_value); - } - - place_value - } - } - - #[inline(always)] - fn insert_non_zero_key( - &mut self, - key: &Key, - hash_value: u64, - inserted: &mut bool, - ) -> *mut Entity { - let place_value = self.find_entity(key, hash_value); - self.insert_non_zero_key_impl(place_value, key, hash_value, inserted) - } - - #[inline(always)] - fn insert_non_zero_key_impl( - &mut self, - place_value: isize, - key: &Key, - hash_value: u64, - inserted: &mut bool, - ) -> *mut Entity { - unsafe { - let entity = self.entities.offset(place_value); - - if !entity.is_zero() { - *inserted = false; - return self.entities.offset(place_value); - } - - self.size += 1; - *inserted = true; - entity.set_key_and_hash(key, hash_value); - - if std::intrinsics::unlikely(self.grower.overflow(self.size)) { - self.resize(); - let new_place = self.find_entity(key, hash_value); - return self.entities.offset(new_place); - } - - self.entities.offset(place_value) - } - } - - #[inline(always)] - fn insert_if_zero_key( - &mut self, - key: &Key, - hash_value: u64, - inserted: &mut bool, - ) -> Option<*mut Entity> { - if key.is_zero() { - return match self.zero_entity { - Some(zero_entity) => { - *inserted = false; - Some(zero_entity) - } - None => unsafe { - let layout = Layout::from_size_align_unchecked( - mem::size_of::(), - mem::align_of::(), - ); - - self.size += 1; - *inserted = true; - - self.zero_entity_raw = Some(ALLOC.alloc_zeroed(layout)); - self.zero_entity = Some(self.zero_entity_raw.unwrap() as *mut Entity); - self.zero_entity.unwrap().set_key_and_hash(key, hash_value); - self.zero_entity - }, - }; - } - - Option::None - } - - unsafe fn resize(&mut self) { - let old_grow_size = self.grower.max_size(); - let mut new_grower = self.grower.clone(); - - new_grower.increase_size(); - - // Realloc memory - if new_grower.max_size() > self.grower.max_size() { - let old_size = (old_grow_size as usize) * std::mem::size_of::(); - let new_size = (new_grower.max_size() as usize) * std::mem::size_of::(); - let layout = - Layout::from_size_align_unchecked(old_size, std::mem::align_of::()); - - self.entities_raw = self - .allocator - .reallocx(self.entities_raw, layout, new_size, true); - - if self.entities_raw.is_null() { - panic!( - "Failed to have enough memory to realloc {new_size} bytes to resize the hashtable" - ); - } - - self.entities = self.entities_raw as *mut Entity; - self.grower = new_grower; - for index in 0..old_grow_size { - let entity_ptr = self.entities.offset(index); - - if !entity_ptr.is_zero() { - self.reinsert(entity_ptr, entity_ptr.get_hash()); - } - } - - // There is also a special case: - // if the element was to be at the end of the old buffer, [ x] - // but is at the beginning because of the collision resolution chain, [o x] - // then after resizing, it will first be out of place again, [ xo ] - // and in order to transfer it where necessary, - // after transferring all the elements from the old halves you need to [ o x ] - // process tail from the collision resolution chain immediately after it [ o x ] - for index in old_grow_size..self.grower.max_size() { - let entity = self.entities.offset(index); - - if entity.is_zero() { - return; - } - - self.reinsert(self.entities.offset(index), entity.get_hash()); - } - } - } - - #[inline(always)] - unsafe fn reinsert(&self, entity: *mut Entity, hash_value: u64) { - if entity != self.entities.offset(self.grower.place(hash_value)) { - let place = self.find_entity(entity.get_key(), hash_value); - let new_entity = self.entities.offset(place); - - if new_entity.is_zero() { - entity.swap(new_entity); - } - } - } -} diff --git a/src/common/hashtable/src/hash_table_entity.rs b/src/common/hashtable/src/hash_table_entity.rs deleted file mode 100644 index 4bf67cdeba222..0000000000000 --- a/src/common/hashtable/src/hash_table_entity.rs +++ /dev/null @@ -1,103 +0,0 @@ -// Copyright 2021 Datafuse Labs. -// -// Licensed 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::HashTableKeyable; - -pub trait HashTableEntity: Sized { - unsafe fn is_zero(self: *mut Self) -> bool; - unsafe fn key_equals(self: *mut Self, key: &Key, hash: u64) -> bool; - unsafe fn set_key_and_hash(self: *mut Self, key: &Key, hash: u64); - - fn get_key<'a>(self: *mut Self) -> &'a Key; - unsafe fn get_hash(self: *mut Self) -> u64; - - unsafe fn not_equals_key(self: *mut Self, other: *mut Self) -> bool; -} - -#[derive(Default)] -pub struct KeyValueEntity -where - Key: HashTableKeyable, - Value: Sized + Clone, -{ - key: Key, - value: Value, - hash: u64, -} - -impl KeyValueEntity -where - Key: HashTableKeyable, - Value: Sized + Clone, -{ - #[inline(always)] - pub fn set_key(self: *mut Self, key: Key) { - unsafe { std::ptr::write(&mut (*self).key as *mut Key, key) } - } - - #[inline(always)] - pub fn set_value(self: *mut Self, value: Value) { - unsafe { std::ptr::write(&mut (*self).value as *mut Value, value) } - } - - #[inline(always)] - pub fn get_value<'a>(self: *mut Self) -> &'a Value { - unsafe { &(*self).value } - } - - #[inline(always)] - pub fn get_mut_value<'a>(self: *mut Self) -> &'a mut Value { - unsafe { &mut (*self).value } - } -} - -impl HashTableEntity for KeyValueEntity -where - Key: HashTableKeyable, - Value: Sized + Clone, -{ - #[inline(always)] - unsafe fn is_zero(self: *mut Self) -> bool { - (*self).key.is_zero() - } - - unsafe fn key_equals(self: *mut Self, key: &Key, hash: u64) -> bool { - if Key::BEFORE_EQ_HASH && (*self).hash != hash { - return false; - } - - (*self).key == *key - } - - unsafe fn set_key_and_hash(self: *mut Self, key: &Key, hash: u64) { - (*self).hash = hash; - (*self).key.set_key(key); - } - - fn get_key<'a>(self: *mut Self) -> &'a Key { - unsafe { &(*self).key } - } - - unsafe fn get_hash(self: *mut Self) -> u64 { - (*self).hash - } - - unsafe fn not_equals_key(self: *mut Self, other: *mut Self) -> bool { - if Key::BEFORE_EQ_HASH && (*self).hash != (*other).hash { - return true; - } - - !((*self).key == (*other).key) - } -} diff --git a/src/common/hashtable/src/hash_table_grower.rs b/src/common/hashtable/src/hash_table_grower.rs deleted file mode 100644 index 7283e4a5cbd4e..0000000000000 --- a/src/common/hashtable/src/hash_table_grower.rs +++ /dev/null @@ -1,107 +0,0 @@ -// Copyright 2021 Datafuse Labs. -// -// Licensed 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. - -pub trait HashTableGrower: Default + Clone { - fn max_size(&self) -> isize; - fn overflow(&self, size: usize) -> bool; - fn place(&self, hash_value: u64) -> isize; - fn next_place(&self, old_place: isize) -> isize; - fn increase_size(&mut self); -} - -#[derive(Clone)] -pub struct SingleLevelGrower { - size_degree: u8, - max_size: isize, -} - -impl Default for SingleLevelGrower { - fn default() -> Self { - SingleLevelGrower { - size_degree: 8, - max_size: 1_isize << 8, - } - } -} - -impl HashTableGrower for SingleLevelGrower { - #[inline(always)] - fn max_size(&self) -> isize { - self.max_size - } - - #[inline(always)] - fn overflow(&self, size: usize) -> bool { - size > ((1_usize) << (self.size_degree - 1)) - } - - #[inline(always)] - fn place(&self, hash_value: u64) -> isize { - hash_value as isize & (self.max_size() - 1) - } - - #[inline(always)] - fn next_place(&self, old_place: isize) -> isize { - (old_place + 1) & (self.max_size() - 1) - } - - #[inline(always)] - fn increase_size(&mut self) { - self.size_degree += if self.size_degree >= 23 { 1 } else { 2 }; - self.max_size = 1_isize << self.size_degree; - } -} - -#[derive(Clone)] -pub struct TwoLevelGrower { - size_degree: u8, - max_size: isize, -} - -impl Default for TwoLevelGrower { - fn default() -> Self { - TwoLevelGrower { - size_degree: 8, - max_size: 1_isize << 8, - } - } -} - -impl HashTableGrower for TwoLevelGrower { - #[inline(always)] - fn max_size(&self) -> isize { - self.max_size - } - - #[inline(always)] - fn overflow(&self, size: usize) -> bool { - size > ((1_usize) << (self.size_degree - 1)) - } - - #[inline(always)] - fn place(&self, hash_value: u64) -> isize { - hash_value as isize & (self.max_size() - 1) - } - - #[inline(always)] - fn next_place(&self, old_place: isize) -> isize { - (old_place + 1) & (self.max_size() - 1) - } - - #[inline(always)] - fn increase_size(&mut self) { - self.size_degree += if self.size_degree >= 15 { 1 } else { 2 }; - self.max_size = 1_isize << self.size_degree; - } -} diff --git a/src/common/hashtable/src/hash_table_iter.rs b/src/common/hashtable/src/hash_table_iter.rs deleted file mode 100644 index ef70995bd23ed..0000000000000 --- a/src/common/hashtable/src/hash_table_iter.rs +++ /dev/null @@ -1,127 +0,0 @@ -// Copyright 2021 Datafuse Labs. -// -// Licensed 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 std::marker::PhantomData; - -use crate::HashTableEntity; - -pub enum HashTableIteratorKind> { - HashMapIterator(HashTableIter), - TwoLevelHashMapIter(TwoLevelHashTableIter), -} - -impl> HashTableIteratorKind { - pub fn create_hash_table_iter( - capacity: isize, - entities: *mut Entity, - zero_entity: Option<*mut Entity>, - ) -> Self { - Self::HashMapIterator(HashTableIter::::create( - capacity, - entities, - zero_entity, - )) - } - - pub fn create_two_level_hash_table_iter(iters: Vec>) -> Self { - Self::TwoLevelHashMapIter(TwoLevelHashTableIter::::create(iters)) - } -} - -impl> Iterator for HashTableIteratorKind { - type Item = *mut Entity; - fn next(&mut self) -> Option { - match self { - HashTableIteratorKind::HashMapIterator(it) => it.next(), - HashTableIteratorKind::TwoLevelHashMapIter(it) => it.next(), - } - } -} - -pub struct HashTableIter> { - idx: isize, - capacity: isize, - entities: *mut Entity, - zero_entity: Option<*mut Entity>, - - phantom: PhantomData, -} - -impl> HashTableIter { - pub fn create( - capacity: isize, - entities: *mut Entity, - zero_entity: Option<*mut Entity>, - ) -> Self { - Self { - idx: -2, - capacity, - entities, - zero_entity, - phantom: PhantomData::default(), - } - } -} - -impl> Iterator for HashTableIter { - type Item = *mut Entity; - - fn next(&mut self) -> Option { - unsafe { - if self.idx == -2 { - self.idx = -1; - if self.zero_entity.is_some() { - return self.zero_entity; - } - } - - self.idx += 1; - while self.idx < self.capacity && self.entities.offset(self.idx).is_zero() { - self.idx += 1; - } - match self.idx == self.capacity { - true => None, - false => Some(self.entities.offset(self.idx)), - } - } - } -} - -pub struct TwoLevelHashTableIter> { - iters: Vec>, - index: usize, -} - -impl> TwoLevelHashTableIter { - pub fn create(iters: Vec>) -> Self { - Self { iters, index: 0 } - } -} - -impl> Iterator for TwoLevelHashTableIter { - type Item = *mut Entity; - fn next(&mut self) -> Option { - match self.iters[self.index].next() { - Some(x) => Some(x), - None => { - if self.index < self.iters.len() - 1 { - self.index += 1; - self.next() - } else { - None - } - } - } - } -} diff --git a/src/common/hashtable/src/hash_table_key.rs b/src/common/hashtable/src/hash_table_key.rs deleted file mode 100644 index 4b6b7727feec9..0000000000000 --- a/src/common/hashtable/src/hash_table_key.rs +++ /dev/null @@ -1,158 +0,0 @@ -// Copyright 2021 Datafuse Labs. -// -// Licensed 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 ordered_float::OrderedFloat; -use primitive_types::U256; -use primitive_types::U512; - -pub trait HashTableKeyable: Eq + Sized { - const BEFORE_EQ_HASH: bool; - - fn is_zero(&self) -> bool; - fn fast_hash(&self) -> u64; - fn set_key(&mut self, new_value: &Self); -} - -macro_rules! primitive_hasher_impl { - ($primitive_type:ty) => { - impl HashTableKeyable for $primitive_type { - const BEFORE_EQ_HASH: bool = false; - - #[inline(always)] - fn is_zero(&self) -> bool { - *self == 0 - } - - #[inline(always)] - fn fast_hash(&self) -> u64 { - let mut hash_value = *self as u64; - hash_value ^= hash_value >> 33; - hash_value = hash_value.wrapping_mul(0xff51afd7ed558ccd_u64); - hash_value ^= hash_value >> 33; - hash_value = hash_value.wrapping_mul(0xc4ceb9fe1a85ec53_u64); - hash_value ^= hash_value >> 33; - hash_value - } - - #[inline(always)] - fn set_key(&mut self, new_value: &$primitive_type) { - *self = *new_value; - } - } - }; -} - -primitive_hasher_impl!(i8); -primitive_hasher_impl!(i16); -primitive_hasher_impl!(i32); -primitive_hasher_impl!(i64); -primitive_hasher_impl!(u8); -primitive_hasher_impl!(u16); -primitive_hasher_impl!(u32); -primitive_hasher_impl!(u64); - -impl HashTableKeyable for u128 { - const BEFORE_EQ_HASH: bool = false; - #[inline(always)] - fn is_zero(&self) -> bool { - *self == 0u128 - } - - #[inline(always)] - fn fast_hash(&self) -> u64 { - let mut hash_value = *self; - hash_value ^= hash_value >> 33; - hash_value = hash_value.wrapping_mul(0xff51afd7ed558ccd_u128); - hash_value ^= hash_value >> 33; - hash_value = hash_value.wrapping_mul(0xc4ceb9fe1a85ec53_u128); - hash_value ^= hash_value >> 33; - - hash_value as u64 - } - #[inline(always)] - fn set_key(&mut self, new_value: &u128) { - *self = *new_value; - } -} - -impl HashTableKeyable for U256 { - const BEFORE_EQ_HASH: bool = false; - #[inline(always)] - fn is_zero(&self) -> bool { - self.is_zero() - } - #[inline(always)] - fn fast_hash(&self) -> u64 { - self.low_u128().fast_hash() ^ (*self >> 128).low_u128().fast_hash() - } - #[inline(always)] - fn set_key(&mut self, new_value: &U256) { - *self = *new_value; - } -} - -impl HashTableKeyable for U512 { - const BEFORE_EQ_HASH: bool = false; - #[inline(always)] - fn is_zero(&self) -> bool { - self.is_zero() - } - #[inline(always)] - fn fast_hash(&self) -> u64 { - self.low_u128().fast_hash() - ^ (*self >> 128).low_u128().fast_hash() - ^ (*self >> 256).low_u128().fast_hash() - ^ (*self >> 384).low_u128().fast_hash() - } - #[inline(always)] - fn set_key(&mut self, new_value: &U512) { - *self = *new_value; - } -} - -impl HashTableKeyable for OrderedFloat { - const BEFORE_EQ_HASH: bool = false; - - #[inline(always)] - fn is_zero(&self) -> bool { - self.0 == 0.0 - } - - #[inline(always)] - fn fast_hash(&self) -> u64 { - self.to_bits() as u64 - } - #[inline(always)] - fn set_key(&mut self, new_value: &OrderedFloat) { - *self = *new_value; - } -} - -impl HashTableKeyable for OrderedFloat { - const BEFORE_EQ_HASH: bool = false; - - #[inline(always)] - fn is_zero(&self) -> bool { - self.0 == 0.0 - } - - #[inline(always)] - fn fast_hash(&self) -> u64 { - self.to_bits() as u64 - } - #[inline(always)] - fn set_key(&mut self, new_value: &OrderedFloat) { - *self = *new_value; - } -} diff --git a/src/common/hashtable/src/hashtable.rs b/src/common/hashtable/src/hashtable.rs new file mode 100644 index 0000000000000..4ef781cdc11fa --- /dev/null +++ b/src/common/hashtable/src/hashtable.rs @@ -0,0 +1,295 @@ +// Copyright 2021 Datafuse Labs. +// +// Licensed 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 std::alloc::Allocator; +use std::intrinsics::unlikely; +use std::mem::MaybeUninit; + +use common_base::mem_allocator::GlobalAllocator; +use common_base::mem_allocator::MmapAllocator; + +use super::container::HeapContainer; +use super::table0::Entry; +use super::table0::Table0; +use super::table0::Table0Iter; +use super::table0::Table0IterMut; +use super::traits::HashtableLike; +use super::traits::Keyable; +use super::utils::ZeroEntry; + +pub struct Hashtable> +where + K: Keyable, + A: Allocator + Clone, +{ + pub(crate) zero: ZeroEntry, + pub(crate) table: Table0, A>, A>, +} + +unsafe impl Send for Hashtable {} + +unsafe impl Sync for Hashtable {} + +impl Hashtable +where + K: Keyable, + A: Allocator + Clone + Default, +{ + pub fn new() -> Self { + Self::new_in(Default::default()) + } + pub fn with_capacity(capacity: usize) -> Self { + Self::with_capacity_in(capacity, Default::default()) + } +} + +impl Default for Hashtable +where + K: Keyable, + A: Allocator + Clone + Default, +{ + fn default() -> Self { + Self::new() + } +} + +impl Hashtable +where + K: Keyable, + A: Allocator + Clone, +{ + pub fn new_in(allocator: A) -> Self { + Self::with_capacity_in(256, allocator) + } + pub fn with_capacity_in(capacity: usize, allocator: A) -> Self { + Self { + table: Table0::with_capacity_in(capacity, allocator), + zero: ZeroEntry(None), + } + } + #[inline(always)] + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + #[inline(always)] + pub fn len(&self) -> usize { + self.zero.is_some() as usize + self.table.len() + } + #[inline(always)] + pub fn capacity(&self) -> usize { + self.zero.is_some() as usize + self.table.capacity() + } + #[inline(always)] + pub fn entry(&self, key: &K) -> Option<&Entry> { + if unlikely(K::equals_zero(key)) { + if let Some(entry) = self.zero.as_ref() { + return Some(entry); + } else { + return None; + } + } + unsafe { self.table.get(key) } + } + #[inline(always)] + pub fn get(&self, key: &K) -> Option<&V> { + unsafe { self.entry(key).map(|e| e.val.assume_init_ref()) } + } + #[inline(always)] + pub fn entry_mut(&mut self, key: &K) -> Option<&mut Entry> { + if unlikely(K::equals_zero(key)) { + if let Some(entry) = self.zero.as_mut() { + return Some(entry); + } else { + return None; + } + } + unsafe { self.table.get_mut(key) } + } + #[inline(always)] + pub fn get_mut(&mut self, key: &K) -> Option<&mut V> { + unsafe { self.entry_mut(key).map(|e| e.val.assume_init_mut()) } + } + #[inline(always)] + pub fn contains(&self, key: &K) -> bool { + self.get(key).is_some() + } + /// # Safety + /// + /// The uninitialized value of returned entry should be written immediately. + #[inline(always)] + pub unsafe fn insert_and_entry( + &mut self, + key: K, + ) -> Result<&mut Entry, &mut Entry> { + if unlikely(K::equals_zero(&key)) { + let res = self.zero.is_some(); + if !res { + *self.zero = Some(MaybeUninit::zeroed().assume_init()); + } + let zero = self.zero.as_mut().unwrap(); + if res { + return Err(zero); + } else { + return Ok(zero); + } + } + if unlikely((self.table.len() + 1) * 2 > self.table.capacity()) { + if (self.table.entries.len() >> 22) == 0 { + self.table.grow(2); + } else { + self.table.grow(1); + } + } + self.table.insert(key) + } + /// # Safety + /// + /// The returned uninitialized value should be written immediately. + #[inline(always)] + pub unsafe fn insert(&mut self, key: K) -> Result<&mut MaybeUninit, &mut V> { + match self.insert_and_entry(key) { + Ok(e) => Ok(&mut e.val), + Err(e) => Err(e.val.assume_init_mut()), + } + } + pub fn iter(&self) -> HashtableIter<'_, K, V> { + HashtableIter { + inner: self.zero.iter().chain(self.table.iter()), + } + } + pub fn iter_mut(&mut self) -> HashtableIterMut<'_, K, V> { + HashtableIterMut { + inner: self.zero.iter_mut().chain(self.table.iter_mut()), + } + } +} + +impl Hashtable +where + K: Keyable, + A: Allocator + Clone, +{ + #[inline(always)] + pub fn set_insert(&mut self, key: K) -> Result<&mut MaybeUninit<()>, &mut ()> { + unsafe { self.insert(key) } + } + #[inline(always)] + pub fn set_merge(&mut self, other: &Self) { + if let Some(entry) = other.zero.0.as_ref() { + self.zero = ZeroEntry(Some(Entry { + key: entry.key, + val: MaybeUninit::uninit(), + _alignment: [0; 0], + })); + } + while (self.table.len() + other.table.len()) * 2 > self.table.capacity() { + if (self.table.entries.len() >> 22) == 0 { + self.table.grow(2); + } else { + self.table.grow(1); + } + } + unsafe { + self.table.set_merge(&other.table); + } + } +} + +pub struct HashtableIter<'a, K, V> { + inner: std::iter::Chain>, Table0Iter<'a, K, V>>, +} + +impl<'a, K, V> Iterator for HashtableIter<'a, K, V> +where K: Keyable +{ + type Item = &'a Entry; + + fn next(&mut self) -> Option { + self.inner.next() + } +} + +pub struct HashtableIterMut<'a, K, V> { + inner: std::iter::Chain>, Table0IterMut<'a, K, V>>, +} + +impl<'a, K, V> Iterator for HashtableIterMut<'a, K, V> +where K: Keyable +{ + type Item = &'a mut Entry; + + fn next(&mut self) -> Option { + self.inner.next() + } +} + +impl HashtableLike for Hashtable +where + K: Keyable, + A: Allocator + Clone + 'static, +{ + type Key = K; + type KeyRef<'a> = K where K:'a; + type Value = V; + + type EntryRef<'a> = &'a Entry where Self:'a, K:'a, V: 'a; + type EntryMutRef<'a> = &'a mut Entry where Self:'a, K:'a, V: 'a; + + type Iterator<'a> = HashtableIter<'a, K, V> where Self:'a, K:'a, V: 'a; + type IteratorMut<'a> = HashtableIterMut<'a, K, V> where Self:'a, K:'a, V: 'a; + + fn entry<'a>(&self, key_ref: Self::KeyRef<'a>) -> Option> + where K: 'a { + self.entry(&key_ref) + } + fn entry_mut<'a>(&mut self, key_ref: Self::KeyRef<'a>) -> Option> + where K: 'a { + self.entry_mut(&key_ref) + } + + fn get<'a>(&self, key_ref: Self::KeyRef<'a>) -> Option<&Self::Value> + where K: 'a { + self.get(&key_ref) + } + fn get_mut<'a>(&mut self, key_ref: Self::KeyRef<'a>) -> Option<&mut Self::Value> + where K: 'a { + self.get_mut(&key_ref) + } + + unsafe fn insert<'a>( + &mut self, + key_ref: Self::KeyRef<'a>, + ) -> Result<&mut MaybeUninit, &mut Self::Value> + where + K: 'a, + { + self.insert(key_ref) + } + unsafe fn insert_and_entry<'a>( + &mut self, + key_ref: Self::KeyRef<'a>, + ) -> Result, Self::EntryMutRef<'_>> + where + K: 'a, + { + self.insert_and_entry(key_ref) + } + + fn iter(&self) -> Self::Iterator<'_> { + self.iter() + } + fn iter_mut(&mut self) -> Self::IteratorMut<'_> { + self.iter_mut() + } +} diff --git a/src/common/hashtable/src/keys_ref.rs b/src/common/hashtable/src/keys_ref.rs index 19472a7c62c64..81ddfd8161406 100644 --- a/src/common/hashtable/src/keys_ref.rs +++ b/src/common/hashtable/src/keys_ref.rs @@ -12,21 +12,28 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::hash::Hash; use std::hash::Hasher; +use std::mem::MaybeUninit; -use ahash::AHasher; - -use super::HashTableKeyable; +use super::FastHash; +use super::HashtableKeyable; #[derive(Clone, Copy)] pub struct KeysRef { pub length: usize, pub address: usize, + pub hash: u64, } impl KeysRef { pub fn create(address: usize, length: usize) -> KeysRef { - KeysRef { length, address } + let hash = unsafe { std::slice::from_raw_parts(address as *const u8, length).fast_hash() }; + KeysRef { + length, + address, + hash, + } } #[allow(clippy::missing_safety_doc)] @@ -48,32 +55,22 @@ impl PartialEq for KeysRef { } } -impl HashTableKeyable for KeysRef { - const BEFORE_EQ_HASH: bool = true; - - fn is_zero(&self) -> bool { - self.length == 0 +unsafe impl HashtableKeyable for KeysRef { + fn is_zero(this: &MaybeUninit) -> bool { + unsafe { this.assume_init_ref().address == 0 } } - fn fast_hash(&self) -> u64 { - unsafe { - // TODO(Winter) We need more efficient hash algorithm - let value = std::slice::from_raw_parts(self.address as *const u8, self.length); - - let mut hasher = AHasher::default(); - hasher.write(value); - hasher.finish() - } + fn equals_zero(this: &Self) -> bool { + this.address == 0 } - fn set_key(&mut self, new_value: &Self) { - self.length = new_value.length; - self.address = new_value.address; + fn hash(&self) -> u64 { + self.hash } } -impl std::hash::Hash for KeysRef { - fn hash(&self, state: &mut H) { +impl Hash for KeysRef { + fn hash(&self, state: &mut H) { let self_value = unsafe { std::slice::from_raw_parts(self.address as *const u8, self.length) }; self_value.hash(state); diff --git a/src/common/hashtable/src/lib.rs b/src/common/hashtable/src/lib.rs index f07d0bf890b20..83c93bb1980af 100644 --- a/src/common/hashtable/src/lib.rs +++ b/src/common/hashtable/src/lib.rs @@ -13,61 +13,69 @@ // limitations under the License. #![feature(core_intrinsics)] +#![feature(allocator_api)] #![feature(arbitrary_self_types)] +#![feature(new_uninit)] +#![feature(ptr_metadata)] +#![feature(maybe_uninit_slice)] -use common_base::mem_allocator::StackfulAllocator; -pub use hash_table::HashTable; -pub use hash_table_entity::HashTableEntity; -pub use hash_table_entity::KeyValueEntity; -pub use hash_table_grower::HashTableGrower; -pub use hash_table_grower::SingleLevelGrower; -pub use hash_table_grower::TwoLevelGrower; -pub use hash_table_iter::HashTableIter; -pub use hash_table_iter::HashTableIteratorKind; -pub use hash_table_iter::TwoLevelHashTableIter; -pub use hash_table_key::HashTableKeyable; -pub use two_level_hash_table::HashTableKind; -pub use two_level_hash_table::TwoLevelHashTable; - -mod hash_set; -mod hash_table; -#[allow(clippy::missing_safety_doc, clippy::not_unsafe_ptr_arg_deref)] -mod hash_table_entity; -mod hash_table_grower; -mod hash_table_iter; -mod hash_table_key; +mod container; +mod hashtable; mod keys_ref; -mod two_level_hash_table; +mod stack_hashtable; +mod table0; +mod table1; +mod traits; +mod twolevel_hashtable; +mod unsized_hashtable; +mod utils; -pub use keys_ref::KeysRef; +pub use table0::Entry as HashtableEntry; +pub use traits::EntryMutRefLike as HashtableEntryMutRefLike; +pub use traits::EntryRefLike as HashtableEntryRefLike; +pub use traits::FastHash; +pub use traits::HashtableLike; +pub use traits::Keyable as HashtableKeyable; +pub use traits::UnsizedKeyable as HashtableUnsizedKeyable; -#[cfg(not(target_os = "linux"))] -type HashTableAllocator = common_base::mem_allocator::JEAllocator; -#[cfg(target_os = "linux")] -type HashTableAllocator = common_base::mem_allocator::MmapAllocator; +pub type Hashed = utils::Hashed; -type HashTableAllocatorWithStackMemory = - StackfulAllocator; +pub type HashMap = hashtable::Hashtable; +pub type HashMapIter<'a, K, V> = hashtable::HashtableIter<'a, K, V>; +pub type HashMapIterMut<'a, K, V> = hashtable::HashtableIter<'a, K, V>; +pub type HashSet = hashtable::Hashtable; +pub type HashSetIter<'a, K> = hashtable::HashtableIter<'a, K, ()>; +pub type HashSetIterMut<'a, K> = hashtable::HashtableIter<'a, K, ()>; -pub type HashMap = - HashTable, Grower, Allocator>; +pub type StackHashMap = stack_hashtable::StackHashtable; +pub type StackHashMapIter<'a, K, V> = stack_hashtable::StackHashtableIter<'a, K, V>; +pub type StackHashMapIterMut<'a, K, V> = stack_hashtable::StackHashtableIter<'a, K, V>; +pub type StackHashSet = stack_hashtable::StackHashtable; +pub type StackHashSetIter<'a, K> = stack_hashtable::StackHashtableIter<'a, K, ()>; +pub type StackHashSetIterMut<'a, K> = stack_hashtable::StackHashtableIter<'a, K, ()>; -pub type HashSet = - HashTable, Grower, Allocator>; +pub type TwolevelHashMap = twolevel_hashtable::TwolevelHashtable; +pub type TwolevelHashMapIter<'a, K, V> = twolevel_hashtable::TwolevelHashtableIter<'a, K, V>; +pub type TwolevelHashMapIterMut<'a, K, V> = twolevel_hashtable::TwolevelHashtableIterMut<'a, K, V>; +pub type TwolevelHashSet = twolevel_hashtable::TwolevelHashtable; +pub type TwolevelHashSetIter<'a, K> = twolevel_hashtable::TwolevelHashtableIter<'a, K, ()>; +pub type TwolevelHashSetIterMut<'a, K> = twolevel_hashtable::TwolevelHashtableIterMut<'a, K, ()>; -pub type TwoLevelHashMap = - TwoLevelHashTable, Grower, Allocator>; -pub type TwoLevelHashSet = - TwoLevelHashTable, Grower, Allocator>; +pub type HashMapKind = twolevel_hashtable::HashtableKind; +pub type HashMapKindIter<'a, K, V> = twolevel_hashtable::HashtableKindIter<'a, K, V>; +pub type HashMapKindIterMut<'a, K, V> = twolevel_hashtable::HashtableKindIterMut<'a, K, V>; +pub type HashSetKind = twolevel_hashtable::HashtableKind; +pub type HashSetKindIter<'a, K> = twolevel_hashtable::HashtableKindIter<'a, K, ()>; +pub type HashSetKindIterMut<'a, K> = twolevel_hashtable::HashtableKindIterMut<'a, K, ()>; -pub type HashMapIteratorKind = HashTableIteratorKind>; -pub type HashMapKind = HashTableKind< - Key, - KeyValueEntity, - SingleLevelGrower, - TwoLevelGrower, - HashTableAllocator, ->; +pub type UnsizedHashMap = unsized_hashtable::UnsizedHashtable; +pub type UnsizedHashMapIter<'a, K, V> = unsized_hashtable::UnsizedHashtableIter<'a, K, V>; +pub type UnsizedHashMapIterMut<'a, K, V> = unsized_hashtable::UnsizedHashtableIterMut<'a, K, V>; +pub type UnsizedHashSet = unsized_hashtable::UnsizedHashtable; +pub type UnsizedHashSetIter<'a, K> = unsized_hashtable::UnsizedHashtableIter<'a, K, ()>; +pub type UnsizedHashSetIterMut<'a, K> = unsized_hashtable::UnsizedHashtableIterMut<'a, K, ()>; +pub type UnsizedHashtableEntryRef<'a, K, V> = unsized_hashtable::UnsizedHashtableEntryRef<'a, K, V>; +pub type UnsizedHashtableEntryMutRef<'a, K, V> = + unsized_hashtable::UnsizedHashtableEntryMutRef<'a, K, V>; -pub type HashSetWithStackMemory = - HashSet>; +pub use keys_ref::KeysRef; diff --git a/src/common/hashtable/src/stack_hashtable.rs b/src/common/hashtable/src/stack_hashtable.rs new file mode 100644 index 0000000000000..738b63c5a81a7 --- /dev/null +++ b/src/common/hashtable/src/stack_hashtable.rs @@ -0,0 +1,240 @@ +// Copyright 2021 Datafuse Labs. +// +// Licensed 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 std::alloc::Allocator; +use std::intrinsics::unlikely; +use std::mem::MaybeUninit; + +use common_base::mem_allocator::GlobalAllocator; +use common_base::mem_allocator::MmapAllocator; + +use super::container::StackContainer; +use super::table0::Entry; +use super::table0::Table0; +use super::table0::Table0Iter; +use super::table0::Table0IterMut; +use super::traits::Keyable; +use super::utils::ZeroEntry; + +pub struct StackHashtable> +where + K: Keyable, + A: Allocator + Clone, +{ + zero: ZeroEntry, + table: Table0, N, A>, A>, +} + +unsafe impl Send + for StackHashtable +{ +} + +unsafe impl Sync + for StackHashtable +{ +} + +impl StackHashtable +where + K: Keyable, + A: Allocator + Clone + Default, +{ + pub fn new() -> Self { + Self::new_in(Default::default()) + } + pub fn with_capacity(capacity: usize) -> Self { + Self::with_capacity_in(capacity, Default::default()) + } +} + +impl Default for StackHashtable +where + K: Keyable, + A: Allocator + Clone + Default, +{ + fn default() -> Self { + Self::new() + } +} + +impl StackHashtable +where + K: Keyable, + A: Allocator + Clone, +{ + pub fn new_in(allocator: A) -> Self { + Self::with_capacity_in(N, allocator) + } + pub fn with_capacity_in(capacity: usize, allocator: A) -> Self { + Self { + table: Table0::with_capacity_in(capacity, allocator), + zero: ZeroEntry(None), + } + } + #[inline(always)] + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + #[inline(always)] + pub fn len(&self) -> usize { + self.zero.is_some() as usize + self.table.len() + } + #[inline(always)] + pub fn capacity(&self) -> usize { + self.zero.is_some() as usize + self.table.capacity() + } + #[inline(always)] + pub fn entry(&self, key: &K) -> Option<&Entry> { + if unlikely(K::equals_zero(key)) { + if let Some(entry) = self.zero.as_ref() { + return Some(entry); + } else { + return None; + } + } + unsafe { self.table.get(key) } + } + #[inline(always)] + pub fn get(&self, key: &K) -> Option<&V> { + unsafe { self.entry(key).map(|e| e.val.assume_init_ref()) } + } + #[inline(always)] + pub fn entry_mut(&mut self, key: &K) -> Option<&mut Entry> { + if unlikely(K::equals_zero(key)) { + if let Some(entry) = self.zero.as_mut() { + return Some(entry); + } else { + return None; + } + } + unsafe { self.table.get_mut(key) } + } + #[inline(always)] + pub fn get_mut(&mut self, key: &K) -> Option<&mut V> { + unsafe { self.entry_mut(key).map(|e| e.val.assume_init_mut()) } + } + #[inline(always)] + pub fn contains(&self, key: &K) -> bool { + self.get(key).is_some() + } + /// # Safety + /// + /// The uninitialized value of returned entry should be written immediately. + #[inline(always)] + pub unsafe fn insert_and_entry( + &mut self, + key: K, + ) -> Result<&mut Entry, &mut Entry> { + if unlikely(K::equals_zero(&key)) { + let res = self.zero.is_some(); + if !res { + *self.zero = Some(MaybeUninit::zeroed().assume_init()); + } + let zero = self.zero.as_mut().unwrap(); + if res { + return Err(zero); + } else { + return Ok(zero); + } + } + if unlikely((self.table.len() + 1) * 2 > self.table.capacity()) { + if (self.table.entries.len() >> 22) == 0 { + self.table.grow(2); + } else { + self.table.grow(1); + } + } + self.table.insert(key) + } + /// # Safety + /// + /// The returned uninitialized value should be written immediately. + #[inline(always)] + pub unsafe fn insert(&mut self, key: K) -> Result<&mut MaybeUninit, &mut V> { + match self.insert_and_entry(key) { + Ok(e) => Ok(&mut e.val), + Err(e) => Err(e.val.assume_init_mut()), + } + } + pub fn iter(&self) -> StackHashtableIter<'_, K, V> { + StackHashtableIter { + inner: self.zero.iter().chain(self.table.iter()), + } + } + pub fn iter_mut(&mut self) -> StackHashtableIterMut<'_, K, V> { + StackHashtableIterMut { + inner: self.zero.iter_mut().chain(self.table.iter_mut()), + } + } +} + +impl StackHashtable +where + K: Keyable, + A: Allocator + Clone, +{ + #[inline(always)] + pub fn set_insert(&mut self, key: K) -> Result<&mut MaybeUninit<()>, &mut ()> { + unsafe { self.insert(key) } + } + #[inline(always)] + pub fn set_merge(&mut self, other: &Self) { + if let Some(entry) = other.zero.0.as_ref() { + self.zero = ZeroEntry(Some(Entry { + key: entry.key, + val: MaybeUninit::uninit(), + _alignment: [0; 0], + })); + } + while (self.table.len() + other.table.len()) * 2 > self.table.capacity() { + if (self.table.entries.len() >> 22) == 0 { + self.table.grow(2); + } else { + self.table.grow(1); + } + } + unsafe { + self.table.set_merge(&other.table); + } + } +} + +pub struct StackHashtableIter<'a, K, V> { + inner: std::iter::Chain>, Table0Iter<'a, K, V>>, +} + +impl<'a, K, V> Iterator for StackHashtableIter<'a, K, V> +where K: Keyable +{ + type Item = &'a Entry; + + fn next(&mut self) -> Option { + self.inner.next() + } +} + +pub struct StackHashtableIterMut<'a, K, V> { + inner: std::iter::Chain>, Table0IterMut<'a, K, V>>, +} + +impl<'a, K, V> Iterator for StackHashtableIterMut<'a, K, V> +where K: Keyable +{ + type Item = &'a mut Entry; + + fn next(&mut self) -> Option { + self.inner.next() + } +} diff --git a/src/common/hashtable/src/table0.rs b/src/common/hashtable/src/table0.rs new file mode 100644 index 0000000000000..addfb0766b473 --- /dev/null +++ b/src/common/hashtable/src/table0.rs @@ -0,0 +1,377 @@ +// Copyright 2021 Datafuse Labs. +// +// Licensed 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 std::alloc::Allocator; +use std::borrow::Borrow; +use std::intrinsics::assume; +use std::mem::MaybeUninit; + +use super::container::Container; +use super::traits::EntryMutRefLike; +use super::traits::EntryRefLike; +use super::traits::Keyable; + +pub struct Entry { + pub(crate) _alignment: [u64; 0], + pub(crate) key: MaybeUninit, + pub(crate) val: MaybeUninit, +} + +impl Entry { + #[inline(always)] + pub(crate) fn is_zero(&self) -> bool { + K::is_zero(&self.key) + } + // this function can only be used in external crates + #[inline(always)] + pub fn key(&self) -> &K { + unsafe { self.key.assume_init_ref() } + } + // this function can only be used in external crates + /// # Safety + /// + /// The new key should be equals the old key. + #[inline(always)] + pub unsafe fn set_key(&mut self, key: K) { + self.key.write(key); + } + // this function can only be used in external crates + #[inline(always)] + pub fn get(&self) -> &V { + unsafe { self.val.assume_init_ref() } + } + // this function can only be used in external crates + #[inline(always)] + pub fn get_mut(&mut self) -> &mut V { + unsafe { self.val.assume_init_mut() } + } + // this function can only be used in external crates + #[inline(always)] + pub fn write(&mut self, val: V) { + self.val.write(val); + } +} + +pub struct Table0 +where + K: Keyable, + C: Container, A = A>, + A: Allocator + Clone, +{ + pub(crate) len: usize, + pub(crate) allocator: A, + pub(crate) entries: C, + pub(crate) dropped: bool, +} + +impl Table0 +where + K: Keyable, + C: Container, A = A>, + A: Allocator + Clone, +{ + pub fn with_capacity_in(capacity: usize, allocator: A) -> Self { + Self { + entries: unsafe { + C::new_zeroed( + std::cmp::max(8, capacity.next_power_of_two()), + allocator.clone(), + ) + }, + len: 0, + allocator, + dropped: false, + } + } + #[inline(always)] + pub fn len(&self) -> usize { + self.len + } + #[inline(always)] + pub fn capacity(&self) -> usize { + self.entries.len() + } + /// # Safety + /// + /// `key` doesn't equal to zero. + #[inline(always)] + pub unsafe fn get(&self, key: &K) -> Option<&Entry> { + self.get_with_hash(key, key.hash()) + } + /// # Safety + /// + /// `key` doesn't equal to zero. + /// Provided hash is correct. + #[inline(always)] + pub unsafe fn get_with_hash(&self, key: &K, hash: u64) -> Option<&Entry> { + assume(!K::equals_zero(key)); + let index = (hash as usize) & (self.entries.len() - 1); + for i in (index..self.entries.len()).chain(0..index) { + assume(i < self.entries.len()); + if self.entries[i].is_zero() { + return None; + } + if self.entries[i].key.assume_init_ref().borrow() == key { + return Some(&self.entries[i]); + } + } + None + } + /// # Safety + /// + /// `key` doesn't equal to zero. + #[inline(always)] + pub unsafe fn get_mut(&mut self, key: &K) -> Option<&mut Entry> { + self.get_with_hash_mut(key, key.hash()) + } + /// # Safety + /// + /// `key` doesn't equal to zero. + /// Provided hash is correct. + #[inline(always)] + pub unsafe fn get_with_hash_mut(&mut self, key: &K, hash: u64) -> Option<&mut Entry> { + assume(!K::equals_zero(key)); + let index = (hash as usize) & (self.entries.len() - 1); + for i in (index..self.entries.len()).chain(0..index) { + assume(i < self.entries.len()); + if self.entries[i].is_zero() { + return None; + } + if self.entries[i].key.assume_init_ref().borrow() == key { + return Some(&mut self.entries[i]); + } + } + None + } + /// # Safety + /// + /// `key` doesn't equal to zero. + /// The resulted `MaybeUninit` should be initialized immedidately. + /// + /// # Panics + /// + /// Panics if the hash table overflows. + #[inline(always)] + pub unsafe fn insert(&mut self, key: K) -> Result<&mut Entry, &mut Entry> { + self.insert_with_hash(key, key.hash()) + } + /// # Safety + /// + /// `key` doesn't equal to zero. + /// The resulted `MaybeUninit` should be initialized immedidately. + /// Provided hash is correct. + /// + /// # Panics + /// The hashtable is full. + #[inline(always)] + pub unsafe fn insert_with_hash( + &mut self, + key: K, + hash: u64, + ) -> Result<&mut Entry, &mut Entry> { + assume(!K::equals_zero(&key)); + let index = (hash as usize) & (self.entries.len() - 1); + for i in (index..self.entries.len()).chain(0..index) { + assume(i < self.entries.len()); + if self.entries[i].is_zero() { + self.len += 1; + self.entries[i].key.write(key); + return Ok(&mut self.entries[i]); + } + if self.entries[i].key.assume_init_ref() == &key { + return Err(&mut self.entries[i]); + } + } + panic!("the hash table overflows") + } + pub fn iter(&self) -> Table0Iter<'_, K, V> { + Table0Iter { + slice: self.entries.as_ref(), + i: 0, + } + } + pub fn iter_mut(&mut self) -> Table0IterMut<'_, K, V> { + Table0IterMut { + slice: self.entries.as_mut(), + i: 0, + } + } + pub fn grow(&mut self, shift: u8) { + let old_capacity = self.entries.len(); + let new_capacity = self.entries.len() << shift; + unsafe { + self.entries.grow_zeroed(new_capacity); + } + for i in 0..old_capacity { + unsafe { + assume(i < self.entries.len()); + } + if K::is_zero(&self.entries[i].key) { + continue; + } + let key = unsafe { self.entries[i].key.assume_init_ref() }; + let hash = K::hash(key); + let index = (hash as usize) & (self.entries.len() - 1); + for j in (index..self.entries.len()).chain(0..index) { + unsafe { + assume(j < self.entries.len()); + } + if j == i { + break; + } + if self.entries[j].is_zero() { + unsafe { + self.entries[j] = std::ptr::read(&self.entries[i]); + self.entries[i].key = MaybeUninit::zeroed(); + } + break; + } + } + } + for i in old_capacity..new_capacity { + unsafe { + assume(i < self.entries.len()); + } + if K::is_zero(&self.entries[i].key) { + break; + } + let key = unsafe { self.entries[i].key.assume_init_ref() }; + let hash = K::hash(key); + let index = (hash as usize) & (self.entries.len() - 1); + for j in (index..self.entries.len()).chain(0..index) { + unsafe { + assume(j < self.entries.len()); + } + if j == i { + break; + } + if self.entries[j].is_zero() { + unsafe { + self.entries[j] = std::ptr::read(&self.entries[i]); + self.entries[i].key = MaybeUninit::zeroed(); + } + break; + } + } + } + } +} + +impl Table0 +where + K: Keyable, + C: Container, A = A>, + A: Allocator + Clone, +{ + pub unsafe fn set_merge(&mut self, other: &Self) { + assert!(self.capacity() >= self.len() + other.len()); + for entry in other.iter() { + let key = entry.key.assume_init(); + let _ = self.insert(key); + } + } +} + +impl Drop for Table0 +where + K: Keyable, + C: Container, A = A>, + A: Allocator + Clone, +{ + fn drop(&mut self) { + if std::mem::needs_drop::() && !self.dropped { + self.iter_mut().for_each(|e| unsafe { + std::ptr::drop_in_place(e.get_mut()); + }); + } + } +} + +pub struct Table0Iter<'a, K, V> { + slice: &'a [Entry], + i: usize, +} + +impl<'a, K, V> Iterator for Table0Iter<'a, K, V> +where K: Keyable +{ + type Item = &'a Entry; + + fn next(&mut self) -> Option { + while self.i < self.slice.len() && self.slice[self.i].is_zero() { + self.i += 1; + } + if self.i == self.slice.len() { + None + } else { + let res = unsafe { &*(self.slice.as_ptr().add(self.i) as *const _) }; + self.i += 1; + Some(res) + } + } +} + +pub struct Table0IterMut<'a, K, V> { + slice: &'a mut [Entry], + i: usize, +} + +impl<'a, K, V> Iterator for Table0IterMut<'a, K, V> +where K: Keyable +{ + type Item = &'a mut Entry; + + fn next(&mut self) -> Option { + while self.i < self.slice.len() && self.slice[self.i].is_zero() { + self.i += 1; + } + if self.i == self.slice.len() { + None + } else { + let res = unsafe { &mut *(self.slice.as_ptr().add(self.i) as *mut _) }; + self.i += 1; + Some(res) + } + } +} + +impl<'a, K: Keyable, V: 'a> EntryRefLike for &'a Entry { + type KeyRef = K; + type ValueRef = &'a V; + + fn key(&self) -> Self::KeyRef { + *(*self).key() + } + fn get(&self) -> Self::ValueRef { + (*self).get() + } +} + +impl<'a, K: Keyable, V> EntryMutRefLike for &'a mut Entry { + type KeyRef = K; + type Value = V; + + fn key(&self) -> Self::KeyRef { + unsafe { self.key.assume_init() } + } + fn get(&self) -> &Self::Value { + unsafe { self.val.assume_init_ref() } + } + fn get_mut(&mut self) -> &mut Self::Value { + unsafe { self.val.assume_init_mut() } + } + fn write(&mut self, value: Self::Value) { + self.val.write(value); + } +} diff --git a/src/common/hashtable/src/table1.rs b/src/common/hashtable/src/table1.rs new file mode 100644 index 0000000000000..617a36d0eccd8 --- /dev/null +++ b/src/common/hashtable/src/table1.rs @@ -0,0 +1,147 @@ +// Copyright 2021 Datafuse Labs. +// +// Licensed 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 std::alloc::Allocator; + +use super::table0::Entry; + +type Ent = Entry<[u8; 2], V>; + +pub struct Table1 { + pub(crate) data: Box<[Entry<[u8; 2], V>; 65536], A>, + pub(crate) len: usize, +} + +impl Table1 { + pub fn new_in(allocator: A) -> Self { + Self { + data: unsafe { + let mut res = + Box::<[Entry<[u8; 2], V>; 65536], A>::new_zeroed_in(allocator).assume_init(); + res[0].key.write([0xff, 0xff]); + res + }, + len: 0, + } + } + pub fn capacity(&self) -> usize { + 65536 + } + pub fn len(&self) -> usize { + self.len + } + pub fn get(&self, key: [u8; 2]) -> Option<&Ent> { + let e = &self.data[key[1] as usize * 256 + key[0] as usize]; + if unsafe { e.key.assume_init() } == key { + Some(e) + } else { + None + } + } + pub fn get_mut(&mut self, key: [u8; 2]) -> Option<&mut Ent> { + let e = &mut self.data[key[1] as usize * 256 + key[0] as usize]; + if unsafe { e.key.assume_init() } == key { + Some(e) + } else { + None + } + } + /// # Safety + /// + /// The resulted `MaybeUninit` should be initialized immedidately. + pub fn insert(&mut self, key: [u8; 2]) -> Result<&mut Ent, &mut Ent> { + let e = &mut self.data[key[1] as usize * 256 + key[0] as usize]; + if unsafe { e.key.assume_init() } == key { + Err(e) + } else { + self.len += 1; + e.key.write(key); + Ok(e) + } + } + pub fn iter(&self) -> Table1Iter<'_, V> { + Table1Iter { + slice: self.data.as_ref(), + i: 0, + } + } + pub fn iter_mut(&mut self) -> Table1IterMut<'_, V> { + Table1IterMut { + slice: self.data.as_mut(), + i: 0, + } + } +} + +impl Drop for Table1 { + fn drop(&mut self) { + if std::mem::needs_drop::() { + self.iter_mut().for_each(|e| unsafe { + e.val.assume_init_drop(); + }); + } + } +} + +pub struct Table1Iter<'a, V> { + slice: &'a [Entry<[u8; 2], V>; 65536], + i: usize, +} + +impl<'a, V> Iterator for Table1Iter<'a, V> { + type Item = &'a Entry<[u8; 2], V>; + + fn next(&mut self) -> Option { + while self.i < 65536 + && unsafe { + u16::from_le_bytes(self.slice[self.i].key.assume_init()) as usize != self.i + } + { + self.i += 1; + } + if self.i == 65536 { + None + } else { + let res = unsafe { &*(self.slice.as_ptr().add(self.i) as *const _) }; + self.i += 1; + Some(res) + } + } +} + +pub struct Table1IterMut<'a, V> { + slice: &'a mut [Entry<[u8; 2], V>; 65536], + i: usize, +} + +impl<'a, V> Iterator for Table1IterMut<'a, V> { + type Item = &'a mut Entry<[u8; 2], V>; + + fn next(&mut self) -> Option { + while self.i < 65536 + && unsafe { + u16::from_le_bytes(self.slice[self.i].key.assume_init()) as usize != self.i + } + { + self.i += 1; + } + if self.i == 65536 { + None + } else { + let res = unsafe { &mut *(self.slice.as_ptr().add(self.i) as *mut _) }; + self.i += 1; + Some(res) + } + } +} diff --git a/src/common/hashtable/src/traits.rs b/src/common/hashtable/src/traits.rs new file mode 100644 index 0000000000000..d2d822618a40c --- /dev/null +++ b/src/common/hashtable/src/traits.rs @@ -0,0 +1,473 @@ +// Copyright 2021 Datafuse Labs. +// +// Licensed 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 std::mem::MaybeUninit; +use std::num::NonZeroU64; + +use ordered_float::OrderedFloat; +use primitive_types::U256; +use primitive_types::U512; + +/// # Safety +/// +/// All functions must be implemented correctly. +pub unsafe trait Keyable: Sized + Copy + Eq { + fn is_zero(this: &MaybeUninit) -> bool; + + fn equals_zero(this: &Self) -> bool; + + fn hash(&self) -> u64; +} + +/// # Safety +/// +/// There shouldn't be any uninitialized bytes or interior mutability in `Self`. +/// The implementation promises `from_bytes(as_bytes(self))` is `self`. +/// `as_bytes` should returns a slice that is valid for reads. +pub unsafe trait UnsizedKeyable { + fn as_bytes(&self) -> &[u8]; + + /// # Safety + /// The argument should be a return value of `as_bytes` in the same implementation. + unsafe fn from_bytes(bytes: &[u8]) -> &Self; +} + +macro_rules! impl_key_for_primitive_types { + ($t: ty) => { + unsafe impl Keyable for $t { + #[inline(always)] + fn equals_zero(this: &Self) -> bool { + *this == 0 + } + + #[inline(always)] + fn is_zero(this: &MaybeUninit) -> bool { + unsafe { this.assume_init() == 0 } + } + + #[inline(always)] + fn hash(&self) -> u64 { + self.fast_hash() + } + } + }; +} + +impl_key_for_primitive_types!(u8); +impl_key_for_primitive_types!(i8); +impl_key_for_primitive_types!(u16); +impl_key_for_primitive_types!(i16); +impl_key_for_primitive_types!(u32); +impl_key_for_primitive_types!(i32); +impl_key_for_primitive_types!(u64); +impl_key_for_primitive_types!(i64); +impl_key_for_primitive_types!(u128); +impl_key_for_primitive_types!(i128); + +unsafe impl Keyable for U256 { + #[inline(always)] + fn equals_zero(this: &Self) -> bool { + U256::is_zero(this) + } + + #[inline(always)] + fn is_zero(this: &MaybeUninit) -> bool { + U256::is_zero(unsafe { this.assume_init_ref() }) + } + + #[inline(always)] + fn hash(&self) -> u64 { + self.fast_hash() + } +} + +unsafe impl Keyable for U512 { + #[inline(always)] + fn equals_zero(this: &Self) -> bool { + U512::is_zero(this) + } + + #[inline(always)] + fn is_zero(this: &MaybeUninit) -> bool { + U512::is_zero(unsafe { this.assume_init_ref() }) + } + + #[inline(always)] + fn hash(&self) -> u64 { + self.fast_hash() + } +} + +unsafe impl Keyable for OrderedFloat { + #[inline(always)] + fn equals_zero(this: &Self) -> bool { + *this == 0.0 + } + + #[inline(always)] + fn is_zero(this: &MaybeUninit) -> bool { + unsafe { this.assume_init() == 0.0 } + } + + #[inline(always)] + fn hash(&self) -> u64 { + self.fast_hash() + } +} + +unsafe impl Keyable for OrderedFloat { + #[inline(always)] + fn equals_zero(this: &Self) -> bool { + *this == 0.0 + } + + #[inline(always)] + fn is_zero(this: &MaybeUninit) -> bool { + unsafe { this.assume_init() == 0.0 } + } + + #[inline(always)] + fn hash(&self) -> u64 { + self.fast_hash() + } +} + +unsafe impl Keyable for [u8; N] { + #[inline(always)] + fn equals_zero(this: &Self) -> bool { + *this == [0; N] + } + + #[inline(always)] + fn is_zero(this: &MaybeUninit) -> bool { + unsafe { this.assume_init() == [0; N] } + } + + #[inline(always)] + fn hash(&self) -> u64 { + (self as &[u8]).fast_hash() + } +} + +unsafe impl UnsizedKeyable for [u8] { + fn as_bytes(&self) -> &[u8] { + self + } + + unsafe fn from_bytes(bytes: &[u8]) -> &Self { + bytes + } +} + +unsafe impl UnsizedKeyable for str { + fn as_bytes(&self) -> &[u8] { + self.as_bytes() + } + + unsafe fn from_bytes(bytes: &[u8]) -> &Self { + std::str::from_utf8_unchecked(bytes) + } +} + +pub trait FastHash { + fn fast_hash(&self) -> u64; +} + +#[allow(dead_code)] +const CRC_A: u32 = u32::MAX; +#[allow(dead_code)] +const CRC_B: u32 = 0; + +macro_rules! impl_fast_hash_for_primitive_types { + ($t: ty) => { + impl FastHash for $t { + #[inline(always)] + fn fast_hash(&self) -> u64 { + cfg_if::cfg_if! { + if #[cfg(target_feature = "sse4.2")] { + use std::arch::x86_64::_mm_crc32_u64; + let mut high = CRC_A; + let mut low = CRC_B; + high = unsafe { _mm_crc32_u64(high as u64, *self as u64) as u32 }; + low = unsafe { _mm_crc32_u64(low as u64, *self as u64) as u32 }; + (high as u64) << 32 | low as u64 + } else { + let mut hasher = *self as u64; + hasher ^= hasher >> 33; + hasher = hasher.wrapping_mul(0xff51afd7ed558ccd_u64); + hasher ^= hasher >> 33; + hasher = hasher.wrapping_mul(0xc4ceb9fe1a85ec53_u64); + hasher ^= hasher >> 33; + hasher + } + } + } + } + }; +} + +impl_fast_hash_for_primitive_types!(u8); +impl_fast_hash_for_primitive_types!(i8); +impl_fast_hash_for_primitive_types!(u16); +impl_fast_hash_for_primitive_types!(i16); +impl_fast_hash_for_primitive_types!(u32); +impl_fast_hash_for_primitive_types!(i32); +impl_fast_hash_for_primitive_types!(u64); +impl_fast_hash_for_primitive_types!(i64); + +impl FastHash for u128 { + #[inline(always)] + fn fast_hash(&self) -> u64 { + cfg_if::cfg_if! { + if #[cfg(target_feature = "sse4.2")] { + use std::arch::x86_64::_mm_crc32_u64; + let mut high = CRC_A; + let mut low = CRC_B; + let y = [*self as u64, (*self >> 64) as u64]; + for x in y { + high = unsafe { _mm_crc32_u64(high as u64, x) as u32 }; + low = unsafe { _mm_crc32_u64(low as u64, x) as u32 }; + } + (high as u64) << 32 | low as u64 + } else { + use std::hash::Hasher; + let mut hasher = ahash::AHasher::default(); + hasher.write_u128(*self); + hasher.finish() + } + } + } +} + +impl FastHash for i128 { + #[inline(always)] + fn fast_hash(&self) -> u64 { + (*self as u128).fast_hash() + } +} + +impl FastHash for U256 { + #[inline(always)] + fn fast_hash(&self) -> u64 { + cfg_if::cfg_if! { + if #[cfg(target_feature = "sse4.2")] { + use std::arch::x86_64::_mm_crc32_u64; + let mut high = CRC_A; + let mut low = CRC_B; + for x in self.0 { + high = unsafe { _mm_crc32_u64(high as u64, x) as u32 }; + low = unsafe { _mm_crc32_u64(low as u64, x) as u32 }; + } + (high as u64) << 32 | low as u64 + } else { + use std::hash::Hasher; + let mut hasher = ahash::AHasher::default(); + for x in self.0 { + hasher.write_u64(x); + } + hasher.finish() + } + } + } +} + +impl FastHash for U512 { + #[inline(always)] + fn fast_hash(&self) -> u64 { + cfg_if::cfg_if! { + if #[cfg(target_feature = "sse4.2")] { + use std::arch::x86_64::_mm_crc32_u64; + let mut high = CRC_A; + let mut low = CRC_B; + for x in self.0 { + high = unsafe { _mm_crc32_u64(high as u64, x) as u32 }; + low = unsafe { _mm_crc32_u64(low as u64, x) as u32 }; + } + (high as u64) << 32 | low as u64 + } else { + use std::hash::Hasher; + let mut hasher = ahash::AHasher::default(); + for x in self.0 { + hasher.write_u64(x); + } + hasher.finish() + } + } + } +} + +impl FastHash for OrderedFloat { + #[inline(always)] + fn fast_hash(&self) -> u64 { + if self.is_nan() { + f32::NAN.to_bits().fast_hash() + } else { + self.to_bits().fast_hash() + } + } +} + +impl FastHash for OrderedFloat { + #[inline(always)] + fn fast_hash(&self) -> u64 { + if self.is_nan() { + f64::NAN.to_bits().fast_hash() + } else { + self.to_bits().fast_hash() + } + } +} + +impl FastHash for [u8] { + #[inline(always)] + fn fast_hash(&self) -> u64 { + cfg_if::cfg_if! { + if #[cfg(target_feature = "sse4.2")] { + use crate::utils::read_le; + use std::arch::x86_64::_mm_crc32_u64; + let mut high = CRC_A; + let mut low = CRC_B; + for i in (0..self.len()).step_by(8) { + if i + 8 < self.len() { + unsafe { + let x = (&self[i] as *const u8 as *const u64).read_unaligned(); + high = _mm_crc32_u64(high as u64, x) as u32; + low = _mm_crc32_u64(low as u64, x) as u32; + } + } else { + unsafe { + let x = read_le(&self[i] as *const u8, self.len() - i); + high = _mm_crc32_u64(high as u64, x) as u32; + low = _mm_crc32_u64(low as u64, x) as u32; + } + } + } + (high as u64) << 32 | low as u64 + } else { + use std::hash::Hasher; + let mut hasher = ahash::AHasher::default(); + hasher.write(self); + hasher.finish() + } + } + } +} + +// trick for unsized_hashtable +impl FastHash for ([u64; N], NonZeroU64) { + #[inline(always)] + fn fast_hash(&self) -> u64 { + cfg_if::cfg_if! { + if #[cfg(target_feature = "sse4.2")] { + use std::arch::x86_64::_mm_crc32_u64; + let mut high = CRC_A; + let mut low = CRC_B; + for x in self.0 { + high = unsafe { _mm_crc32_u64(high as u64, x) as u32 }; + low = unsafe { _mm_crc32_u64(low as u64, x) as u32 }; + } + high = unsafe { _mm_crc32_u64(high as u64, self.1.get()) as u32 }; + low = unsafe { _mm_crc32_u64(low as u64, self.1.get()) as u32 }; + (high as u64) << 32 | low as u64 + } else { + use std::hash::Hasher; + let mut hasher = ahash::AHasher::default(); + for x in self.0 { + hasher.write_u64(x); + } + hasher.write_u64(self.1.get()); + hasher.finish() + } + } + } +} + +pub trait EntryRefLike: Copy { + type KeyRef; + type ValueRef; + + fn key(&self) -> Self::KeyRef; + fn get(&self) -> Self::ValueRef; +} + +pub trait EntryMutRefLike { + type KeyRef; + type Value; + + fn key(&self) -> Self::KeyRef; + fn get(&self) -> &Self::Value; + fn get_mut(&mut self) -> &mut Self::Value; + fn write(&mut self, value: Self::Value); +} + +pub trait HashtableLike { + type Key: ?Sized; + type KeyRef<'a>: Copy + where Self::Key: 'a; + type Value; + + type EntryRef<'a>: EntryRefLike, ValueRef = &'a Self::Value> + where + Self: 'a, + Self::Key: 'a, + Self::Value: 'a; + type EntryMutRef<'a>: EntryMutRefLike, Value = Self::Value> + where + Self: 'a, + Self::Key: 'a, + Self::Value: 'a; + + type Iterator<'a>: Iterator> + where + Self: 'a, + Self::Key: 'a, + Self::Value: 'a; + type IteratorMut<'a>: Iterator> + where + Self: 'a, + Self::Key: 'a, + Self::Value: 'a; + + fn entry<'a>(&self, key_ref: Self::KeyRef<'a>) -> Option> + where Self::Key: 'a; + fn entry_mut<'a>(&mut self, key_ref: Self::KeyRef<'a>) -> Option> + where Self::Key: 'a; + + fn get<'a>(&self, key_ref: Self::KeyRef<'a>) -> Option<&Self::Value> + where Self::Key: 'a; + fn get_mut<'a>(&mut self, key_ref: Self::KeyRef<'a>) -> Option<&mut Self::Value> + where Self::Key: 'a; + + /// # Safety + /// + /// The uninitialized value of returned entry should be written immediately. + unsafe fn insert<'a>( + &mut self, + key_ref: Self::KeyRef<'a>, + ) -> Result<&mut MaybeUninit, &mut Self::Value> + where + Self::Key: 'a; + /// # Safety + /// + /// The uninitialized value of returned entry should be written immediately. + unsafe fn insert_and_entry<'a>( + &mut self, + key_ref: Self::KeyRef<'a>, + ) -> Result, Self::EntryMutRef<'_>> + where + Self::Key: 'a; + + fn iter(&self) -> Self::Iterator<'_>; + fn iter_mut(&mut self) -> Self::IteratorMut<'_>; +} diff --git a/src/common/hashtable/src/two_level_hash_table.rs b/src/common/hashtable/src/two_level_hash_table.rs deleted file mode 100644 index fd0443a1cdacd..0000000000000 --- a/src/common/hashtable/src/two_level_hash_table.rs +++ /dev/null @@ -1,225 +0,0 @@ -// Copyright 2022 Datafuse Labs. -// -// Licensed 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. -// -// Reference the ClickHouse HashTable to implement the Databend HashTable - -use common_base::mem_allocator::Allocator as AllocatorTrait; - -use crate::HashTable; -use crate::HashTableEntity; -use crate::HashTableGrower; -use crate::HashTableIteratorKind; -use crate::HashTableKeyable; -use crate::TwoLevelHashTableIter; - -static BITS_FOR_BUCKET: u8 = 8; -static NUM_BUCKETS: usize = 1 << BITS_FOR_BUCKET; -static MAX_BUCKECT: usize = NUM_BUCKETS - 1; - -pub enum HashTableKind< - Key: HashTableKeyable, - Entity: HashTableEntity, - SingleLevelGrower: HashTableGrower, - TwoLevelGrower: HashTableGrower, - Allocator: AllocatorTrait, -> { - HashTable(HashTable), - TwoLevelHashTable(TwoLevelHashTable), -} - -impl< - Key: HashTableKeyable, - Entity: HashTableEntity, - SingleLevelGrower: HashTableGrower, - TwoLevelGrower: HashTableGrower, - Allocator: AllocatorTrait + Default, -> HashTableKind -{ - pub fn create_hash_table() -> Self { - Self::HashTable(HashTable::create()) - } - - pub fn create_two_level_hash_table() -> Self { - Self::TwoLevelHashTable(TwoLevelHashTable::create()) - } - - #[inline(always)] - pub fn len(&self) -> usize { - match self { - HashTableKind::HashTable(data) => data.len(), - HashTableKind::TwoLevelHashTable(data) => data.len(), - } - } - - #[inline(always)] - pub fn is_empty(&self) -> bool { - self.len() == 0 - } - - #[inline(always)] - pub fn iter(&self) -> HashTableIteratorKind { - match self { - HashTableKind::HashTable(data) => data.enum_iter(), - HashTableKind::TwoLevelHashTable(data) => data.enum_iter(), - } - } - - #[inline(always)] - pub fn insert_key(&mut self, key: &Key, inserted: &mut bool) -> *mut Entity { - match self { - HashTableKind::HashTable(data) => data.insert_key(key, inserted), - HashTableKind::TwoLevelHashTable(data) => data.insert_key(key, inserted), - } - } - - #[inline(always)] - pub fn insert_hash_key(&mut self, key: &Key, hash: u64, inserted: &mut bool) -> *mut Entity { - match self { - HashTableKind::HashTable(data) => data.insert_hash_key(key, hash, inserted), - HashTableKind::TwoLevelHashTable(data) => data.insert_hash_key(key, hash, inserted), - } - } - - #[allow(clippy::missing_safety_doc)] - pub unsafe fn convert_to_two_level(&mut self) { - let mut two_level_hash_table = Self::create_two_level_hash_table(); - if !self.is_empty() { - let mut inserted = true; - for old_entity in self.iter() { - let new_entity = - two_level_hash_table.insert_key(old_entity.get_key(), &mut inserted); - if inserted { - new_entity.swap(old_entity); - } - } - } - self.set_entity_swapped(true); - std::mem::swap(self, &mut two_level_hash_table); - } - - pub fn set_entity_swapped(&mut self, entity_swapped: bool) { - match self { - HashTableKind::HashTable(data) => data.entity_swapped = entity_swapped, - HashTableKind::TwoLevelHashTable(data) => { - for table in data.hash_tables.iter_mut() { - table.entity_swapped = entity_swapped; - } - } - } - } -} - -pub struct TwoLevelHashTable< - Key: HashTableKeyable, - Entity: HashTableEntity, - Grower: HashTableGrower, - Allocator: AllocatorTrait, -> { - hash_tables: Vec>, -} - -impl< - Key: HashTableKeyable, - Entity: HashTableEntity, - Grower: HashTableGrower, - Allocator: AllocatorTrait + Default, -> TwoLevelHashTable -{ - pub fn create() -> TwoLevelHashTable { - let mut hash_tables: Vec> = - Vec::with_capacity(NUM_BUCKETS); - - for _ in 0..NUM_BUCKETS { - hash_tables.push(HashTable::::create()); - } - - TwoLevelHashTable { hash_tables } - } - - pub fn with_capacity(capacity: usize) -> TwoLevelHashTable { - let per_capcity = (capacity / NUM_BUCKETS).max(1 << 8); - let mut hash_tables: Vec> = - Vec::with_capacity(NUM_BUCKETS); - for _ in 0..NUM_BUCKETS { - hash_tables.push(HashTable::::with_capacity( - per_capcity, - )); - } - TwoLevelHashTable { hash_tables } - } - - pub fn inner_hash_tables(&self) -> &[HashTable] { - self.hash_tables.as_slice() - } - - pub fn inner_hash_tables_mut(&mut self) -> &mut [HashTable] { - self.hash_tables.as_mut_slice() - } - - #[inline(always)] - pub fn len(&self) -> usize { - self.hash_tables - .iter() - .map(|hash_table| hash_table.len()) - .sum() - } - - #[inline(always)] - pub fn is_empty(&self) -> bool { - self.len() == 0 - } - - #[inline(always)] - pub fn enum_iter(&self) -> HashTableIteratorKind { - let mut iters = Vec::with_capacity(NUM_BUCKETS); - for i in 0..NUM_BUCKETS { - iters.push(self.hash_tables[i].iter()) - } - HashTableIteratorKind::::create_two_level_hash_table_iter(iters) - } - - #[inline(always)] - pub fn iter(&self) -> TwoLevelHashTableIter { - let mut iters = Vec::with_capacity(NUM_BUCKETS); - for i in 0..NUM_BUCKETS { - iters.push(self.hash_tables[i].iter()) - } - TwoLevelHashTableIter::::create(iters) - } - - #[inline(always)] - pub fn insert_key(&mut self, key: &Key, inserted: &mut bool) -> *mut Entity { - let hash = key.fast_hash(); - let bucket = self.get_bucket_from_hash(&hash); - self.hash_tables[bucket].insert_hash_key(key, hash, inserted) - } - - #[inline(always)] - pub fn insert_hash_key(&mut self, key: &Key, hash: u64, inserted: &mut bool) -> *mut Entity { - let bucket = self.get_bucket_from_hash(&hash); - self.hash_tables[bucket].insert_hash_key(key, hash, inserted) - } - - #[inline(always)] - pub fn find_key(&self, key: &Key) -> Option<*mut Entity> { - let hash = key.fast_hash(); - let bucket = self.get_bucket_from_hash(&hash); - self.hash_tables[bucket].find_key(key) - } - - #[inline(always)] - fn get_bucket_from_hash(&self, hash_value: &u64) -> usize { - ((hash_value >> (64 - BITS_FOR_BUCKET)) & (MAX_BUCKECT as u64)) as usize - } -} diff --git a/src/common/hashtable/src/twolevel_hashtable.rs b/src/common/hashtable/src/twolevel_hashtable.rs new file mode 100644 index 0000000000000..ad21a6d826f7a --- /dev/null +++ b/src/common/hashtable/src/twolevel_hashtable.rs @@ -0,0 +1,521 @@ +// Copyright 2021 Datafuse Labs. +// +// Licensed 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 std::alloc::Allocator; +use std::intrinsics::unlikely; +use std::mem::MaybeUninit; + +use common_base::mem_allocator::GlobalAllocator; +use common_base::mem_allocator::MmapAllocator; + +use super::container::HeapContainer; +use super::hashtable::Hashtable; +use super::hashtable::HashtableIter; +use super::hashtable::HashtableIterMut; +use super::table0::Entry; +use super::table0::Table0; +use super::table0::Table0Iter; +use super::table0::Table0IterMut; +use super::traits::Keyable; +use super::utils::ZeroEntry; + +const BUCKETS: usize = 256; +const BUCKETS_LG2: u32 = 8; + +type Tables = [Table0, A>, A>; BUCKETS]; + +pub struct TwolevelHashtable> +where + K: Keyable, + A: Allocator + Clone, +{ + zero: ZeroEntry, + tables: Tables, +} + +unsafe impl Send + for TwolevelHashtable +{ +} + +unsafe impl Sync + for TwolevelHashtable +{ +} + +impl TwolevelHashtable +where + K: Keyable, + A: Allocator + Clone + Default, +{ + pub fn new() -> Self { + Self::new_in(Default::default()) + } +} + +impl Default for TwolevelHashtable +where + K: Keyable, + A: Allocator + Clone + Default, +{ + fn default() -> Self { + Self::new() + } +} + +impl TwolevelHashtable +where + K: Keyable, + A: Allocator + Clone, +{ + pub fn new_in(allocator: A) -> Self { + Self { + zero: ZeroEntry(None), + tables: std::array::from_fn(|_| Table0::with_capacity_in(256, allocator.clone())), + } + } + #[inline(always)] + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + #[inline(always)] + pub fn len(&self) -> usize { + self.zero.is_some() as usize + self.tables.iter().map(|x| x.len()).sum::() + } + #[inline(always)] + pub fn capacity(&self) -> usize { + self.zero.is_some() as usize + self.tables.iter().map(|x| x.capacity()).sum::() + } + #[inline(always)] + pub fn entry(&self, key: &K) -> Option<&Entry> { + if unlikely(K::equals_zero(key)) { + if let Some(entry) = self.zero.as_ref() { + return Some(entry); + } else { + return None; + } + } + let hash = K::hash(key); + let index = hash as usize >> (64u32 - BUCKETS_LG2); + unsafe { self.tables[index].get_with_hash(key, hash) } + } + #[inline(always)] + pub fn get(&self, key: &K) -> Option<&V> { + unsafe { self.entry(key).map(|e| e.val.assume_init_ref()) } + } + #[inline(always)] + pub fn entry_mut(&mut self, key: &K) -> Option<&mut Entry> { + if unlikely(K::equals_zero(key)) { + if let Some(entry) = self.zero.as_mut() { + return Some(entry); + } else { + return None; + } + } + let hash = K::hash(key); + let index = hash as usize >> (64u32 - BUCKETS_LG2); + unsafe { self.tables[index].get_with_hash_mut(key, hash) } + } + #[inline(always)] + pub fn get_mut(&mut self, key: &K) -> Option<&mut V> { + unsafe { self.entry_mut(key).map(|e| e.val.assume_init_mut()) } + } + #[inline(always)] + pub fn contains(&self, key: &K) -> bool { + self.get(key).is_some() + } + /// # Safety + /// + /// The uninitialized value of returned entry should be written immediately. + #[inline(always)] + pub unsafe fn insert_and_entry( + &mut self, + key: K, + ) -> Result<&mut Entry, &mut Entry> { + if unlikely(K::equals_zero(&key)) { + let res = self.zero.is_some(); + if !res { + *self.zero = Some(MaybeUninit::zeroed().assume_init()); + } + let zero = self.zero.as_mut().unwrap(); + if res { + return Err(zero); + } else { + return Ok(zero); + } + } + let hash = K::hash(&key); + let index = hash as usize >> (64u32 - BUCKETS_LG2); + if unlikely((self.tables[index].len() + 1) * 2 > self.tables[index].capacity()) { + if (self.tables[index].entries.len() >> 14) == 0 { + self.tables[index].grow(2); + } else { + self.tables[index].grow(1); + } + } + self.tables[index].insert_with_hash(key, hash) + } + /// # Safety + /// + /// The returned uninitialized value should be written immediately. + #[inline(always)] + pub unsafe fn insert(&mut self, key: K) -> Result<&mut MaybeUninit, &mut V> { + match self.insert_and_entry(key) { + Ok(e) => Ok(&mut e.val), + Err(e) => Err(e.val.assume_init_mut()), + } + } + pub fn iter(&self) -> TwolevelHashtableIter<'_, K, V, A> { + TwolevelHashtableIter { + inner: self.zero.iter().chain(self.tables.iter().flat_map( + Table0::iter + as fn( + &'_ Table0, A>, A>, + ) -> Table0Iter<'_, _, _>, + )), + } + } + pub fn iter_mut(&mut self) -> TwolevelHashtableIterMut<'_, K, V, A> { + TwolevelHashtableIterMut { + inner: self.zero.iter_mut().chain(self.tables.iter_mut().flat_map( + Table0::iter_mut + as fn( + &'_ mut Table0, A>, A>, + ) -> Table0IterMut<'_, _, _>, + )), + } + } +} + +impl From> for TwolevelHashtable +where + K: Keyable, + A: Allocator + Clone, +{ + fn from(mut src: Hashtable) -> Self { + let mut res = TwolevelHashtable::new_in(src.table.allocator.clone()); + unsafe { + src.table.dropped = true; + res.zero = ZeroEntry(src.zero.take()); + for entry in src.table.iter() { + let key = entry.key.assume_init(); + let val = std::ptr::read(entry.val.assume_init_ref()); + let hash = K::hash(&key); + let index = hash as usize >> (64u32 - BUCKETS_LG2); + if unlikely((res.tables[index].len() + 1) * 2 > res.tables[index].capacity()) { + if (res.tables[index].entries.len() >> 14) == 0 { + res.tables[index].grow(2); + } else { + res.tables[index].grow(1); + } + } + res.tables[index] + .insert_with_hash(key, hash) + .ok() + .unwrap() + .write(val); + } + } + res + } +} + +impl TwolevelHashtable +where + K: Keyable, + A: Allocator + Clone, +{ + #[inline(always)] + pub fn set_insert(&mut self, key: K) -> Result<&mut MaybeUninit<()>, &mut ()> { + unsafe { self.insert(key) } + } + #[inline(always)] + pub fn set_merge(&mut self, other: &Self) { + if let Some(entry) = other.zero.0.as_ref() { + self.zero = ZeroEntry(Some(Entry { + key: entry.key, + val: MaybeUninit::uninit(), + _alignment: [0; 0], + })); + } + for (i, table) in other.tables.iter().enumerate() { + while (self.tables[i].len() + table.len()) * 2 > self.tables[i].capacity() { + if (self.tables[i].entries.len() >> 22) == 0 { + self.tables[i].grow(2); + } else { + self.tables[i].grow(1); + } + } + unsafe { + self.tables[i].set_merge(table); + } + } + } +} + +type TwolevelHashtableIterInner<'a, K, V, A> = std::iter::Chain< + std::option::Iter<'a, Entry>, + std::iter::FlatMap< + std::slice::Iter<'a, Table0, A>, A>>, + Table0Iter<'a, K, V>, + fn(&'a Table0, A>, A>) -> Table0Iter<'a, K, V>, + >, +>; + +pub struct TwolevelHashtableIter<'a, K, V, A = MmapAllocator> +where + K: Keyable, + A: Allocator + Clone, +{ + inner: TwolevelHashtableIterInner<'a, K, V, A>, +} + +impl<'a, K, V, A> Iterator for TwolevelHashtableIter<'a, K, V, A> +where + K: Keyable, + A: Allocator + Clone, +{ + type Item = &'a Entry; + + fn next(&mut self) -> Option { + self.inner.next() + } +} + +type TwolevelHashtableIterMutInner<'a, K, V, A> = std::iter::Chain< + std::option::IterMut<'a, Entry>, + std::iter::FlatMap< + std::slice::IterMut<'a, Table0, A>, A>>, + Table0IterMut<'a, K, V>, + fn(&'a mut Table0, A>, A>) -> Table0IterMut<'a, K, V>, + >, +>; + +pub struct TwolevelHashtableIterMut<'a, K, V, A = MmapAllocator> +where + K: Keyable, + A: Allocator + Clone, +{ + inner: TwolevelHashtableIterMutInner<'a, K, V, A>, +} + +impl<'a, K, V, A> Iterator for TwolevelHashtableIterMut<'a, K, V, A> +where + K: Keyable, + A: Allocator + Clone, +{ + type Item = &'a mut Entry; + + fn next(&mut self) -> Option { + self.inner.next() + } +} + +pub enum HashtableKind> +where + K: Keyable, + A: Allocator + Clone, +{ + Onelevel(Hashtable), + Twolevel(Box>), +} + +impl HashtableKind +where + K: Keyable, + A: Allocator + Clone + Default, +{ + pub fn new() -> Self { + Self::new_in(Default::default()) + } + pub fn new_twolevel() -> Self { + Self::new_twolevel_in(Default::default()) + } +} + +impl Default for HashtableKind +where + K: Keyable, + A: Allocator + Clone + Default, +{ + fn default() -> Self { + Self::new() + } +} + +impl HashtableKind +where + K: Keyable, + A: Allocator + Clone, +{ + pub fn new_in(allocator: A) -> Self { + Self::Onelevel(Hashtable::new_in(allocator)) + } + pub fn new_twolevel_in(allocator: A) -> Self { + Self::Twolevel(Box::new(TwolevelHashtable::new_in(allocator))) + } + #[inline(always)] + pub fn is_empty(&self) -> bool { + use HashtableKind::*; + match self { + Onelevel(x) => Hashtable::is_empty(x), + Twolevel(x) => TwolevelHashtable::is_empty(x), + } + } + #[inline(always)] + pub fn len(&self) -> usize { + use HashtableKind::*; + match self { + Onelevel(x) => Hashtable::len(x), + Twolevel(x) => TwolevelHashtable::len(x), + } + } + #[inline(always)] + pub fn capacity(&self) -> usize { + use HashtableKind::*; + match self { + Onelevel(x) => Hashtable::capacity(x), + Twolevel(x) => TwolevelHashtable::capacity(x), + } + } + #[inline(always)] + pub fn entry(&self, key: &K) -> Option<&Entry> { + use HashtableKind::*; + match self { + Onelevel(x) => Hashtable::entry(x, key), + Twolevel(x) => TwolevelHashtable::entry(x, key), + } + } + #[inline(always)] + pub fn get(&self, key: &K) -> Option<&V> { + use HashtableKind::*; + match self { + Onelevel(x) => Hashtable::get(x, key), + Twolevel(x) => TwolevelHashtable::get(x, key), + } + } + #[inline(always)] + pub fn entry_mut(&mut self, key: &K) -> Option<&mut Entry> { + use HashtableKind::*; + match self { + Onelevel(x) => Hashtable::entry_mut(x, key), + Twolevel(x) => TwolevelHashtable::entry_mut(x, key), + } + } + #[inline(always)] + pub fn get_mut(&mut self, key: &K) -> Option<&mut V> { + use HashtableKind::*; + match self { + Onelevel(x) => Hashtable::get_mut(x, key), + Twolevel(x) => TwolevelHashtable::get_mut(x, key), + } + } + #[inline(always)] + pub fn contains(&self, key: &K) -> bool { + self.get(key).is_some() + } + #[inline(always)] + pub unsafe fn insert_and_entry( + &mut self, + key: K, + ) -> Result<&mut Entry, &mut Entry> { + use HashtableKind::*; + match self { + Onelevel(x) => Hashtable::insert_and_entry(x, key), + Twolevel(x) => TwolevelHashtable::insert_and_entry(x, key), + } + } + #[inline(always)] + pub unsafe fn insert(&mut self, key: K) -> Result<&mut MaybeUninit, &mut V> { + use HashtableKind::*; + match self { + Onelevel(x) => Hashtable::insert(x, key), + Twolevel(x) => TwolevelHashtable::insert(x, key), + } + } + pub fn iter(&self) -> HashtableKindIter<'_, K, V, A> { + use HashtableKind::*; + match self { + Onelevel(x) => HashtableKindIter::Onelevel(Hashtable::iter(x)), + Twolevel(x) => HashtableKindIter::Twolevel(TwolevelHashtable::iter(x)), + } + } + pub fn iter_mut(&mut self) -> HashtableKindIterMut<'_, K, V, A> { + use HashtableKind::*; + match self { + Onelevel(x) => HashtableKindIterMut::Onelevel(Hashtable::iter_mut(x)), + Twolevel(x) => HashtableKindIterMut::Twolevel(TwolevelHashtable::iter_mut(x)), + } + } + pub fn convert_to_twolevel(&mut self) { + use HashtableKind::*; + unsafe { + if let Onelevel(x) = self { + let onelevel = std::ptr::read(x); + let twolevel = Box::new(TwolevelHashtable::::from(onelevel)); + std::ptr::write(self, Twolevel(twolevel)); + } + } + } +} + +pub enum HashtableKindIter<'a, K, V, A = MmapAllocator> +where + K: Keyable, + A: Allocator + Clone, +{ + Onelevel(HashtableIter<'a, K, V>), + Twolevel(TwolevelHashtableIter<'a, K, V, A>), +} + +impl<'a, K, V, A> Iterator for HashtableKindIter<'a, K, V, A> +where + K: Keyable, + A: Allocator + Clone, +{ + type Item = &'a Entry; + + fn next(&mut self) -> Option { + use HashtableKindIter::*; + match self { + Onelevel(x) => x.next(), + Twolevel(x) => x.next(), + } + } +} + +pub enum HashtableKindIterMut<'a, K, V, A = MmapAllocator> +where + K: Keyable, + A: Allocator + Clone, +{ + Onelevel(HashtableIterMut<'a, K, V>), + Twolevel(TwolevelHashtableIterMut<'a, K, V, A>), +} + +impl<'a, K, V, A> Iterator for HashtableKindIterMut<'a, K, V, A> +where + K: Keyable, + A: Allocator + Clone, +{ + type Item = &'a mut Entry; + + fn next(&mut self) -> Option { + use HashtableKindIterMut::*; + match self { + Onelevel(x) => x.next(), + Twolevel(x) => x.next(), + } + } +} diff --git a/src/common/hashtable/src/unsized_hashtable.rs b/src/common/hashtable/src/unsized_hashtable.rs new file mode 100644 index 0000000000000..1e20a42dd5c59 --- /dev/null +++ b/src/common/hashtable/src/unsized_hashtable.rs @@ -0,0 +1,1132 @@ +// Copyright 2021 Datafuse Labs. +// +// Licensed 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 std::alloc::Allocator; +use std::intrinsics::unlikely; +use std::marker::PhantomData; +use std::mem::MaybeUninit; +use std::num::NonZeroU64; +use std::ptr::NonNull; + +use bumpalo::Bump; +use common_base::mem_allocator::GlobalAllocator; +use common_base::mem_allocator::MmapAllocator; + +use super::container::HeapContainer; +use super::table0::Entry; +use super::table0::Table0; +use super::table1::Table1; +use super::traits::EntryMutRefLike; +use super::traits::EntryRefLike; +use super::traits::FastHash; +use super::traits::HashtableLike; +use super::traits::Keyable; +use super::traits::UnsizedKeyable; +use super::utils::read_le; +use crate::table0::Table0Iter; +use crate::table0::Table0IterMut; +use crate::table1::Table1Iter; +use crate::table1::Table1IterMut; + +pub struct UnsizedHashtable> +where + K: UnsizedKeyable + ?Sized, + A: Allocator + Clone, +{ + pub(crate) arena: Bump, + pub(crate) table0: Table1, + pub(crate) table1: Table0, V, HeapContainer, V>, A>, A>, + pub(crate) table2: Table0, V, HeapContainer, V>, A>, A>, + pub(crate) table3: Table0, V, HeapContainer, V>, A>, A>, + pub(crate) table4: Table0, A>, A>, + pub(crate) _phantom: PhantomData, +} + +unsafe impl Send + for UnsizedHashtable +{ +} + +unsafe impl Sync + for UnsizedHashtable +{ +} + +impl UnsizedHashtable +where + K: UnsizedKeyable + ?Sized, + A: Allocator + Clone + Default, +{ + pub fn new() -> Self { + Self::new_in(Default::default()) + } +} + +impl Default for UnsizedHashtable +where + K: UnsizedKeyable + ?Sized, + A: Allocator + Clone + Default, +{ + fn default() -> Self { + Self::new() + } +} + +impl UnsizedHashtable +where + K: UnsizedKeyable + ?Sized, + A: Allocator + Clone + Default, +{ + /// The bump for strings doesn't allocate memory by `A`. + pub fn new_in(allocator: A) -> Self { + Self { + arena: Bump::new(), + table0: Table1::new_in(allocator.clone()), + table1: Table0::with_capacity_in(128, allocator.clone()), + table2: Table0::with_capacity_in(128, allocator.clone()), + table3: Table0::with_capacity_in(128, allocator.clone()), + table4: Table0::with_capacity_in(128, allocator), + _phantom: PhantomData, + } + } + #[inline(always)] + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + #[inline(always)] + pub fn len(&self) -> usize { + self.table0.len() + + self.table1.len() + + self.table2.len() + + self.table3.len() + + self.table4.len() + } + #[inline(always)] + pub fn capacity(&self) -> usize { + self.table0.capacity() + + self.table1.capacity() + + self.table2.capacity() + + self.table3.capacity() + + self.table4.capacity() + } + #[inline(always)] + pub fn entry(&self, key: &K) -> Option> { + let key = key.as_bytes(); + match key.len() { + _ if key.last().copied() == Some(0) => unsafe { + self.table4 + .get(&FallbackKey::new(key)) + .map(|x| UnsizedHashtableEntryRef(UnsizedHashtableEntryRefInner::Table4(x))) + }, + 0 => self.table0.get([0, 0]).map(|x| { + UnsizedHashtableEntryRef(UnsizedHashtableEntryRefInner::Table0(x, PhantomData)) + }), + 1 => self.table0.get([key[0], 0]).map(|x| { + UnsizedHashtableEntryRef(UnsizedHashtableEntryRefInner::Table0(x, PhantomData)) + }), + 2 => self.table0.get([key[0], key[1]]).map(|x| { + UnsizedHashtableEntryRef(UnsizedHashtableEntryRefInner::Table0(x, PhantomData)) + }), + 3..=8 => unsafe { + let mut t = [0u64; 1]; + t[0] = read_le(key.as_ptr(), key.len()); + let t = std::mem::transmute::<_, InlineKey<0>>(t); + self.table1 + .get(&t) + .map(|x| UnsizedHashtableEntryRef(UnsizedHashtableEntryRefInner::Table1(x))) + }, + 9..=16 => unsafe { + let mut t = [0u64; 2]; + t[0] = (key.as_ptr() as *const u64).read_unaligned(); + t[1] = read_le(key.as_ptr().offset(8), key.len() - 8); + let t = std::mem::transmute::<_, InlineKey<1>>(t); + self.table2 + .get(&t) + .map(|x| UnsizedHashtableEntryRef(UnsizedHashtableEntryRefInner::Table2(x))) + }, + 17..=24 => unsafe { + let mut t = [0u64; 3]; + t[0] = (key.as_ptr() as *const u64).read_unaligned(); + t[1] = (key.as_ptr() as *const u64).offset(1).read_unaligned(); + t[2] = read_le(key.as_ptr().offset(16), key.len() - 16); + let t = std::mem::transmute::<_, InlineKey<2>>(t); + self.table3 + .get(&t) + .map(|x| UnsizedHashtableEntryRef(UnsizedHashtableEntryRefInner::Table3(x))) + }, + _ => unsafe { + self.table4 + .get(&FallbackKey::new(key)) + .map(|x| UnsizedHashtableEntryRef(UnsizedHashtableEntryRefInner::Table4(x))) + }, + } + } + #[inline(always)] + pub fn entry_mut(&mut self, key: &K) -> Option> { + let key = key.as_bytes(); + match key.len() { + _ if key.last().copied() == Some(0) => unsafe { + self.table4.get_mut(&FallbackKey::new(key)).map(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table4(x)) + }) + }, + 0 => self.table0.get_mut([0, 0]).map(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table0( + x, + PhantomData, + )) + }), + 1 => self.table0.get_mut([key[0], 0]).map(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table0( + x, + PhantomData, + )) + }), + 2 => self.table0.get_mut([key[0], key[1]]).map(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table0( + x, + PhantomData, + )) + }), + 3..=8 => unsafe { + let mut t = [0u64; 1]; + t[0] = read_le(key.as_ptr(), key.len()); + let t = std::mem::transmute::<_, InlineKey<0>>(t); + self.table1.get_mut(&t).map(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table1(x)) + }) + }, + 9..=16 => unsafe { + let mut t = [0u64; 2]; + t[0] = (key.as_ptr() as *const u64).read_unaligned(); + t[1] = read_le(key.as_ptr().offset(8), key.len() - 8); + let t = std::mem::transmute::<_, InlineKey<1>>(t); + self.table2.get_mut(&t).map(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table2(x)) + }) + }, + 17..=24 => unsafe { + let mut t = [0u64; 3]; + t[0] = (key.as_ptr() as *const u64).read_unaligned(); + t[1] = (key.as_ptr() as *const u64).offset(1).read_unaligned(); + t[2] = read_le(key.as_ptr().offset(16), key.len() - 16); + let t = std::mem::transmute::<_, InlineKey<2>>(t); + self.table3.get_mut(&t).map(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table3(x)) + }) + }, + _ => unsafe { + self.table4.get_mut(&FallbackKey::new(key)).map(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table4(x)) + }) + }, + } + } + #[inline(always)] + pub fn get(&self, key: &K) -> Option<&V> { + self.entry(key).map(|e| e.get()) + } + #[inline(always)] + pub fn get_mut(&mut self, key: &K) -> Option<&mut V> { + self.entry_mut(key) + .map(|e| unsafe { &mut *(e.get_mut_ptr() as *mut V) }) + } + #[inline(always)] + pub fn contains(&self, key: &K) -> bool { + self.get(key).is_some() + } + /// # Safety + /// + /// The uninitialized value of returned entry should be written immediately. + #[inline(always)] + pub unsafe fn insert_and_entry( + &mut self, + key: &K, + ) -> Result, UnsizedHashtableEntryMutRef<'_, K, V>> { + let key = key.as_bytes(); + match key.len() { + _ if key.last().copied() == Some(0) => { + if unlikely((self.table4.len() + 1) * 2 > self.table4.capacity()) { + if (self.table4.entries.len() >> 22) == 0 { + self.table4.grow(2); + } else { + self.table4.grow(1); + } + } + let s = self.arena.alloc_slice_copy(key); + self.table4 + .insert(FallbackKey::new(s)) + .map(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table4(x)) + }) + .map_err(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table4(x)) + }) + } + 0 => self + .table0 + .insert([0, 0]) + .map(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table0( + x, + PhantomData, + )) + }) + .map_err(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table0( + x, + PhantomData, + )) + }), + 1 => self + .table0 + .insert([key[0], 0]) + .map(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table0( + x, + PhantomData, + )) + }) + .map_err(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table0( + x, + PhantomData, + )) + }), + 2 => self + .table0 + .insert([key[0], key[1]]) + .map(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table0( + x, + PhantomData, + )) + }) + .map_err(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table0( + x, + PhantomData, + )) + }), + 3..=8 => { + if unlikely((self.table1.len() + 1) * 2 > self.table1.capacity()) { + if (self.table1.entries.len() >> 22) == 0 { + self.table1.grow(2); + } else { + self.table1.grow(1); + } + } + let mut t = [0u64; 1]; + t[0] = read_le(key.as_ptr(), key.len()); + let t = std::mem::transmute::<_, InlineKey<0>>(t); + self.table1 + .insert(t) + .map(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table1(x)) + }) + .map_err(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table1(x)) + }) + } + 9..=16 => { + if unlikely((self.table2.len() + 1) * 2 > self.table2.capacity()) { + if (self.table2.entries.len() >> 22) == 0 { + self.table2.grow(2); + } else { + self.table2.grow(1); + } + } + let mut t = [0u64; 2]; + t[0] = (key.as_ptr() as *const u64).read_unaligned(); + t[1] = read_le(key.as_ptr().offset(8), key.len() - 8); + let t = std::mem::transmute::<_, InlineKey<1>>(t); + self.table2 + .insert(t) + .map(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table2(x)) + }) + .map_err(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table2(x)) + }) + } + 17..=24 => { + if unlikely((self.table3.len() + 1) * 2 > self.table3.capacity()) { + if (self.table3.entries.len() >> 22) == 0 { + self.table3.grow(2); + } else { + self.table3.grow(1); + } + } + let mut t = [0u64; 3]; + t[0] = (key.as_ptr() as *const u64).read_unaligned(); + t[1] = (key.as_ptr() as *const u64).offset(1).read_unaligned(); + t[2] = read_le(key.as_ptr().offset(16), key.len() - 16); + let t = std::mem::transmute::<_, InlineKey<2>>(t); + self.table3 + .insert(t) + .map(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table3(x)) + }) + .map_err(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table3(x)) + }) + } + _ => { + if unlikely((self.table4.len() + 1) * 2 > self.table4.capacity()) { + if (self.table4.entries.len() >> 22) == 0 { + self.table4.grow(2); + } else { + self.table4.grow(1); + } + } + let s = self.arena.alloc_slice_copy(key); + self.table4 + .insert(FallbackKey::new(s)) + .map(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table4(x)) + }) + .map_err(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table4(x)) + }) + } + } + } + /// # Safety + /// + /// The uninitialized value of returned entry should be written immediately. + #[inline(always)] + pub unsafe fn insert(&mut self, key: &K) -> Result<&mut MaybeUninit, &mut V> { + match self.insert_and_entry(key) { + Ok(e) => Ok(&mut *(e.get_mut_ptr() as *mut MaybeUninit)), + Err(e) => Err(&mut *e.get_mut_ptr()), + } + } + /// # Safety + /// + /// * The uninitialized value of returned entry should be written immediately. + /// * The lifetime of key lives longer than the hashtable. + #[inline(always)] + pub unsafe fn insert_and_entry_borrowing( + &mut self, + key: *const K, + ) -> Result, UnsizedHashtableEntryMutRef<'_, K, V>> { + let key = (*key).as_bytes(); + match key.len() { + _ if key.last().copied() == Some(0) => { + if unlikely((self.table4.len() + 1) * 2 > self.table4.capacity()) { + if (self.table4.entries.len() >> 22) == 0 { + self.table4.grow(2); + } else { + self.table4.grow(1); + } + } + self.table4 + .insert(FallbackKey::new(key)) + .map(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table4(x)) + }) + .map_err(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table4(x)) + }) + } + 0 => self + .table0 + .insert([0, 0]) + .map(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table0( + x, + PhantomData, + )) + }) + .map_err(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table0( + x, + PhantomData, + )) + }), + 1 => self + .table0 + .insert([key[0], 0]) + .map(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table0( + x, + PhantomData, + )) + }) + .map_err(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table0( + x, + PhantomData, + )) + }), + 2 => self + .table0 + .insert([key[0], key[1]]) + .map(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table0( + x, + PhantomData, + )) + }) + .map_err(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table0( + x, + PhantomData, + )) + }), + 3..=8 => { + if unlikely((self.table1.len() + 1) * 2 > self.table1.capacity()) { + if (self.table1.entries.len() >> 22) == 0 { + self.table1.grow(2); + } else { + self.table1.grow(1); + } + } + let mut t = [0u64; 1]; + t[0] = read_le(key.as_ptr(), key.len()); + let t = std::mem::transmute::<_, InlineKey<0>>(t); + self.table1 + .insert(t) + .map(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table1(x)) + }) + .map_err(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table1(x)) + }) + } + 9..=16 => { + if unlikely((self.table2.len() + 1) * 2 > self.table2.capacity()) { + if (self.table2.entries.len() >> 22) == 0 { + self.table2.grow(2); + } else { + self.table2.grow(1); + } + } + let mut t = [0u64; 2]; + t[0] = (key.as_ptr() as *const u64).read_unaligned(); + t[1] = read_le(key.as_ptr().offset(8), key.len() - 8); + let t = std::mem::transmute::<_, InlineKey<1>>(t); + self.table2 + .insert(t) + .map(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table2(x)) + }) + .map_err(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table2(x)) + }) + } + 17..=24 => { + if unlikely((self.table3.len() + 1) * 2 > self.table3.capacity()) { + if (self.table3.entries.len() >> 22) == 0 { + self.table3.grow(2); + } else { + self.table3.grow(1); + } + } + let mut t = [0u64; 3]; + t[0] = (key.as_ptr() as *const u64).read_unaligned(); + t[1] = (key.as_ptr() as *const u64).offset(1).read_unaligned(); + t[2] = read_le(key.as_ptr().offset(16), key.len() - 16); + let t = std::mem::transmute::<_, InlineKey<2>>(t); + self.table3 + .insert(t) + .map(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table3(x)) + }) + .map_err(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table3(x)) + }) + } + _ => { + if unlikely((self.table4.len() + 1) * 2 > self.table4.capacity()) { + if (self.table4.entries.len() >> 22) == 0 { + self.table4.grow(2); + } else { + self.table4.grow(1); + } + } + self.table4 + .insert(FallbackKey::new(key)) + .map(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table4(x)) + }) + .map_err(|x| { + UnsizedHashtableEntryMutRef(UnsizedHashtableEntryMutRefInner::Table4(x)) + }) + } + } + } + /// # Safety + /// + /// * The uninitialized value of returned entry should be written immediately. + /// * The lifetime of key lives longer than the hashtable. + #[inline(always)] + pub unsafe fn insert_borrowing(&mut self, key: &K) -> Result<&mut MaybeUninit, &mut V> { + match self.insert_and_entry_borrowing(key) { + Ok(e) => Ok(&mut *(e.get_mut_ptr() as *mut MaybeUninit)), + Err(e) => Err(&mut *e.get_mut_ptr()), + } + } + pub fn iter(&self) -> UnsizedHashtableIter<'_, K, V> { + UnsizedHashtableIter { + it_0: Some(self.table0.iter()), + it_1: Some(self.table1.iter()), + it_2: Some(self.table2.iter()), + it_3: Some(self.table3.iter()), + it_4: Some(self.table4.iter()), + _phantom: PhantomData, + } + } + pub fn iter_mut(&mut self) -> UnsizedHashtableIterMut<'_, K, V> { + UnsizedHashtableIterMut { + it_0: Some(self.table0.iter_mut()), + it_1: Some(self.table1.iter_mut()), + it_2: Some(self.table2.iter_mut()), + it_3: Some(self.table3.iter_mut()), + it_4: Some(self.table4.iter_mut()), + _phantom: PhantomData, + } + } +} + +pub struct UnsizedHashtableIter<'a, K, V> +where K: UnsizedKeyable + ?Sized +{ + it_0: Option>, + it_1: Option, V>>, + it_2: Option, V>>, + it_3: Option, V>>, + it_4: Option>, + _phantom: PhantomData<&'a mut K>, +} + +impl<'a, K, V> Iterator for UnsizedHashtableIter<'a, K, V> +where K: UnsizedKeyable + ?Sized +{ + type Item = UnsizedHashtableEntryRef<'a, K, V>; + + fn next(&mut self) -> Option { + if let Some(it) = self.it_0.as_mut() { + if let Some(e) = it.next() { + return Some(UnsizedHashtableEntryRef( + UnsizedHashtableEntryRefInner::Table0(e, PhantomData), + )); + } + self.it_0 = None; + } + if let Some(it) = self.it_1.as_mut() { + if let Some(e) = it.next() { + return Some(UnsizedHashtableEntryRef( + UnsizedHashtableEntryRefInner::Table1(e), + )); + } + self.it_1 = None; + } + if let Some(it) = self.it_2.as_mut() { + if let Some(e) = it.next() { + return Some(UnsizedHashtableEntryRef( + UnsizedHashtableEntryRefInner::Table2(e), + )); + } + self.it_2 = None; + } + if let Some(it) = self.it_3.as_mut() { + if let Some(e) = it.next() { + return Some(UnsizedHashtableEntryRef( + UnsizedHashtableEntryRefInner::Table3(e), + )); + } + self.it_3 = None; + } + if let Some(it) = self.it_4.as_mut() { + if let Some(e) = it.next() { + return Some(UnsizedHashtableEntryRef( + UnsizedHashtableEntryRefInner::Table4(e), + )); + } + self.it_4 = None; + } + None + } +} + +pub struct UnsizedHashtableIterMut<'a, K, V> +where K: UnsizedKeyable + ?Sized +{ + it_0: Option>, + it_1: Option, V>>, + it_2: Option, V>>, + it_3: Option, V>>, + it_4: Option>, + _phantom: PhantomData<&'a mut K>, +} + +impl<'a, K, V> Iterator for UnsizedHashtableIterMut<'a, K, V> +where K: UnsizedKeyable + ?Sized +{ + type Item = UnsizedHashtableEntryMutRef<'a, K, V>; + + fn next(&mut self) -> Option { + if let Some(it) = self.it_0.as_mut() { + if let Some(e) = it.next() { + return Some(UnsizedHashtableEntryMutRef( + UnsizedHashtableEntryMutRefInner::Table0(e, PhantomData), + )); + } + self.it_0 = None; + } + if let Some(it) = self.it_1.as_mut() { + if let Some(e) = it.next() { + return Some(UnsizedHashtableEntryMutRef( + UnsizedHashtableEntryMutRefInner::Table1(e), + )); + } + self.it_1 = None; + } + if let Some(it) = self.it_2.as_mut() { + if let Some(e) = it.next() { + return Some(UnsizedHashtableEntryMutRef( + UnsizedHashtableEntryMutRefInner::Table2(e), + )); + } + self.it_2 = None; + } + if let Some(it) = self.it_3.as_mut() { + if let Some(e) = it.next() { + return Some(UnsizedHashtableEntryMutRef( + UnsizedHashtableEntryMutRefInner::Table3(e), + )); + } + self.it_3 = None; + } + if let Some(it) = self.it_4.as_mut() { + if let Some(e) = it.next() { + return Some(UnsizedHashtableEntryMutRef( + UnsizedHashtableEntryMutRefInner::Table4(e), + )); + } + self.it_4 = None; + } + None + } +} + +enum UnsizedHashtableEntryRefInner<'a, K: ?Sized, V> { + Table0(&'a Entry<[u8; 2], V>, PhantomData), + Table1(&'a Entry, V>), + Table2(&'a Entry, V>), + Table3(&'a Entry, V>), + Table4(&'a Entry), +} + +impl<'a, K: ?Sized, V> Copy for UnsizedHashtableEntryRefInner<'a, K, V> {} + +impl<'a, K: ?Sized, V> Clone for UnsizedHashtableEntryRefInner<'a, K, V> { + fn clone(&self) -> Self { + use UnsizedHashtableEntryRefInner::*; + match self { + Table0(a, b) => Table0(a, *b), + Table1(a) => Table1(a), + Table2(a) => Table2(a), + Table3(a) => Table3(a), + Table4(a) => Table4(a), + } + } +} + +impl<'a, K: ?Sized + UnsizedKeyable, V> UnsizedHashtableEntryRefInner<'a, K, V> { + fn key(self) -> &'a K { + use UnsizedHashtableEntryRefInner::*; + match self { + Table0(e, _) => unsafe { + let key = e.key.assume_init_ref(); + if key[1] != 0 { + UnsizedKeyable::from_bytes(&key[..2]) + } else if key[0] != 0 { + UnsizedKeyable::from_bytes(&key[..1]) + } else { + UnsizedKeyable::from_bytes(&key[..0]) + } + }, + Table1(e) => unsafe { + let bytes = e.key().1.get().to_le_bytes(); + for i in (0..=7).rev() { + if bytes[i] != 0 { + return UnsizedKeyable::from_bytes(std::slice::from_raw_parts( + e.key() as *const _ as *const u8, + i + 1, + )); + } + } + unreachable!() + }, + Table2(e) => unsafe { + let bytes = e.key().1.get().to_le_bytes(); + for i in (0..=7).rev() { + if bytes[i] != 0 { + return UnsizedKeyable::from_bytes(std::slice::from_raw_parts( + e.key() as *const _ as *const u8, + i + 9, + )); + } + } + unreachable!() + }, + Table3(e) => unsafe { + let bytes = e.key().1.get().to_le_bytes(); + for i in (0..=7).rev() { + if bytes[i] != 0 { + return UnsizedKeyable::from_bytes(std::slice::from_raw_parts( + e.key() as *const _ as *const u8, + i + 17, + )); + } + } + unreachable!() + }, + Table4(e) => unsafe { + UnsizedKeyable::from_bytes(e.key.assume_init().key.unwrap().as_ref()) + }, + } + } + fn get(self) -> &'a V { + use UnsizedHashtableEntryRefInner::*; + match self { + Table0(e, _) => e.get(), + Table1(e) => e.get(), + Table2(e) => e.get(), + Table3(e) => e.get(), + Table4(e) => e.get(), + } + } + fn get_ptr(self) -> *const V { + use UnsizedHashtableEntryRefInner::*; + match self { + Table0(e, _) => e.val.as_ptr(), + Table1(e) => e.val.as_ptr(), + Table2(e) => e.val.as_ptr(), + Table3(e) => e.val.as_ptr(), + Table4(e) => e.val.as_ptr(), + } + } +} + +pub struct UnsizedHashtableEntryRef<'a, K: ?Sized, V>(UnsizedHashtableEntryRefInner<'a, K, V>); + +impl<'a, K: ?Sized, V> Copy for UnsizedHashtableEntryRef<'a, K, V> {} + +impl<'a, K: ?Sized, V> Clone for UnsizedHashtableEntryRef<'a, K, V> { + fn clone(&self) -> Self { + Self(self.0) + } +} + +impl<'a, K: ?Sized + UnsizedKeyable, V> UnsizedHashtableEntryRef<'a, K, V> { + pub fn key(self) -> &'a K { + self.0.key() + } + pub fn get(self) -> &'a V { + self.0.get() + } + pub fn get_ptr(self) -> *const V { + self.0.get_ptr() + } +} + +enum UnsizedHashtableEntryMutRefInner<'a, K: ?Sized, V> { + Table0(&'a mut Entry<[u8; 2], V>, PhantomData), + Table1(&'a mut Entry, V>), + Table2(&'a mut Entry, V>), + Table3(&'a mut Entry, V>), + Table4(&'a mut Entry), +} + +impl<'a, K: ?Sized + UnsizedKeyable, V> UnsizedHashtableEntryMutRefInner<'a, K, V> { + fn key(&self) -> &'a K { + use UnsizedHashtableEntryMutRefInner::*; + match self { + Table0(e, _) => unsafe { + let key = e.key.assume_init_ref(); + if key[1] != 0 { + &*(UnsizedKeyable::from_bytes(&key[..2]) as *const K) + } else if key[0] != 0 { + &*(UnsizedKeyable::from_bytes(&key[..1]) as *const K) + } else { + &*(UnsizedKeyable::from_bytes(&key[..0]) as *const K) + } + }, + Table1(e) => unsafe { + let bytes = e.key().1.get().to_le_bytes(); + for i in (0..=7).rev() { + if bytes[i] != 0 { + return UnsizedKeyable::from_bytes(std::slice::from_raw_parts( + e.key.assume_init_ref() as *const _ as *const u8, + i + 1, + )); + } + } + unreachable!() + }, + Table2(e) => unsafe { + let bytes = e.key().1.get().to_le_bytes(); + for i in (0..=7).rev() { + if bytes[i] != 0 { + return UnsizedKeyable::from_bytes(std::slice::from_raw_parts( + e.key.assume_init_ref() as *const _ as *const u8, + i + 9, + )); + } + } + unreachable!() + }, + Table3(e) => unsafe { + let bytes = e.key().1.get().to_le_bytes(); + for i in (0..=7).rev() { + if bytes[i] != 0 { + return UnsizedKeyable::from_bytes(std::slice::from_raw_parts( + e.key.assume_init_ref() as *const _ as *const u8, + i + 17, + )); + } + } + unreachable!() + }, + Table4(e) => unsafe { + UnsizedKeyable::from_bytes(e.key.assume_init().key.unwrap().as_ref()) + }, + } + } + fn get(&self) -> &V { + use UnsizedHashtableEntryMutRefInner::*; + match self { + Table0(e, _) => e.get(), + Table1(e) => e.get(), + Table2(e) => e.get(), + Table3(e) => e.get(), + Table4(e) => e.get(), + } + } + fn get_ptr(&self) -> *const V { + use UnsizedHashtableEntryMutRefInner::*; + match self { + Table0(e, _) => e.val.as_ptr(), + Table1(e) => e.val.as_ptr(), + Table2(e) => e.val.as_ptr(), + Table3(e) => e.val.as_ptr(), + Table4(e) => e.val.as_ptr(), + } + } + fn get_mut(&mut self) -> &mut V { + use UnsizedHashtableEntryMutRefInner::*; + match self { + Table0(e, _) => e.get_mut(), + Table1(e) => e.get_mut(), + Table2(e) => e.get_mut(), + Table3(e) => e.get_mut(), + Table4(e) => e.get_mut(), + } + } + fn write(&mut self, val: V) { + use UnsizedHashtableEntryMutRefInner::*; + match self { + Table0(e, _) => e.write(val), + Table1(e) => e.write(val), + Table2(e) => e.write(val), + Table3(e) => e.write(val), + Table4(e) => e.write(val), + } + } +} + +pub struct UnsizedHashtableEntryMutRef<'a, K: ?Sized, V>( + UnsizedHashtableEntryMutRefInner<'a, K, V>, +); + +impl<'a, K: ?Sized + UnsizedKeyable, V> UnsizedHashtableEntryMutRef<'a, K, V> { + pub fn key(&self) -> &'a K { + self.0.key() + } + pub fn get(&self) -> &V { + self.0.get() + } + pub fn get_ptr(&self) -> *const V { + self.0.get_ptr() + } + pub fn get_mut_ptr(&self) -> *mut V { + self.get_ptr() as *mut V + } + pub fn get_mut(&mut self) -> &mut V { + self.0.get_mut() + } + pub fn write(&mut self, val: V) { + self.0.write(val) + } +} + +#[repr(C)] +#[derive(Clone, Copy, PartialEq, Eq)] +pub(crate) struct InlineKey(pub [u64; N], pub NonZeroU64); + +unsafe impl Keyable for InlineKey { + #[inline(always)] + fn equals_zero(_: &Self) -> bool { + false + } + + #[inline(always)] + fn is_zero(this: &MaybeUninit) -> bool { + unsafe { *(this as *const _ as *const u64).add(N) == 0 } + } + + #[inline(always)] + fn hash(&self) -> u64 { + (self.0, self.1).fast_hash() + } +} + +#[derive(Copy, Clone)] +pub(crate) struct FallbackKey { + key: Option>, + hash: u64, +} + +impl FallbackKey { + unsafe fn new(key: &[u8]) -> Self { + Self { + key: Some(NonNull::from(key)), + hash: key.fast_hash(), + } + } +} + +impl PartialEq for FallbackKey { + fn eq(&self, other: &Self) -> bool { + if self.hash == other.hash { + unsafe { self.key.map(|x| x.as_ref()) == other.key.map(|x| x.as_ref()) } + } else { + false + } + } +} + +impl Eq for FallbackKey {} + +unsafe impl Keyable for FallbackKey { + #[inline(always)] + fn equals_zero(_: &Self) -> bool { + false + } + + #[inline(always)] + fn is_zero(this: &MaybeUninit) -> bool { + unsafe { this.assume_init_ref().key.is_none() } + } + + #[inline(always)] + fn hash(&self) -> u64 { + self.hash + } +} + +impl<'a, K: UnsizedKeyable + ?Sized + 'a, V: 'a> EntryRefLike + for UnsizedHashtableEntryRef<'a, K, V> +{ + type KeyRef = &'a K; + type ValueRef = &'a V; + + fn key(&self) -> Self::KeyRef { + (*self).key() + } + fn get(&self) -> Self::ValueRef { + (*self).get() + } +} + +impl<'a, K: UnsizedKeyable + ?Sized + 'a, V: 'a> EntryMutRefLike + for UnsizedHashtableEntryMutRef<'a, K, V> +{ + type KeyRef = &'a K; + type Value = V; + + fn key(&self) -> Self::KeyRef { + self.key() + } + fn get(&self) -> &Self::Value { + self.get() + } + fn get_mut(&mut self) -> &mut Self::Value { + self.get_mut() + } + fn write(&mut self, value: Self::Value) { + self.write(value); + } +} + +impl HashtableLike for UnsizedHashtable<[u8], V, A> +where A: Allocator + Clone + Default +{ + type Key = [u8]; + type KeyRef<'a> = &'a [u8] where Self::Key:'a; + type Value = V; + + type EntryRef<'a> = UnsizedHashtableEntryRef<'a, [u8], V> where Self:'a, V: 'a; + type EntryMutRef<'a> = UnsizedHashtableEntryMutRef<'a, [u8], V> where Self:'a, V: 'a; + + type Iterator<'a> = UnsizedHashtableIter<'a, [u8], V> where Self:'a, V: 'a; + type IteratorMut<'a> = UnsizedHashtableIterMut<'a, [u8], V> where Self:'a, V: 'a; + + fn entry<'a>(&self, key_ref: Self::KeyRef<'a>) -> Option> + where Self::Key: 'a { + self.entry(key_ref) + } + fn entry_mut<'a>(&mut self, key_ref: Self::KeyRef<'a>) -> Option> + where Self::Key: 'a { + self.entry_mut(key_ref) + } + + fn get<'a>(&self, key_ref: Self::KeyRef<'a>) -> Option<&Self::Value> + where Self::Key: 'a { + self.get(key_ref) + } + fn get_mut<'a>(&mut self, key_ref: Self::KeyRef<'a>) -> Option<&mut Self::Value> + where Self::Key: 'a { + self.get_mut(key_ref) + } + + unsafe fn insert<'a>( + &mut self, + key_ref: Self::KeyRef<'a>, + ) -> Result<&mut MaybeUninit, &mut Self::Value> + where + Self::Key: 'a, + { + self.insert(key_ref) + } + unsafe fn insert_and_entry<'a>( + &mut self, + key_ref: Self::KeyRef<'a>, + ) -> Result, Self::EntryMutRef<'_>> + where + Self::Key: 'a, + { + self.insert_and_entry(key_ref) + } + + fn iter(&self) -> Self::Iterator<'_> { + self.iter() + } + fn iter_mut(&mut self) -> Self::IteratorMut<'_> { + self.iter_mut() + } +} diff --git a/src/common/hashtable/src/utils.rs b/src/common/hashtable/src/utils.rs new file mode 100644 index 0000000000000..2b24fbb7f26b3 --- /dev/null +++ b/src/common/hashtable/src/utils.rs @@ -0,0 +1,98 @@ +// Copyright 2021 Datafuse Labs. +// +// Licensed 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 std::intrinsics::assume; +use std::mem::MaybeUninit; +use std::ops::Deref; +use std::ops::DerefMut; + +use super::table0::Entry; +use super::traits::Keyable; + +#[derive(Clone, Copy, PartialEq, Eq)] +pub struct Hashed { + hash: u64, + key: K, +} + +impl Hashed { + #[inline(always)] + pub fn new(key: K) -> Self { + Self { + hash: K::hash(&key), + key, + } + } + #[inline(always)] + pub fn key(self) -> K { + self.key + } +} + +unsafe impl Keyable for Hashed { + #[inline(always)] + fn is_zero(this: &MaybeUninit) -> bool { + unsafe { + let addr = std::ptr::addr_of!((*this.as_ptr()).key); + K::is_zero(&*(addr as *const MaybeUninit)) + } + } + + #[inline(always)] + fn equals_zero(this: &Self) -> bool { + K::equals_zero(&this.key) + } + + #[inline(always)] + fn hash(&self) -> u64 { + self.hash + } +} + +pub struct ZeroEntry(pub Option>); + +impl Deref for ZeroEntry { + type Target = Option>; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl DerefMut for ZeroEntry { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.0 + } +} + +impl Drop for ZeroEntry { + fn drop(&mut self) { + if let Some(e) = self.0.as_mut() { + unsafe { + e.val.assume_init_drop(); + } + } + } +} + +#[inline(always)] +pub unsafe fn read_le(data: *const u8, len: usize) -> u64 { + assume(0 < len && len <= 8); + let s = 64 - 8 * len as isize; + if data as usize & 2048 == 0 { + (data as *const u64).read_unaligned() & (u64::MAX >> s) + } else { + (data.offset(len as isize - 8) as *const u64).read_unaligned() >> s + } +} diff --git a/src/common/hashtable/tests/it/main.rs b/src/common/hashtable/tests/it/main.rs index 24259b361b541..27dc316ecc6d2 100644 --- a/src/common/hashtable/tests/it/main.rs +++ b/src/common/hashtable/tests/it/main.rs @@ -14,115 +14,127 @@ use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; -use std::sync::Arc; use common_hashtable::HashMap; -use common_hashtable::HashMapKind; -use common_hashtable::HashTableGrower; -use common_hashtable::SingleLevelGrower; +use common_hashtable::StackHashMap; +use common_hashtable::TwolevelHashMap; +use common_hashtable::UnsizedHashMap; use rand::Rng; -#[test] -fn test_hash_table_grower() { - let mut grower = SingleLevelGrower::default(); - - assert_eq!(grower.max_size(), 256); - - assert!(grower.overflow(129)); - assert!(!grower.overflow(128)); - - assert_eq!(grower.place(1), 1); - assert_eq!(grower.place(255), 255); - assert_eq!(grower.place(256), 0); - assert_eq!(grower.place(257), 1); - - assert_eq!(grower.next_place(1), 2); - assert_eq!(grower.next_place(2), 3); - assert_eq!(grower.next_place(254), 255); - assert_eq!(grower.next_place(255), 0); - - grower.increase_size(); - assert_eq!(grower.max_size(), 1024); +macro_rules! simple_test { + ($t: tt) => { + static COUNT: AtomicUsize = AtomicUsize::new(0); + #[derive(Debug)] + struct U64(u64); + impl U64 { + fn new(x: u64) -> Self { + COUNT.fetch_add(1, Ordering::Relaxed); + Self(x) + } + } + impl Drop for U64 { + fn drop(&mut self) { + COUNT.fetch_sub(1, Ordering::Relaxed); + } + } + let mut sequence = vec![0u64; 1 << 12]; + sequence.fill_with(|| rand::thread_rng().gen_range(0..1 << 10)); + let mut standard = std::collections::HashMap::::new(); + let mut hashtable = $t::::new(); + for &s in sequence.iter() { + match standard.get_mut(&s) { + Some(x) => { + *x += 1; + } + None => { + standard.insert(s, 1); + } + } + } + for &s in sequence.iter() { + match unsafe { hashtable.insert(s) } { + Ok(x) => { + x.write(U64::new(1)); + } + Err(x) => { + x.0 += 1; + } + } + } + assert_eq!(standard.len(), hashtable.len()); + let mut check = std::collections::HashSet::new(); + for e in hashtable.iter() { + assert!(check.insert(e.key())); + assert_eq!(standard[e.key()], e.get().0); + } + drop(hashtable); + assert_eq!(COUNT.load(Ordering::Relaxed), 0); + }; } #[test] -fn test_two_level_hash_table() { - let mut hashtable = HashMapKind::::create_hash_table(); - let mut inserted = true; - let entity = hashtable.insert_key(&1u64, &mut inserted); - if inserted { - entity.set_value(2); - } - - unsafe { - hashtable.convert_to_two_level(); - } - - let is_two_level = match hashtable { - HashMapKind::HashTable(_) => false, - HashMapKind::TwoLevelHashTable(_) => true, - }; - assert!(is_two_level); - - let entity = hashtable.insert_key(&1u64, &mut inserted); - assert!(!inserted); +fn test_hash_map() { + simple_test!(HashMap); +} - assert_eq!(entity.get_value(), &2); +#[test] +fn test_stack_hash_map() { + simple_test!(StackHashMap); } #[test] -fn test_hash_map_drop() { - #[derive(Debug, Clone)] - struct A { - item: Arc, - } +fn test_twolevel_hash_map() { + simple_test!(TwolevelHashMap); +} - impl A { - pub fn new(item: Arc) -> Self { - item.fetch_add(1, Ordering::Relaxed); - Self { item } +#[test] +fn test_unsized_hash_map() { + static COUNT: AtomicUsize = AtomicUsize::new(0); + #[derive(Debug)] + struct U64(u64); + impl U64 { + fn new(x: u64) -> Self { + COUNT.fetch_add(1, Ordering::Relaxed); + Self(x) } } - - impl Drop for A { + impl Drop for U64 { fn drop(&mut self) { - self.item.fetch_sub(1, Ordering::Relaxed); + COUNT.fetch_sub(1, Ordering::Relaxed); } } - - let item = Arc::new(AtomicUsize::new(0)); - let two_level_item = Arc::new(AtomicUsize::new(0)); - - { - let mut table = HashMap::>::create(); - let mut two_level_table = HashMapKind::>::create_hash_table(); - let mut rng = rand::thread_rng(); - - for _ in 0..1200 { - let key = rng.gen::() % 300; - - let mut inserted = false; - let entity = table.insert_key(&key, &mut inserted); - if inserted { - entity.set_value(vec![A::new(item.clone())]); - } else { - entity.get_mut_value().push(A::new(item.clone())); + let mut sequence = Vec::new(); + for _ in 0..1000 { + let length = rand::thread_rng().gen_range(0..64); + let mut array = vec![0u8; length]; + rand::thread_rng().fill(&mut array[..]); + sequence.push(array); + } + let mut standard = std::collections::HashMap::<&[u8], u64>::new(); + for s in sequence.iter() { + if let Some(x) = standard.get_mut(&s[..]) { + *x += 1; + } else { + standard.insert(s, 1); + } + } + let mut hashtable = UnsizedHashMap::<[u8], U64>::new(); + for s in sequence.iter() { + match unsafe { hashtable.insert_and_entry(s) } { + Ok(mut e) => { + e.write(U64::new(1u64)); } - - inserted = false; - let entity = two_level_table.insert_key(&key, &mut inserted); - if inserted { - entity.set_value(vec![A::new(item.clone())]); - } else { - entity.get_mut_value().push(A::new(item.clone())); + Err(mut e) => { + e.get_mut().0 += 1; } } - unsafe { - two_level_table.convert_to_two_level(); - } } - let count = item.load(Ordering::Relaxed); - let count2 = two_level_item.load(Ordering::Relaxed); - assert_eq!(count, 0); - assert_eq!(count2, 0); + assert_eq!(standard.len(), hashtable.len()); + let mut check = std::collections::HashSet::new(); + for e in hashtable.iter() { + assert!(check.insert(e.key())); + assert_eq!(standard.get(e.key()).copied().unwrap(), e.get().0); + } + drop(hashtable); + assert_eq!(COUNT.load(Ordering::Relaxed), 0); } diff --git a/src/query/functions-v2/src/aggregates/aggregate_distinct_state.rs b/src/query/functions-v2/src/aggregates/aggregate_distinct_state.rs index a4bcaac95effc..d16e22da5ebfd 100644 --- a/src/query/functions-v2/src/aggregates/aggregate_distinct_state.rs +++ b/src/query/functions-v2/src/aggregates/aggregate_distinct_state.rs @@ -32,10 +32,9 @@ use common_expression::Column; use common_expression::ColumnBuilder; use common_expression::Scalar; use common_hashtable::HashSet as CommonHashSet; -use common_hashtable::HashSetWithStackMemory; -use common_hashtable::HashTableEntity; -use common_hashtable::HashTableKeyable; +use common_hashtable::HashtableKeyable; use common_hashtable::KeysRef; +use common_hashtable::StackHashSet; use common_io::prelude::*; use serde::de::DeserializeOwned; use serde::Serialize; @@ -63,8 +62,8 @@ pub struct AggregateDistinctState { set: HashSet, RandomState>, } -pub struct AggregateDistinctNumberState { - set: HashSetWithStackMemory<{ 16 * 8 }, T>, +pub struct AggregateDistinctNumberState { + set: StackHashSet, inserted: bool, } @@ -159,27 +158,36 @@ impl DistinctStateFunc for AggregateDistinctState { impl AggregateDistinctStringState { #[inline] fn insert_and_materialize(&mut self, key: &KeysRef) { - let entity = self.set.insert_key(key, &mut self.inserted); - if self.inserted { - let data = unsafe { key.as_slice() }; - - let holder = self.holders.last_mut().unwrap(); - // TODO(sundy): may cause memory fragmentation, refactor this using arena - if holder.may_resize(data.len()) { - let mut holder = StringColumnBuilder::with_capacity( - HOLDER_CAPACITY, - HOLDER_BYTES_CAPACITY.max(data.len()), - ); - holder.put_slice(data); - holder.commit_row(); - let value = unsafe { holder.index_unchecked(holder.len() - 1) }; - entity.set_key(KeysRef::create(value.as_ptr() as usize, value.len())); - self.holders.push(holder); - } else { - holder.put_slice(data); - holder.commit_row(); - let value = unsafe { holder.index_unchecked(holder.len() - 1) }; - entity.set_key(KeysRef::create(value.as_ptr() as usize, value.len())); + match unsafe { self.set.insert_and_entry(*key) } { + Ok(entity) => { + self.inserted = true; + let data = unsafe { key.as_slice() }; + + let holder = self.holders.last_mut().unwrap(); + // TODO(sundy): may cause memory fragmentation, refactor this using arena + if holder.may_resize(data.len()) { + let mut holder = StringColumnBuilder::with_capacity( + HOLDER_CAPACITY, + HOLDER_BYTES_CAPACITY.max(data.len()), + ); + holder.put_slice(data); + holder.commit_row(); + let value = unsafe { holder.index_unchecked(holder.len() - 1) }; + unsafe { + entity.set_key(KeysRef::create(value.as_ptr() as usize, value.len())); + } + self.holders.push(holder); + } else { + holder.put_slice(data); + holder.commit_row(); + let value = unsafe { holder.index_unchecked(holder.len() - 1) }; + unsafe { + entity.set_key(KeysRef::create(value.as_ptr() as usize, value.len())); + } + } + } + Err(_) => { + self.inserted = false; } } } @@ -188,7 +196,7 @@ impl AggregateDistinctStringState { impl DistinctStateFunc for AggregateDistinctStringState { fn new() -> Self { AggregateDistinctStringState { - set: CommonHashSet::create(), + set: CommonHashSet::new(), inserted: false, holders: vec![StringColumnBuilder::with_capacity( HOLDER_CAPACITY, @@ -209,7 +217,7 @@ impl DistinctStateFunc for AggregateDistinctStringState { for index in 0..holder.len() { let data = unsafe { holder.index_unchecked(index) }; let key = KeysRef::create(data.as_ptr() as usize, data.len()); - self.set.insert_key(&key, &mut self.inserted); + self.inserted = self.set.set_insert(key).is_ok(); } } Ok(()) @@ -262,7 +270,7 @@ impl DistinctStateFunc for AggregateDistinctStringState { fn merge(&mut self, rhs: &Self) -> Result<()> { for value in rhs.set.iter() { - self.insert_and_materialize(value.get_key()); + self.insert_and_materialize(value.key()); } Ok(()) } @@ -297,29 +305,29 @@ impl DistinctStateFunc for AggregateDistinctStringState { } impl DistinctStateFunc for AggregateDistinctNumberState -where T: Number + Serialize + DeserializeOwned + HashTableKeyable +where T: Number + Serialize + DeserializeOwned + HashtableKeyable { fn new() -> Self { AggregateDistinctNumberState { - set: HashSetWithStackMemory::create(), + set: StackHashSet::new(), inserted: false, } } fn serialize(&self, writer: &mut BytesMut) -> Result<()> { writer.write_uvarint(self.set.len() as u64)?; - for value in self.set.iter() { - serialize_into_buf(writer, value.get_key())? + for e in self.set.iter() { + serialize_into_buf(writer, e.key())? } Ok(()) } fn deserialize(&mut self, reader: &mut &[u8]) -> Result<()> { let size = reader.read_uvarint()?; - self.set = HashSetWithStackMemory::with_capacity(size as usize); + self.set = StackHashSet::with_capacity(size as usize); for _ in 0..size { let t: T = deserialize_from_slice(reader)?; - let _ = self.set.insert_key(&t, &mut self.inserted); + self.inserted = self.set.set_insert(t).is_ok(); } Ok(()) } @@ -335,7 +343,7 @@ where T: Number + Serialize + DeserializeOwned + HashTableKeyable fn add(&mut self, columns: &[Column], row: usize) -> Result<()> { let col = NumberType::::try_downcast_column(&columns[0]).unwrap(); let v = unsafe { col.get_unchecked(row) }; - self.set.insert_key(v, &mut self.inserted); + self.inserted = self.set.set_insert(*v).is_ok(); Ok(()) } @@ -350,14 +358,14 @@ where T: Number + Serialize + DeserializeOwned + HashTableKeyable Some(bitmap) => { for (t, v) in col.iter().zip(bitmap.iter()) { if v { - self.set.insert_key(t, &mut self.inserted); + self.inserted = self.set.set_insert(*t).is_ok(); } } } None => { for row in 0..input_rows { let v = unsafe { col.get_unchecked(row) }; - self.set.insert_key(v, &mut self.inserted); + self.inserted = self.set.set_insert(*v).is_ok(); } } } @@ -365,26 +373,26 @@ where T: Number + Serialize + DeserializeOwned + HashTableKeyable } fn merge(&mut self, rhs: &Self) -> Result<()> { - self.set.merge(&rhs.set); + self.set.set_merge(&rhs.set); Ok(()) } fn build_columns(&mut self, _types: &[DataType]) -> Result> { - let values: Buffer = self.set.iter().map(|e| *e.get_key()).collect(); + let values: Buffer = self.set.iter().map(|e| *e.key()).collect(); Ok(vec![NumberType::::upcast_column(values)]) } } // For count(distinct string) and uniq(string) pub struct AggregateUniqStringState { - set: HashSetWithStackMemory<{ 16 * 8 }, u128>, + set: StackHashSet, inserted: bool, } impl DistinctStateFunc for AggregateUniqStringState { fn new() -> Self { AggregateUniqStringState { - set: HashSetWithStackMemory::create(), + set: StackHashSet::new(), inserted: false, } } @@ -392,17 +400,17 @@ impl DistinctStateFunc for AggregateUniqStringState { fn serialize(&self, writer: &mut BytesMut) -> Result<()> { writer.write_uvarint(self.set.len() as u64)?; for value in self.set.iter() { - serialize_into_buf(writer, value.get_key())? + serialize_into_buf(writer, value.key())? } Ok(()) } fn deserialize(&mut self, reader: &mut &[u8]) -> Result<()> { let size = reader.read_uvarint()?; - self.set = HashSetWithStackMemory::with_capacity(size as usize); + self.set = StackHashSet::with_capacity(size as usize); for _ in 0..size { let e = deserialize_from_slice(reader)?; - let _ = self.set.insert_key(&e, &mut self.inserted); + self.inserted = self.set.set_insert(e).is_ok(); } Ok(()) } @@ -421,7 +429,7 @@ impl DistinctStateFunc for AggregateUniqStringState { let mut hasher = SipHasher24::new(); hasher.write(data); let hash128 = hasher.finish128(); - self.set.insert_key(&hash128.into(), &mut self.inserted); + self.inserted = self.set.set_insert(hash128.into()).is_ok(); Ok(()) } @@ -439,7 +447,7 @@ impl DistinctStateFunc for AggregateUniqStringState { let mut hasher = SipHasher24::new(); hasher.write(t); let hash128 = hasher.finish128(); - self.set.insert_key(&hash128.into(), &mut self.inserted); + self.inserted = self.set.set_insert(hash128.into()).is_ok(); } } } @@ -449,7 +457,7 @@ impl DistinctStateFunc for AggregateUniqStringState { let mut hasher = SipHasher24::new(); hasher.write(data); let hash128 = hasher.finish128(); - self.set.insert_key(&hash128.into(), &mut self.inserted); + let _ = self.set.set_insert(hash128.into()).is_ok(); } } } @@ -457,7 +465,7 @@ impl DistinctStateFunc for AggregateUniqStringState { } fn merge(&mut self, rhs: &Self) -> Result<()> { - self.set.merge(&rhs.set); + self.set.set_merge(&rhs.set); Ok(()) } diff --git a/src/query/functions/src/aggregates/aggregate_distinct_state.rs b/src/query/functions/src/aggregates/aggregate_distinct_state.rs index 8da43f807845e..112a2a7eb9e98 100644 --- a/src/query/functions/src/aggregates/aggregate_distinct_state.rs +++ b/src/query/functions/src/aggregates/aggregate_distinct_state.rs @@ -25,10 +25,9 @@ use common_arrow::arrow::bitmap::Bitmap; use common_datavalues::prelude::*; use common_exception::Result; use common_hashtable::HashSet as CommonHashSet; -use common_hashtable::HashSetWithStackMemory; -use common_hashtable::HashTableEntity; -use common_hashtable::HashTableKeyable; +use common_hashtable::HashtableKeyable; use common_hashtable::KeysRef; +use common_hashtable::StackHashSet; use common_io::prelude::*; use serde::Deserialize; use serde::Serialize; @@ -61,10 +60,9 @@ pub struct AggregateDistinctState { set: HashSet, } -pub struct AggregateDistinctPrimitiveState + HashTableKeyable> { - set: HashSetWithStackMemory<{ 16 * 8 }, E>, +pub struct AggregateDistinctPrimitiveState + HashtableKeyable> { + set: StackHashSet, _t: PhantomData, - inserted: bool, } const HOLDER_CAPACITY: usize = 256; @@ -185,25 +183,34 @@ impl DistinctStateFunc for AggregateDistinctState { impl AggregateDistinctStringState { #[inline] fn insert_and_materialize(&mut self, key: &KeysRef) { - let entity = self.set.insert_key(key, &mut self.inserted); - if self.inserted { - let data = unsafe { key.as_slice() }; - - let holder = self.holders.last_mut().unwrap(); - // TODO(sundy): may cause memory fragmentation, refactor this using arena - if holder.may_resize(data.len()) { - let mut holder = MutableStringColumn::with_values_capacity( - HOLDER_BYTES_CAPACITY.max(data.len()), - HOLDER_CAPACITY, - ); - holder.push(data); - let value = unsafe { holder.value_unchecked(holder.len() - 1) }; - entity.set_key(KeysRef::create(value.as_ptr() as usize, value.len())); - self.holders.push(holder); - } else { - holder.push(data); - let value = unsafe { holder.value_unchecked(holder.len() - 1) }; - entity.set_key(KeysRef::create(value.as_ptr() as usize, value.len())); + match unsafe { self.set.insert_and_entry(*key) } { + Ok(entity) => { + self.inserted = true; + let data = unsafe { key.as_slice() }; + + let holder = self.holders.last_mut().unwrap(); + // TODO(sundy): may cause memory fragmentation, refactor this using arena + if holder.may_resize(data.len()) { + let mut holder = MutableStringColumn::with_values_capacity( + HOLDER_BYTES_CAPACITY.max(data.len()), + HOLDER_CAPACITY, + ); + holder.push(data); + let value = unsafe { holder.value_unchecked(holder.len() - 1) }; + unsafe { + entity.set_key(KeysRef::create(value.as_ptr() as usize, value.len())); + } + self.holders.push(holder); + } else { + holder.push(data); + let value = unsafe { holder.value_unchecked(holder.len() - 1) }; + unsafe { + entity.set_key(KeysRef::create(value.as_ptr() as usize, value.len())); + } + } + } + Err(_) => { + self.inserted = false; } } } @@ -212,7 +219,7 @@ impl AggregateDistinctStringState { impl DistinctStateFunc for AggregateDistinctStringState { fn new() -> Self { AggregateDistinctStringState { - set: CommonHashSet::create(), + set: CommonHashSet::new(), inserted: false, holders: vec![MutableStringColumn::with_values_capacity( HOLDER_BYTES_CAPACITY, @@ -233,7 +240,7 @@ impl DistinctStateFunc for AggregateDistinctStringState { for index in 0..holder.len() { let data = unsafe { holder.value_unchecked(index) }; let key = KeysRef::create(data.as_ptr() as usize, data.len()); - self.set.insert_key(&key, &mut self.inserted); + self.inserted = self.set.set_insert(key).is_ok(); } } Ok(()) @@ -286,7 +293,7 @@ impl DistinctStateFunc for AggregateDistinctStringState { fn merge(&mut self, rhs: &Self) -> Result<()> { for value in rhs.set.iter() { - self.insert_and_materialize(value.get_key()); + self.insert_and_materialize(value.key()); } Ok(()) } @@ -322,20 +329,19 @@ impl DistinctStateFunc for AggregateDistinctStringState { impl DistinctStateFunc for AggregateDistinctPrimitiveState where T: PrimitiveType + From, - E: From + Sync + Send + Clone + std::fmt::Debug + HashTableKeyable, + E: From + Sync + Send + Clone + std::fmt::Debug + HashtableKeyable, { fn new() -> Self { AggregateDistinctPrimitiveState { - set: HashSetWithStackMemory::create(), + set: StackHashSet::new(), _t: PhantomData, - inserted: false, } } fn serialize(&self, writer: &mut BytesMut) -> Result<()> { writer.write_uvarint(self.set.len() as u64)?; for value in self.set.iter() { - let t: T = value.get_key().clone().into(); + let t: T = (*value.key()).into(); serialize_into_buf(writer, &t)? } Ok(()) @@ -343,11 +349,11 @@ where fn deserialize(&mut self, reader: &mut &[u8]) -> Result<()> { let size = reader.read_uvarint()?; - self.set = HashSetWithStackMemory::with_capacity(size as usize); + self.set = StackHashSet::with_capacity(size as usize); for _ in 0..size { let t: T = deserialize_from_slice(reader)?; let e = E::from(t); - let _ = self.set.insert_key(&e, &mut self.inserted); + let _ = self.set.set_insert(e); } Ok(()) } @@ -363,7 +369,7 @@ where fn add(&mut self, columns: &[ColumnRef], row: usize) -> Result<()> { let array: &PrimitiveColumn = unsafe { Series::static_cast(&columns[0]) }; let v = unsafe { array.value_unchecked(row) }; - self.set.insert_key(&E::from(v), &mut self.inserted); + let _ = self.set.set_insert(E::from(v)); Ok(()) } @@ -378,14 +384,14 @@ where Some(bitmap) => { for (t, v) in array.iter().zip(bitmap.iter()) { if v { - self.set.insert_key(&E::from(*t), &mut self.inserted); + let _ = self.set.set_insert(E::from(*t)); } } } None => { for row in 0..input_rows { let v = unsafe { array.value_unchecked(row) }; - self.set.insert_key(&E::from(v), &mut self.inserted); + let _ = self.set.set_insert(E::from(v)); } } } @@ -393,16 +399,14 @@ where } fn merge(&mut self, rhs: &Self) -> Result<()> { - self.set.merge(&rhs.set); + for x in rhs.set.iter() { + let _ = self.set.set_insert(*x.key()); + } Ok(()) } fn build_columns(&mut self, _fields: &[DataField]) -> Result> { - let values: Vec = self - .set - .iter() - .map(|e| e.get_key().clone().into()) - .collect(); + let values: Vec = self.set.iter().map(|e| (*e.key()).into()).collect(); let result = PrimitiveColumn::::new_from_vec(values); Ok(vec![result.arc()]) } @@ -410,14 +414,14 @@ where // For count(distinct string) and uniq(string) pub struct AggregateUniqStringState { - set: HashSetWithStackMemory<{ 16 * 8 }, u128>, + set: StackHashSet, inserted: bool, } impl DistinctStateFunc for AggregateUniqStringState { fn new() -> Self { AggregateUniqStringState { - set: HashSetWithStackMemory::create(), + set: StackHashSet::new(), inserted: false, } } @@ -425,17 +429,17 @@ impl DistinctStateFunc for AggregateUniqStringState { fn serialize(&self, writer: &mut BytesMut) -> Result<()> { writer.write_uvarint(self.set.len() as u64)?; for value in self.set.iter() { - serialize_into_buf(writer, value.get_key())? + serialize_into_buf(writer, value.key())? } Ok(()) } fn deserialize(&mut self, reader: &mut &[u8]) -> Result<()> { let size = reader.read_uvarint()?; - self.set = HashSetWithStackMemory::with_capacity(size as usize); + self.set = StackHashSet::with_capacity(size as usize); for _ in 0..size { let e = deserialize_from_slice(reader)?; - let _ = self.set.insert_key(&e, &mut self.inserted); + self.inserted = self.set.set_insert(e).is_ok(); } Ok(()) } @@ -454,7 +458,7 @@ impl DistinctStateFunc for AggregateUniqStringState { let mut hasher = SipHasher24::new(); hasher.write(data); let hash128 = hasher.finish128(); - self.set.insert_key(&hash128.into(), &mut self.inserted); + self.inserted = self.set.set_insert(hash128.into()).is_ok(); Ok(()) } @@ -472,7 +476,7 @@ impl DistinctStateFunc for AggregateUniqStringState { let mut hasher = SipHasher24::new(); hasher.write(t); let hash128 = hasher.finish128(); - self.set.insert_key(&hash128.into(), &mut self.inserted); + self.inserted = self.set.set_insert(hash128.into()).is_ok(); } } } @@ -482,7 +486,7 @@ impl DistinctStateFunc for AggregateUniqStringState { let mut hasher = SipHasher24::new(); hasher.write(data); let hash128 = hasher.finish128(); - self.set.insert_key(&hash128.into(), &mut self.inserted); + self.inserted = self.set.set_insert(hash128.into()).is_ok(); } } } @@ -490,7 +494,7 @@ impl DistinctStateFunc for AggregateUniqStringState { } fn merge(&mut self, rhs: &Self) -> Result<()> { - self.set.merge(&rhs.set); + self.set.set_merge(&rhs.set); Ok(()) } diff --git a/src/query/functions/src/scalars/conditionals/in_evaluator.rs b/src/query/functions/src/scalars/conditionals/in_evaluator.rs index 279fbee440222..9bccec94a9bd2 100644 --- a/src/query/functions/src/scalars/conditionals/in_evaluator.rs +++ b/src/query/functions/src/scalars/conditionals/in_evaluator.rs @@ -21,9 +21,9 @@ use common_datavalues::type_coercion::numerical_coercion; use common_datavalues::with_match_physical_primitive_type; use common_exception::ErrorCode; use common_exception::Result; -use common_hashtable::HashSetWithStackMemory; -use common_hashtable::HashTableKeyable; +use common_hashtable::HashtableKeyable; use common_hashtable::KeysRef; +use common_hashtable::StackHashSet; use crate::scalars::default_column_cast; use crate::scalars::AlwaysNullFunction; @@ -31,10 +31,10 @@ use crate::scalars::Function; use crate::scalars::FunctionContext; #[derive(Clone)] -pub struct InEvaluatorImpl { +pub struct InEvaluatorImpl { input_type: DataTypeImpl, nonull_least_super_dt: DataTypeImpl, - set: Arc>, + set: Arc>, _col: ColumnRef, _s: PhantomData, } @@ -114,7 +114,7 @@ pub fn create_by_values( impl Function for InEvaluatorImpl where S: Scalar + Sync + Send + Clone, - Key: HashTableKeyable + Clone + 'static + Sync + Send, + Key: HashtableKeyable + Clone + 'static + Sync + Send, { fn name(&self) -> &str { if NEGATED { "NOT IN" } else { "IN" } @@ -164,7 +164,7 @@ where impl InEvaluatorImpl where S: Scalar + Clone + Sync + Send, - Key: HashTableKeyable + Clone + Send + Sync + 'static, + Key: HashtableKeyable + Clone + Send + Sync + 'static, { fn try_create( input_type: DataTypeImpl, @@ -175,10 +175,9 @@ where let col = col.convert_full_column(); let c = Series::check_get_scalar::(&col)?; - let mut set = HashSetWithStackMemory::with_capacity(c.len()); - let mut inserted = false; + let mut set = StackHashSet::with_capacity(c.len()); for data in c.scalar_iter() { - set.insert_key(&data.to_key(), &mut inserted); + let _ = set.set_insert(data.to_key()); } Ok(Box::new(Self { input_type, @@ -213,7 +212,7 @@ where impl fmt::Display for InEvaluatorImpl where S: Clone + Scalar + Sync + Send, - Key: HashTableKeyable + Clone + Send + Sync + 'static, + Key: HashtableKeyable + Clone + Send + Sync + 'static, { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { if NEGATED { diff --git a/src/query/service/src/pipelines/processors/mod.rs b/src/query/service/src/pipelines/processors/mod.rs index 747f9f3383646..40803ae3beb11 100644 --- a/src/query/service/src/pipelines/processors/mod.rs +++ b/src/query/service/src/pipelines/processors/mod.rs @@ -36,17 +36,11 @@ pub use transforms::AggregatorParams; pub use transforms::AggregatorTransformParams; pub use transforms::BlockCompactor; pub use transforms::ExpressionTransform; +pub use transforms::FixedKeyHashTable; pub use transforms::HashJoinDesc; pub use transforms::HashJoinState; pub use transforms::HashTable; pub use transforms::JoinHashTable; -pub use transforms::KeyU128HashTable; -pub use transforms::KeyU16HashTable; -pub use transforms::KeyU256HashTable; -pub use transforms::KeyU32HashTable; -pub use transforms::KeyU512HashTable; -pub use transforms::KeyU64HashTable; -pub use transforms::KeyU8HashTable; pub use transforms::MarkJoinCompactor; pub use transforms::RightJoinCompactor; pub use transforms::SerializerHashTable; diff --git a/src/query/service/src/pipelines/processors/transforms/aggregator/aggregator_final.rs b/src/query/service/src/pipelines/processors/transforms/aggregator/aggregator_final.rs index a1fe205804e5a..79b293ce79021 100644 --- a/src/query/service/src/pipelines/processors/transforms/aggregator/aggregator_final.rs +++ b/src/query/service/src/pipelines/processors/transforms/aggregator/aggregator_final.rs @@ -38,7 +38,8 @@ use crate::pipelines::processors::transforms::group_by::AggregatorState; use crate::pipelines::processors::transforms::group_by::GroupColumnsBuilder; use crate::pipelines::processors::transforms::group_by::KeysColumnIter; use crate::pipelines::processors::transforms::group_by::PolymorphicKeysHelper; -use crate::pipelines::processors::transforms::group_by::StateEntity; +use crate::pipelines::processors::transforms::group_by::StateEntityMutRef; +use crate::pipelines::processors::transforms::group_by::StateEntityRef; use crate::pipelines::processors::transforms::transform_aggregator::Aggregator; use crate::pipelines::processors::AggregatorParams; use crate::sessions::QueryContext; @@ -106,23 +107,24 @@ impl + Send> FinalAggregator< fn lookup_state( params: &AggregatorParams, state: &mut Method::State, - keys: &[>::Key], + keys: &[>::KeyRef<'_>], ) -> StateAddrs { let mut places = Vec::with_capacity(keys.len()); let mut inserted = true; for key in keys { - let entity = state.entity_by_key(key, &mut inserted); + let unsafe_state = state as *mut Method::State; + let mut entity = unsafe { (*unsafe_state).entity_by_key(*key, &mut inserted) }; match inserted { true => { - if let Some(place) = state.alloc_layout(params) { + if let Some(place) = unsafe { (*unsafe_state).alloc_layout(params) } { places.push(place); entity.set_state_value(place.addr()); } } false => { - let place: StateAddr = (*entity.get_state_value()).into(); + let place: StateAddr = entity.get_state_value().into(); places.push(place); } } @@ -145,7 +147,7 @@ impl + Send> Aggregator let group_by_two_level_threshold = self.ctx.get_settings().get_group_by_two_level_threshold()? as usize; if !self.state.is_two_level() && self.state.len() >= group_by_two_level_threshold { - self.state.convert_to_two_level(); + self.state.convert_to_twolevel(); } // first state places of current block @@ -203,7 +205,7 @@ impl + Send> Aggregator }; for group_entity in self.state.iter() { - let place: StateAddr = (*group_entity.get_state_value()).into(); + let place: StateAddr = group_entity.get_state_value().into(); for (idx, aggregate_function) in aggregate_functions.iter().enumerate() { let arg_place = place.next(offsets_aggregate_states[idx]); @@ -245,12 +247,12 @@ impl + Send> Aggregator let group_by_two_level_threshold = self.ctx.get_settings().get_group_by_two_level_threshold()? as usize; if !self.state.is_two_level() && self.state.len() >= group_by_two_level_threshold { - self.state.convert_to_two_level(); + self.state.convert_to_twolevel(); } let mut inserted = true; for keys_ref in keys_iter.get_slice() { - self.state.entity_by_key(keys_ref, &mut inserted); + self.state.entity_by_key(*keys_ref, &mut inserted); } Ok(()) @@ -300,7 +302,7 @@ impl + Sen .collect::>(); for group_entity in self.state.iter() { - let place: StateAddr = (*group_entity.get_state_value()).into(); + let place: StateAddr = group_entity.get_state_value().into(); for (function, state_offset) in functions.iter().zip(state_offsets.iter()) { unsafe { function.drop_state(place.next(*state_offset)) } diff --git a/src/query/service/src/pipelines/processors/transforms/aggregator/aggregator_partial.rs b/src/query/service/src/pipelines/processors/transforms/aggregator/aggregator_partial.rs index b24c3b547550d..61a9d6e837874 100644 --- a/src/query/service/src/pipelines/processors/transforms/aggregator/aggregator_partial.rs +++ b/src/query/service/src/pipelines/processors/transforms/aggregator/aggregator_partial.rs @@ -36,7 +36,8 @@ use common_functions::aggregates::StateAddrs; use crate::pipelines::processors::transforms::group_by::AggregatorState; use crate::pipelines::processors::transforms::group_by::KeysColumnBuilder; use crate::pipelines::processors::transforms::group_by::PolymorphicKeysHelper; -use crate::pipelines::processors::transforms::group_by::StateEntity; +use crate::pipelines::processors::transforms::group_by::StateEntityMutRef; +use crate::pipelines::processors::transforms::group_by::StateEntityRef; use crate::pipelines::processors::transforms::transform_aggregator::Aggregator; use crate::pipelines::processors::AggregatorParams; use crate::sessions::QueryContext; @@ -109,17 +110,18 @@ impl + S let mut inserted = true; for key in keys_iter { - let entity = state.entity(key, &mut inserted); + let unsafe_state = state as *mut Method::State; + let mut entity = state.entity(key, &mut inserted); match inserted { true => { - if let Some(place) = state.alloc_layout(params) { + if let Some(place) = unsafe { (*unsafe_state).alloc_layout(params) } { places.push(place); entity.set_state_value(place.addr()); } } false => { - let place: StateAddr = (*entity.get_state_value()).into(); + let place: StateAddr = entity.get_state_value().into(); places.push(place); } } @@ -205,7 +207,7 @@ impl + S let mut bytes = BytesMut::new(); for group_entity in self.state.iter() { - let place: StateAddr = (*group_entity.get_state_value()).into(); + let place: StateAddr = group_entity.get_state_value().into(); for (idx, func) in funcs.iter().enumerate() { let arg_place = place.next(offsets_aggregate_states[idx]); @@ -245,7 +247,7 @@ impl + Send> Aggregator let group_by_two_level_threshold = self.ctx.get_settings().get_group_by_two_level_threshold()? as usize; if !self.state.is_two_level() && self.state.len() >= group_by_two_level_threshold { - self.state.convert_to_two_level(); + self.state.convert_to_twolevel(); } let places = Self::lookup_state(&self.params, group_keys_iter, &mut self.state); @@ -274,7 +276,7 @@ impl + Send> Aggregator let group_by_two_level_threshold = self.ctx.get_settings().get_group_by_two_level_threshold()? as usize; if !self.state.is_two_level() && self.state.len() >= group_by_two_level_threshold { - self.state.convert_to_two_level(); + self.state.convert_to_twolevel(); } Self::lookup_key(group_keys_iter, &mut self.state); @@ -326,7 +328,7 @@ impl> .collect::>(); for group_entity in self.state.iter() { - let place: StateAddr = (*group_entity.get_state_value()).into(); + let place: StateAddr = group_entity.get_state_value().into(); for (function, state_offset) in functions.iter().zip(states.iter()) { unsafe { function.drop_state(place.next(*state_offset)) } diff --git a/src/query/service/src/pipelines/processors/transforms/group_by/aggregator_groups_builder.rs b/src/query/service/src/pipelines/processors/transforms/group_by/aggregator_groups_builder.rs index c42ec8e069331..95a02f54238fd 100644 --- a/src/query/service/src/pipelines/processors/transforms/group_by/aggregator_groups_builder.rs +++ b/src/query/service/src/pipelines/processors/transforms/group_by/aggregator_groups_builder.rs @@ -25,11 +25,11 @@ use common_datavalues::TypeID; use common_exception::Result; use common_io::prelude::FormatSettings; -use crate::pipelines::processors::transforms::group_by::keys_ref::KeysRef; use crate::pipelines::processors::AggregatorParams; -pub trait GroupColumnsBuilder { - fn append_value(&mut self, v: &Key); +pub trait GroupColumnsBuilder { + type T; + fn append_value(&mut self, v: Self::T); fn finish(self) -> Result>; } @@ -51,12 +51,14 @@ where for<'a> HashMethodFixedKeys: HashMethod } } -impl GroupColumnsBuilder for FixedKeysGroupColumnsBuilder +impl GroupColumnsBuilder for FixedKeysGroupColumnsBuilder where for<'a> HashMethodFixedKeys: HashMethod { + type T = T; + #[inline] - fn append_value(&mut self, v: &T) { - self.data.push(*v); + fn append_value(&mut self, v: T) { + self.data.push(v); } #[inline] @@ -74,12 +76,12 @@ where for<'a> HashMethodFixedKeys: HashMethod } } -pub struct SerializedKeysGroupColumnsBuilder { - data: Vec, +pub struct SerializedKeysGroupColumnsBuilder<'a> { + data: Vec<&'a [u8]>, group_data_types: Vec, } -impl SerializedKeysGroupColumnsBuilder { +impl<'a> SerializedKeysGroupColumnsBuilder<'a> { pub fn create(capacity: usize, params: &AggregatorParams) -> Self { Self { data: Vec::with_capacity(capacity), @@ -88,29 +90,24 @@ impl SerializedKeysGroupColumnsBuilder { } } -impl GroupColumnsBuilder for SerializedKeysGroupColumnsBuilder { - fn append_value(&mut self, v: &KeysRef) { - self.data.push(*v); - } +impl<'a> GroupColumnsBuilder for SerializedKeysGroupColumnsBuilder<'a> { + type T = &'a [u8]; - fn finish(self) -> Result> { - let mut keys = Vec::with_capacity(self.data.len()); + fn append_value(&mut self, v: &'a [u8]) { + self.data.push(v); + } - for v in &self.data { - unsafe { - let value = std::slice::from_raw_parts(v.address as *const u8, v.length); - keys.push(value); - } - } + fn finish(mut self) -> Result> { + let rows = self.data.len(); + let keys = self.data.as_mut_slice(); if self.group_data_types.len() == 1 && self.group_data_types[0].data_type_id() == TypeID::String { - let col = StringColumn::from_slice(&keys); + let col = StringColumn::from_slice(keys); return Ok(vec![col.arc()]); } - let rows = self.data.len(); let mut res = Vec::with_capacity(self.group_data_types.len()); let format = FormatSettings::default(); for data_type in self.group_data_types.iter() { diff --git a/src/query/service/src/pipelines/processors/transforms/group_by/aggregator_keys_builder.rs b/src/query/service/src/pipelines/processors/transforms/group_by/aggregator_keys_builder.rs index 2dff1b01c38e8..8a7abbbd194b5 100644 --- a/src/query/service/src/pipelines/processors/transforms/group_by/aggregator_keys_builder.rs +++ b/src/query/service/src/pipelines/processors/transforms/group_by/aggregator_keys_builder.rs @@ -18,12 +18,12 @@ use common_datablocks::HashMethod; use common_datablocks::HashMethodFixedKeys; use common_datavalues::prelude::*; -use crate::pipelines::processors::transforms::group_by::keys_ref::KeysRef; - /// Remove the group by key from the state and rebuild it into a column -pub trait KeysColumnBuilder { +pub trait KeysColumnBuilder { + type T; + + fn append_value(&mut self, v: Self::T); fn finish(self) -> ColumnRef; - fn append_value(&mut self, v: &Key); } pub struct FixedKeysColumnBuilder @@ -32,36 +32,47 @@ where T: PrimitiveType pub inner_builder: MutablePrimitiveColumn, } -impl KeysColumnBuilder for FixedKeysColumnBuilder +impl KeysColumnBuilder for FixedKeysColumnBuilder where T: PrimitiveType, for<'a> HashMethodFixedKeys: HashMethod, { + type T = T; + #[inline] fn finish(mut self) -> ColumnRef { self.inner_builder.to_column() } #[inline] - fn append_value(&mut self, v: &T) { - self.inner_builder.append_value(*v) + fn append_value(&mut self, v: T) { + self.inner_builder.append_value(v) } } -pub struct SerializedKeysColumnBuilder { +pub struct SerializedKeysColumnBuilder<'a> { pub inner_builder: MutableStringColumn, + _phantom: PhantomData<&'a ()>, +} + +impl<'a> SerializedKeysColumnBuilder<'a> { + pub fn create(capacity: usize) -> Self { + SerializedKeysColumnBuilder { + inner_builder: MutableStringColumn::with_capacity(capacity), + _phantom: PhantomData, + } + } } -impl KeysColumnBuilder for SerializedKeysColumnBuilder { +impl<'a> KeysColumnBuilder for SerializedKeysColumnBuilder<'a> { + type T = &'a [u8]; + fn finish(mut self) -> ColumnRef { self.inner_builder.to_column() } - fn append_value(&mut self, v: &KeysRef) { - unsafe { - let value = std::slice::from_raw_parts(v.address as *const u8, v.length); - self.inner_builder.append_value(value); - } + fn append_value(&mut self, v: &'a [u8]) { + self.inner_builder.append_value(v); } } @@ -72,18 +83,20 @@ where T: LargePrimitive pub _t: PhantomData, } -impl KeysColumnBuilder for LargeFixedKeysColumnBuilder +impl KeysColumnBuilder for LargeFixedKeysColumnBuilder where T: LargePrimitive, for<'a> HashMethodFixedKeys: HashMethod, { + type T = T; + #[inline] fn finish(mut self) -> ColumnRef { self.inner_builder.to_column() } #[inline] - fn append_value(&mut self, v: &T) { + fn append_value(&mut self, v: T) { let values = self.inner_builder.values_mut(); let new_len = values.len() + T::BYTE_SIZE; values.resize(new_len, 0); diff --git a/src/query/service/src/pipelines/processors/transforms/group_by/aggregator_keys_iter.rs b/src/query/service/src/pipelines/processors/transforms/group_by/aggregator_keys_iter.rs index a6d31493d6add..812084b7684d2 100644 --- a/src/query/service/src/pipelines/processors/transforms/group_by/aggregator_keys_iter.rs +++ b/src/query/service/src/pipelines/processors/transforms/group_by/aggregator_keys_iter.rs @@ -20,8 +20,6 @@ use common_datavalues::ScalarColumn; use common_datavalues::StringColumn; use common_exception::Result; -use crate::pipelines::processors::transforms::group_by::keys_ref::KeysRef; - pub trait KeysColumnIter { fn get_slice(&self) -> &[T]; } @@ -77,14 +75,14 @@ where T: LargePrimitive } } -pub struct SerializedKeysColumnIter { - inner: Vec, +pub struct SerializedKeysColumnIter<'a> { + inner: Vec<&'a [u8]>, #[allow(unused)] column: StringColumn, } -impl SerializedKeysColumnIter { - pub fn create(column: &StringColumn) -> Result { +impl<'a> SerializedKeysColumnIter<'a> { + pub fn create(column: &'a StringColumn) -> Result> { let values = column.values(); let offsets = column.offsets(); @@ -92,8 +90,10 @@ impl SerializedKeysColumnIter { for index in 0..(offsets.len() - 1) { let offset = offsets[index] as usize; let offset_1 = offsets[index + 1] as usize; - let address = unsafe { values.as_ptr().add(offset) as usize }; - inner.push(KeysRef::create(address, offset_1 - offset)); + unsafe { + let address = values.as_ptr().add(offset) as *const u8; + inner.push(std::slice::from_raw_parts(address, offset_1 - offset)); + } } Ok(SerializedKeysColumnIter { @@ -103,8 +103,8 @@ impl SerializedKeysColumnIter { } } -impl KeysColumnIter for SerializedKeysColumnIter { - fn get_slice(&self) -> &[KeysRef] { +impl<'a> KeysColumnIter<&'a [u8]> for SerializedKeysColumnIter<'a> { + fn get_slice(&self) -> &[&'a [u8]] { self.inner.as_slice() } } diff --git a/src/query/service/src/pipelines/processors/transforms/group_by/aggregator_polymorphic_keys.rs b/src/query/service/src/pipelines/processors/transforms/group_by/aggregator_polymorphic_keys.rs index 12e112c53f0a4..c90435391f854 100644 --- a/src/query/service/src/pipelines/processors/transforms/group_by/aggregator_polymorphic_keys.rs +++ b/src/query/service/src/pipelines/processors/transforms/group_by/aggregator_polymorphic_keys.rs @@ -23,7 +23,7 @@ use common_datablocks::HashMethodKeysU512; use common_datablocks::HashMethodSerializer; use common_datavalues::prelude::*; use common_exception::Result; -use common_hashtable::HashMapKind; +use common_hashtable::UnsizedHashMap; use primitive_types::U256; use primitive_types::U512; @@ -84,18 +84,34 @@ pub trait PolymorphicKeysHelper { type State: AggregatorState; fn aggregate_state(&self) -> Self::State; - type ColumnBuilder: KeysColumnBuilder<>::Key>; - fn keys_column_builder(&self, capacity: usize) -> Self::ColumnBuilder; + type ColumnBuilder<'a>: KeysColumnBuilder< + T = >::KeyRef<'a>, + > + where + Self: 'a, + Self::State: 'a; - type KeysColumnIter: KeysColumnIter<>::Key>; - fn keys_iter_from_column(&self, column: &ColumnRef) -> Result; + fn keys_column_builder(&self, capacity: usize) -> Self::ColumnBuilder<'_>; - type GroupColumnsBuilder: GroupColumnsBuilder<>::Key>; - fn group_columns_builder( - &self, + type KeysColumnIter<'a>: KeysColumnIter<>::KeyRef<'a>> + where + Self: 'a, + Self::State: 'a; + + fn keys_iter_from_column<'a>(&self, column: &'a ColumnRef) -> Result>; + + type GroupColumnsBuilder<'a>: GroupColumnsBuilder< + T = >::KeyRef<'a>, + > + where + Self: 'a, + Self::State: 'a; + + fn group_columns_builder<'a>( + &'a self, capacity: usize, params: &AggregatorParams, - ) -> Self::GroupColumnsBuilder; + ) -> Self::GroupColumnsBuilder<'a>; } impl PolymorphicKeysHelper> for HashMethodFixedKeys { @@ -103,22 +119,22 @@ impl PolymorphicKeysHelper> for HashMethodFixedKeys fn aggregate_state(&self) -> Self::State { Self::State::create((u8::MAX as usize) + 1) } - type ColumnBuilder = FixedKeysColumnBuilder; - fn keys_column_builder(&self, capacity: usize) -> Self::ColumnBuilder { + type ColumnBuilder<'a> = FixedKeysColumnBuilder; + fn keys_column_builder(&self, capacity: usize) -> FixedKeysColumnBuilder { FixedKeysColumnBuilder:: { inner_builder: MutablePrimitiveColumn::::with_capacity(capacity), } } - type KeysColumnIter = FixedKeysColumnIter; - fn keys_iter_from_column(&self, column: &ColumnRef) -> Result { + type KeysColumnIter<'a> = FixedKeysColumnIter; + fn keys_iter_from_column<'a>(&self, column: &'a ColumnRef) -> Result> { FixedKeysColumnIter::create(Series::check_get::>(column)?) } - type GroupColumnsBuilder = FixedKeysGroupColumnsBuilder; + type GroupColumnsBuilder<'a> = FixedKeysGroupColumnsBuilder; fn group_columns_builder( &self, capacity: usize, params: &AggregatorParams, - ) -> Self::GroupColumnsBuilder { + ) -> FixedKeysGroupColumnsBuilder { FixedKeysGroupColumnsBuilder::::create(capacity, params) } } @@ -128,22 +144,22 @@ impl PolymorphicKeysHelper> for HashMethodFixedKeys Self::State { Self::State::create((u16::MAX as usize) + 1) } - type ColumnBuilder = FixedKeysColumnBuilder; - fn keys_column_builder(&self, capacity: usize) -> Self::ColumnBuilder { + type ColumnBuilder<'a> = FixedKeysColumnBuilder; + fn keys_column_builder(&self, capacity: usize) -> FixedKeysColumnBuilder { FixedKeysColumnBuilder:: { inner_builder: MutablePrimitiveColumn::::with_capacity(capacity), } } - type KeysColumnIter = FixedKeysColumnIter; - fn keys_iter_from_column(&self, column: &ColumnRef) -> Result { + type KeysColumnIter<'a> = FixedKeysColumnIter; + fn keys_iter_from_column<'a>(&self, column: &'a ColumnRef) -> Result> { FixedKeysColumnIter::create(Series::check_get::>(column)?) } - type GroupColumnsBuilder = FixedKeysGroupColumnsBuilder; + type GroupColumnsBuilder<'a> = FixedKeysGroupColumnsBuilder; fn group_columns_builder( &self, capacity: usize, params: &AggregatorParams, - ) -> Self::GroupColumnsBuilder { + ) -> FixedKeysGroupColumnsBuilder { FixedKeysGroupColumnsBuilder::::create(capacity, params) } } @@ -153,22 +169,22 @@ impl PolymorphicKeysHelper> for HashMethodFixedKeys Self::State { Self::State::default() } - type ColumnBuilder = FixedKeysColumnBuilder; - fn keys_column_builder(&self, capacity: usize) -> Self::ColumnBuilder { + type ColumnBuilder<'a> = FixedKeysColumnBuilder; + fn keys_column_builder(&self, capacity: usize) -> FixedKeysColumnBuilder { FixedKeysColumnBuilder:: { inner_builder: MutablePrimitiveColumn::::with_capacity(capacity), } } - type KeysColumnIter = FixedKeysColumnIter; - fn keys_iter_from_column(&self, column: &ColumnRef) -> Result { + type KeysColumnIter<'a> = FixedKeysColumnIter; + fn keys_iter_from_column<'a>(&self, column: &'a ColumnRef) -> Result> { FixedKeysColumnIter::create(Series::check_get::>(column)?) } - type GroupColumnsBuilder = FixedKeysGroupColumnsBuilder; + type GroupColumnsBuilder<'a> = FixedKeysGroupColumnsBuilder; fn group_columns_builder( &self, capacity: usize, params: &AggregatorParams, - ) -> Self::GroupColumnsBuilder { + ) -> FixedKeysGroupColumnsBuilder { FixedKeysGroupColumnsBuilder::::create(capacity, params) } } @@ -178,22 +194,22 @@ impl PolymorphicKeysHelper> for HashMethodFixedKeys Self::State { Self::State::default() } - type ColumnBuilder = FixedKeysColumnBuilder; - fn keys_column_builder(&self, capacity: usize) -> Self::ColumnBuilder { + type ColumnBuilder<'a> = FixedKeysColumnBuilder; + fn keys_column_builder(&self, capacity: usize) -> FixedKeysColumnBuilder { FixedKeysColumnBuilder:: { inner_builder: MutablePrimitiveColumn::::with_capacity(capacity), } } - type KeysColumnIter = FixedKeysColumnIter; - fn keys_iter_from_column(&self, column: &ColumnRef) -> Result { + type KeysColumnIter<'a> = FixedKeysColumnIter; + fn keys_iter_from_column<'a>(&self, column: &'a ColumnRef) -> Result> { FixedKeysColumnIter::create(Series::check_get::>(column)?) } - type GroupColumnsBuilder = FixedKeysGroupColumnsBuilder; + type GroupColumnsBuilder<'a> = FixedKeysGroupColumnsBuilder; fn group_columns_builder( &self, capacity: usize, params: &AggregatorParams, - ) -> Self::GroupColumnsBuilder { + ) -> FixedKeysGroupColumnsBuilder { FixedKeysGroupColumnsBuilder::::create(capacity, params) } } @@ -204,25 +220,25 @@ impl PolymorphicKeysHelper for HashMethodKeysU128 { Self::State::default() } - type ColumnBuilder = LargeFixedKeysColumnBuilder; - fn keys_column_builder(&self, capacity: usize) -> Self::ColumnBuilder { + type ColumnBuilder<'a> = LargeFixedKeysColumnBuilder; + fn keys_column_builder(&self, capacity: usize) -> LargeFixedKeysColumnBuilder { LargeFixedKeysColumnBuilder { inner_builder: MutableStringColumn::with_capacity(capacity * 16), _t: PhantomData, } } - type KeysColumnIter = LargeFixedKeysColumnIter; - fn keys_iter_from_column(&self, column: &ColumnRef) -> Result { + type KeysColumnIter<'a> = LargeFixedKeysColumnIter; + fn keys_iter_from_column<'a>(&self, column: &'a ColumnRef) -> Result> { LargeFixedKeysColumnIter::create(Series::check_get::(column)?) } - type GroupColumnsBuilder = FixedKeysGroupColumnsBuilder; + type GroupColumnsBuilder<'a> = FixedKeysGroupColumnsBuilder; fn group_columns_builder( &self, capacity: usize, params: &AggregatorParams, - ) -> Self::GroupColumnsBuilder { + ) -> FixedKeysGroupColumnsBuilder { FixedKeysGroupColumnsBuilder::create(capacity, params) } } @@ -233,25 +249,25 @@ impl PolymorphicKeysHelper for HashMethodKeysU256 { Self::State::default() } - type ColumnBuilder = LargeFixedKeysColumnBuilder; - fn keys_column_builder(&self, capacity: usize) -> Self::ColumnBuilder { + type ColumnBuilder<'a> = LargeFixedKeysColumnBuilder; + fn keys_column_builder(&self, capacity: usize) -> LargeFixedKeysColumnBuilder { LargeFixedKeysColumnBuilder { inner_builder: MutableStringColumn::with_capacity(capacity * 32), _t: PhantomData, } } - type KeysColumnIter = LargeFixedKeysColumnIter; - fn keys_iter_from_column(&self, column: &ColumnRef) -> Result { + type KeysColumnIter<'a> = LargeFixedKeysColumnIter; + fn keys_iter_from_column<'a>(&self, column: &'a ColumnRef) -> Result> { LargeFixedKeysColumnIter::create(Series::check_get::(column)?) } - type GroupColumnsBuilder = FixedKeysGroupColumnsBuilder; + type GroupColumnsBuilder<'a> = FixedKeysGroupColumnsBuilder; fn group_columns_builder( &self, capacity: usize, params: &AggregatorParams, - ) -> Self::GroupColumnsBuilder { + ) -> FixedKeysGroupColumnsBuilder { FixedKeysGroupColumnsBuilder::create(capacity, params) } } @@ -262,25 +278,25 @@ impl PolymorphicKeysHelper for HashMethodKeysU512 { Self::State::default() } - type ColumnBuilder = LargeFixedKeysColumnBuilder; - fn keys_column_builder(&self, capacity: usize) -> Self::ColumnBuilder { + type ColumnBuilder<'a> = LargeFixedKeysColumnBuilder; + fn keys_column_builder(&self, capacity: usize) -> LargeFixedKeysColumnBuilder { LargeFixedKeysColumnBuilder { inner_builder: MutableStringColumn::with_capacity(capacity * 64), _t: PhantomData, } } - type KeysColumnIter = LargeFixedKeysColumnIter; - fn keys_iter_from_column(&self, column: &ColumnRef) -> Result { + type KeysColumnIter<'a> = LargeFixedKeysColumnIter; + fn keys_iter_from_column<'a>(&self, column: &'a ColumnRef) -> Result> { LargeFixedKeysColumnIter::create(Series::check_get::(column)?) } - type GroupColumnsBuilder = FixedKeysGroupColumnsBuilder; + type GroupColumnsBuilder<'a> = FixedKeysGroupColumnsBuilder; fn group_columns_builder( &self, capacity: usize, params: &AggregatorParams, - ) -> Self::GroupColumnsBuilder { + ) -> FixedKeysGroupColumnsBuilder { FixedKeysGroupColumnsBuilder::create(capacity, params) } } @@ -289,31 +305,28 @@ impl PolymorphicKeysHelper for HashMethodSerializer { type State = SerializedKeysAggregatorState; fn aggregate_state(&self) -> Self::State { SerializedKeysAggregatorState { - keys_area: Bump::new(), - state_area: Bump::new(), - data_state_map: HashMapKind::create_hash_table(), + area: Bump::new(), + data_state_map: UnsizedHashMap::new(), two_level_flag: false, } } - type ColumnBuilder = SerializedKeysColumnBuilder; - fn keys_column_builder(&self, capacity: usize) -> Self::ColumnBuilder { - SerializedKeysColumnBuilder { - inner_builder: MutableStringColumn::with_capacity(capacity), - } + type ColumnBuilder<'a> = SerializedKeysColumnBuilder<'a>; + fn keys_column_builder(&self, capacity: usize) -> SerializedKeysColumnBuilder<'_> { + SerializedKeysColumnBuilder::create(capacity) } - type KeysColumnIter = SerializedKeysColumnIter; - fn keys_iter_from_column(&self, column: &ColumnRef) -> Result { + type KeysColumnIter<'a> = SerializedKeysColumnIter<'a>; + fn keys_iter_from_column<'a>(&self, column: &'a ColumnRef) -> Result> { SerializedKeysColumnIter::create(Series::check_get::(column)?) } - type GroupColumnsBuilder = SerializedKeysGroupColumnsBuilder; + type GroupColumnsBuilder<'a> = SerializedKeysGroupColumnsBuilder<'a>; fn group_columns_builder( &self, capacity: usize, params: &AggregatorParams, - ) -> Self::GroupColumnsBuilder { + ) -> SerializedKeysGroupColumnsBuilder<'_> { SerializedKeysGroupColumnsBuilder::create(capacity, params) } } diff --git a/src/query/service/src/pipelines/processors/transforms/group_by/aggregator_state.rs b/src/query/service/src/pipelines/processors/transforms/group_by/aggregator_state.rs index 8cb89c5236438..62ee7ae4bd668 100644 --- a/src/query/service/src/pipelines/processors/transforms/group_by/aggregator_state.rs +++ b/src/query/service/src/pipelines/processors/transforms/group_by/aggregator_state.rs @@ -17,23 +17,26 @@ use std::alloc::Layout; use std::intrinsics::likely; use bumpalo::Bump; -use common_base::mem_allocator::ALLOC; +use common_base::mem_allocator::GlobalAllocator; use common_datablocks::HashMethod; use common_datablocks::HashMethodFixedKeys; use common_datablocks::HashMethodSerializer; use common_datavalues::prelude::*; use common_functions::aggregates::StateAddr; -use common_hashtable::HashMapIteratorKind; use common_hashtable::HashMapKind; -use common_hashtable::HashTableEntity; -use common_hashtable::HashTableKeyable; -use common_hashtable::KeyValueEntity; +use common_hashtable::HashMapKindIter; +use common_hashtable::HashtableEntry; +use common_hashtable::HashtableKeyable; +use common_hashtable::UnsizedHashMap; +use common_hashtable::UnsizedHashMapIter; +use common_hashtable::UnsizedHashtableEntryMutRef; +use common_hashtable::UnsizedHashtableEntryRef; use crate::pipelines::processors::transforms::group_by::aggregator_state_entity::ShortFixedKeyable; use crate::pipelines::processors::transforms::group_by::aggregator_state_entity::ShortFixedKeysStateEntity; -use crate::pipelines::processors::transforms::group_by::aggregator_state_entity::StateEntity; +use crate::pipelines::processors::transforms::group_by::aggregator_state_entity::StateEntityMutRef; +use crate::pipelines::processors::transforms::group_by::aggregator_state_entity::StateEntityRef; use crate::pipelines::processors::transforms::group_by::aggregator_state_iterator::ShortFixedKeysStateIterator; -use crate::pipelines::processors::transforms::group_by::keys_ref::KeysRef; use crate::pipelines::processors::AggregatorParams as NewAggregatorParams; /// Aggregate state of the SELECT query, destroy when group by is completed. @@ -44,13 +47,19 @@ use crate::pipelines::processors::AggregatorParams as NewAggregatorParams; /// - Group by key data memory pool (if necessary) #[allow(clippy::len_without_is_empty)] pub trait AggregatorState: Sync + Send { - type Key; - type Entity: StateEntity; - type Iterator: Iterator; + type Key: ?Sized; + type KeyRef<'a>: Copy + 'a + where Self: 'a; + type EntityRef<'a>: StateEntityRef> + where Self: 'a; + type EntityMutRef<'a>: StateEntityMutRef> + where Self: 'a; + type Iterator<'a>: Iterator> + where Self: 'a; fn len(&self) -> usize; - fn iter(&self) -> Self::Iterator; + fn iter(&self) -> Self::Iterator<'_>; fn alloc_place(&self, layout: Layout) -> StateAddr; @@ -66,15 +75,23 @@ pub trait AggregatorState: Sync + Send { Some(place) } - fn entity(&mut self, key: Method::HashKeyRef<'_>, inserted: &mut bool) -> *mut Self::Entity; + fn entity( + &mut self, + key: Method::HashKeyRef<'_>, + inserted: &mut bool, + ) -> Self::EntityMutRef<'_>; - fn entity_by_key(&mut self, key: &Self::Key, inserted: &mut bool) -> *mut Self::Entity; + fn entity_by_key<'a>( + &mut self, + key: Self::KeyRef<'a>, + inserted: &mut bool, + ) -> Self::EntityMutRef<'_>; fn is_two_level(&self) -> bool { false } - fn convert_to_two_level(&mut self) {} + fn convert_to_twolevel(&mut self) {} } /// The fixed length array is used as the data structure to locate the key by subscript @@ -103,7 +120,7 @@ impl ShortFixedKeysAggregatorState { let entity_align = std::mem::align_of::>(); let entity_layout = Layout::from_size_align_unchecked(size, entity_align); - let raw_ptr = ALLOC.alloc_zeroed(entity_layout); + let raw_ptr = GlobalAllocator.alloc_zeroed(entity_layout); ShortFixedKeysAggregatorState:: { area: Default::default(), @@ -122,7 +139,7 @@ impl Drop for ShortFixedKeysAggregatorState { let size = self.max_size * std::mem::size_of::>(); let entity_align = std::mem::align_of::>(); let layout = Layout::from_size_align_unchecked(size, entity_align); - ALLOC.dealloc(self.data as *mut u8, layout); + GlobalAllocator.dealloc(self.data as *mut u8, layout); } } } @@ -131,11 +148,13 @@ impl AggregatorState> for ShortFixedKeysAggregatorStat where T: PrimitiveType + ShortFixedKeyable, for<'a> HashMethodFixedKeys: HashMethod = T>, - for<'a> as HashMethod>::HashKey: HashTableKeyable, + for<'a> as HashMethod>::HashKey: HashtableKeyable, { type Key = T; - type Entity = ShortFixedKeysStateEntity; - type Iterator = ShortFixedKeysStateIterator; + type KeyRef<'a> = T; + type EntityRef<'a> = &'a ShortFixedKeysStateEntity; + type EntityMutRef<'a> = &'a mut ShortFixedKeysStateEntity; + type Iterator<'a> = ShortFixedKeysStateIterator<'a, T>; #[inline(always)] fn len(&self) -> usize { @@ -143,8 +162,11 @@ where } #[inline(always)] - fn iter(&self) -> Self::Iterator { - Self::Iterator::create(self.data, self.max_size as isize) + fn iter(&self) -> Self::Iterator<'_> { + unsafe { + let data = std::slice::from_raw_parts_mut(self.data, self.max_size); + Self::Iterator::create(data) + } } #[inline(always)] @@ -153,27 +175,32 @@ where } #[inline(always)] - fn entity(&mut self, key: T, inserted: &mut bool) -> *mut Self::Entity { + fn entity(&mut self, key: T, inserted: &mut bool) -> Self::EntityMutRef<'_> { unsafe { let index = key.lookup(); let value = self.data.offset(index); if likely((*value).fill) { *inserted = false; - return value; + return &mut *(value); } *inserted = true; self.size += 1; (*value).key = key; (*value).fill = true; - value + // It's an undefined behavior, but uninitialized integers rarely lead to problems. + &mut (*value) } } #[inline(always)] - fn entity_by_key(&mut self, key: &Self::Key, inserted: &mut bool) -> *mut Self::Entity { - self.entity(*key, inserted) + fn entity_by_key<'a>( + &mut self, + key: Self::KeyRef<'a>, + inserted: &mut bool, + ) -> Self::EntityMutRef<'_> { + self.entity(key, inserted) } #[inline(always)] @@ -182,22 +209,22 @@ where } #[inline(always)] - fn convert_to_two_level(&mut self) { + fn convert_to_twolevel(&mut self) { self.two_level_flag = true; } } -pub struct LongerFixedKeysAggregatorState { +pub struct LongerFixedKeysAggregatorState { pub area: Bump, pub data: HashMapKind, pub two_level_flag: bool, } -impl Default for LongerFixedKeysAggregatorState { +impl Default for LongerFixedKeysAggregatorState { fn default() -> Self { Self { area: Bump::new(), - data: HashMapKind::create_hash_table(), + data: HashMapKind::new(), two_level_flag: false, } } @@ -207,22 +234,24 @@ impl Default for LongerFixedKeysAggregatorState { // The *mut KeyValueEntity needs to be used externally, but we can ensure that *mut KeyValueEntity // will not be used multiple async, so KeyValueEntity is Send #[allow(clippy::non_send_fields_in_send_ty)] -unsafe impl Send for LongerFixedKeysAggregatorState {} +unsafe impl Send for LongerFixedKeysAggregatorState {} // TODO:(Winter) Hack: // The *mut KeyValueEntity needs to be used externally, but we can ensure that &*mut KeyValueEntity // will not be used multiple async, so KeyValueEntity is Sync -unsafe impl Sync for LongerFixedKeysAggregatorState {} +unsafe impl Sync for LongerFixedKeysAggregatorState {} impl AggregatorState> for LongerFixedKeysAggregatorState where for<'a> HashMethodFixedKeys: HashMethod = T>, - for<'a> as HashMethod>::HashKey: HashTableKeyable, + for<'a> as HashMethod>::HashKey: HashtableKeyable, { type Key = T; - type Entity = KeyValueEntity; - type Iterator = HashMapIteratorKind; + type KeyRef<'a> = T; + type EntityRef<'a> = &'a HashtableEntry; + type EntityMutRef<'a> = &'a mut HashtableEntry; + type Iterator<'a> = HashMapKindIter<'a, T, usize>; #[inline(always)] fn len(&self) -> usize { @@ -230,7 +259,7 @@ where } #[inline(always)] - fn iter(&self) -> Self::Iterator { + fn iter(&self) -> Self::Iterator<'_> { self.data.iter() } @@ -240,13 +269,26 @@ where } #[inline(always)] - fn entity(&mut self, key: Self::Key, inserted: &mut bool) -> *mut Self::Entity { - self.data.insert_key(&key, inserted) + fn entity(&mut self, key: Self::Key, inserted: &mut bool) -> Self::EntityMutRef<'_> { + match unsafe { self.data.insert_and_entry(key) } { + Ok(e) => { + *inserted = true; + e + } + Err(e) => { + *inserted = false; + e + } + } } #[inline(always)] - fn entity_by_key(&mut self, key: &Self::Key, inserted: &mut bool) -> *mut Self::Entity { - self.entity(*key, inserted) + fn entity_by_key<'a>( + &mut self, + key: Self::KeyRef<'a>, + inserted: &mut bool, + ) -> Self::EntityMutRef<'_> { + self.entity(key, inserted) } #[inline(always)] @@ -255,18 +297,15 @@ where } #[inline(always)] - fn convert_to_two_level(&mut self) { - unsafe { - self.data.convert_to_two_level(); - } + fn convert_to_twolevel(&mut self) { + self.data.convert_to_twolevel(); self.two_level_flag = true; } } pub struct SerializedKeysAggregatorState { - pub keys_area: Bump, - pub state_area: Bump, - pub data_state_map: HashMapKind, + pub area: Bump, + pub data_state_map: UnsizedHashMap<[u8], usize>, pub two_level_flag: bool, } @@ -282,59 +321,58 @@ unsafe impl Send for SerializedKeysAggregatorState {} unsafe impl Sync for SerializedKeysAggregatorState {} impl AggregatorState for SerializedKeysAggregatorState { - type Key = KeysRef; - type Entity = KeyValueEntity; - type Iterator = HashMapIteratorKind; + type Key = [u8]; + type KeyRef<'a> = &'a [u8]; + type EntityRef<'a> = UnsizedHashtableEntryRef<'a, [u8], usize>; + type EntityMutRef<'a> = UnsizedHashtableEntryMutRef<'a, [u8], usize>; + type Iterator<'a> = UnsizedHashMapIter<'a, [u8], usize>; fn len(&self) -> usize { self.data_state_map.len() } - fn iter(&self) -> Self::Iterator { + fn iter(&self) -> Self::Iterator<'_> { self.data_state_map.iter() } #[inline(always)] fn alloc_place(&self, layout: Layout) -> StateAddr { - self.state_area.alloc_layout(layout).into() + self.area.alloc_layout(layout).into() } #[inline(always)] - fn entity(&mut self, keys: &[u8], inserted: &mut bool) -> *mut Self::Entity { - let mut keys_ref = KeysRef::create(keys.as_ptr() as usize, keys.len()); - let state_entity = self.data_state_map.insert_key(&keys_ref, inserted); - - if *inserted { - unsafe { - // Keys will be destroyed after call we need copy the keys to the memory pool. - let global_keys = self.keys_area.alloc_slice_copy(keys); - let inserted_hash = state_entity.get_hash(); - keys_ref.address = global_keys.as_ptr() as usize; - // TODO: maybe need set key method. - state_entity.set_key_and_hash(&keys_ref, inserted_hash) + fn entity(&mut self, keys: &[u8], inserted: &mut bool) -> Self::EntityMutRef<'_> { + unsafe { + match self.data_state_map.insert_and_entry(keys) { + Ok(e) => { + *inserted = true; + e + } + Err(e) => { + *inserted = false; + e + } } } - - state_entity } #[inline(always)] - fn entity_by_key(&mut self, keys_ref: &KeysRef, inserted: &mut bool) -> *mut Self::Entity { - let state_entity = self.data_state_map.insert_key(keys_ref, inserted); - - if *inserted { - unsafe { - // Keys will be destroyed after call we need copy the keys to the memory pool. - let data_ptr = keys_ref.address as *mut u8; - let keys = std::slice::from_raw_parts_mut(data_ptr, keys_ref.length); - let global_keys = self.keys_area.alloc_slice_copy(keys); - let inserted_hash = state_entity.get_hash(); - let address = global_keys.as_ptr() as usize; - let new_keys_ref = KeysRef::create(address, keys_ref.length); - state_entity.set_key_and_hash(&new_keys_ref, inserted_hash) + fn entity_by_key<'a>( + &mut self, + key_ref: Self::KeyRef<'a>, + inserted: &mut bool, + ) -> Self::EntityMutRef<'_> { + unsafe { + match self.data_state_map.insert_and_entry(key_ref) { + Ok(e) => { + *inserted = true; + e + } + Err(e) => { + *inserted = false; + e + } } } - - state_entity } #[inline(always)] @@ -343,10 +381,8 @@ impl AggregatorState for SerializedKeysAggregatorState { } #[inline(always)] - fn convert_to_two_level(&mut self) { - unsafe { - self.data_state_map.convert_to_two_level(); - } + fn convert_to_twolevel(&mut self) { + // self.data_state_map.convert_to_twolevel(); self.two_level_flag = true; } } diff --git a/src/query/service/src/pipelines/processors/transforms/group_by/aggregator_state_entity.rs b/src/query/service/src/pipelines/processors/transforms/group_by/aggregator_state_entity.rs index f9af0dc4961f2..1ef1fc3551079 100644 --- a/src/query/service/src/pipelines/processors/transforms/group_by/aggregator_state_entity.rs +++ b/src/query/service/src/pipelines/processors/transforms/group_by/aggregator_state_entity.rs @@ -12,14 +12,24 @@ // See the License for the specific language governing permissions and // limitations under the License. -use common_hashtable::HashTableEntity; -use common_hashtable::HashTableKeyable; -use common_hashtable::KeyValueEntity; - -pub trait StateEntity { - fn get_state_key<'a>(self: *mut Self) -> &'a Key; - fn set_state_value(self: *mut Self, value: usize); - fn get_state_value<'a>(self: *mut Self) -> &'a usize; +use common_hashtable::HashtableEntry; +use common_hashtable::HashtableKeyable; +use common_hashtable::UnsizedHashtableEntryMutRef; +use common_hashtable::UnsizedHashtableEntryRef; + +pub trait StateEntityRef { + type KeyRef: Copy; + + fn get_state_key(&self) -> Self::KeyRef; + fn get_state_value(&self) -> usize; +} + +pub trait StateEntityMutRef { + type KeyRef: Copy; + + fn get_state_key(&self) -> Self::KeyRef; + fn get_state_value(&self) -> usize; + fn set_state_value(&mut self, value: usize); } pub trait ShortFixedKeyable: Sized + Clone { @@ -33,37 +43,104 @@ pub struct ShortFixedKeysStateEntity { pub fill: bool, } -impl StateEntity for ShortFixedKeysStateEntity { +impl<'a, Key: ShortFixedKeyable + Copy> StateEntityRef for &'a ShortFixedKeysStateEntity { + type KeyRef = Key; + + #[inline(always)] + fn get_state_key(&self) -> Key { + self.key + } + + #[inline(always)] + fn get_state_value(&self) -> usize { + self.value + } +} + +impl<'a, Key: ShortFixedKeyable + Copy> StateEntityMutRef + for &'a mut ShortFixedKeysStateEntity +{ + type KeyRef = Key; + + #[inline(always)] + fn get_state_key(&self) -> Key { + self.key + } + + #[inline(always)] + fn get_state_value(&self) -> usize { + self.value + } + + #[inline(always)] + fn set_state_value(&mut self, value: usize) { + self.value = value; + } +} + +impl<'a, Key: HashtableKeyable> StateEntityRef for &'a HashtableEntry { + type KeyRef = Key; + #[inline(always)] - fn get_state_key<'a>(self: *mut Self) -> &'a Key { - unsafe { &(*self).key } + fn get_state_key(&self) -> Key { + *self.key() } #[inline(always)] - fn set_state_value(self: *mut Self, value: usize) { - unsafe { (*self).value = value } + fn get_state_value(&self) -> usize { + *self.get() } +} + +impl<'a, Key: HashtableKeyable> StateEntityMutRef for &'a mut HashtableEntry { + type KeyRef = Key; #[inline(always)] - fn get_state_value<'a>(self: *mut Self) -> &'a usize { - unsafe { &(*self).value } + fn get_state_key(&self) -> Key { + *self.key() + } + + #[inline(always)] + fn get_state_value(&self) -> usize { + *self.get() + } + + #[inline(always)] + fn set_state_value(&mut self, value: usize) { + *self.get_mut() = value; + } +} + +impl<'a> StateEntityRef for UnsizedHashtableEntryRef<'a, [u8], usize> { + type KeyRef = &'a [u8]; + + #[inline(always)] + fn get_state_key(&self) -> &'a [u8] { + self.key() + } + + #[inline(always)] + fn get_state_value(&self) -> usize { + *self.get() } } -impl StateEntity for KeyValueEntity { +impl<'a> StateEntityMutRef for UnsizedHashtableEntryMutRef<'a, [u8], usize> { + type KeyRef = &'a [u8]; + #[inline(always)] - fn get_state_key<'a>(self: *mut Self) -> &'a Key { - self.get_key() + fn get_state_key(&self) -> &'a [u8] { + self.key() } #[inline(always)] - fn set_state_value(self: *mut Self, value: usize) { - self.set_value(value) + fn get_state_value(&self) -> usize { + *self.get() } #[inline(always)] - fn get_state_value<'a>(self: *mut Self) -> &'a usize { - self.get_value() + fn set_state_value(&mut self, value: usize) { + *self.get_mut() = value; } } diff --git a/src/query/service/src/pipelines/processors/transforms/group_by/aggregator_state_iterator.rs b/src/query/service/src/pipelines/processors/transforms/group_by/aggregator_state_iterator.rs index f017548a61d72..9ee49ff5f5acf 100644 --- a/src/query/service/src/pipelines/processors/transforms/group_by/aggregator_state_iterator.rs +++ b/src/query/service/src/pipelines/processors/transforms/group_by/aggregator_state_iterator.rs @@ -15,39 +15,30 @@ use crate::pipelines::processors::transforms::group_by::aggregator_state_entity::ShortFixedKeyable; use crate::pipelines::processors::transforms::group_by::aggregator_state_entity::ShortFixedKeysStateEntity; -type Entities = *mut ShortFixedKeysStateEntity; - -pub struct ShortFixedKeysStateIterator { - index: isize, - capacity: isize, - entities: *mut ShortFixedKeysStateEntity, +pub struct ShortFixedKeysStateIterator<'a, Key: ShortFixedKeyable> { + index: usize, + entities: &'a [ShortFixedKeysStateEntity], } -impl ShortFixedKeysStateIterator { - pub fn create(entities: Entities, capacity: isize) -> Self { - ShortFixedKeysStateIterator:: { - index: 0, - capacity, - entities, - } +impl<'a, Key: ShortFixedKeyable> ShortFixedKeysStateIterator<'a, Key> { + pub fn create(entities: &'a [ShortFixedKeysStateEntity]) -> Self { + Self { index: 0, entities } } } -impl Iterator for ShortFixedKeysStateIterator { - type Item = *mut ShortFixedKeysStateEntity; +impl<'a, Key: ShortFixedKeyable> Iterator for ShortFixedKeysStateIterator<'a, Key> { + type Item = &'a ShortFixedKeysStateEntity; fn next(&mut self) -> Option { - unsafe { - while self.index < self.capacity { - let entity = self.entities.offset(self.index); - self.index += 1; + while self.index < self.entities.len() { + let entity = &self.entities[self.index]; + self.index += 1; - if (*entity).fill { - return Some(entity); - } + if entity.fill { + return Some(entity); } - - None } + + None } } diff --git a/src/query/service/src/pipelines/processors/transforms/group_by/keys_ref.rs b/src/query/service/src/pipelines/processors/transforms/group_by/keys_ref.rs deleted file mode 100644 index 0b182d77a4153..0000000000000 --- a/src/query/service/src/pipelines/processors/transforms/group_by/keys_ref.rs +++ /dev/null @@ -1,70 +0,0 @@ -// Copyright 2021 Datafuse Labs. -// -// Licensed 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 std::hash::Hasher; - -use ahash::AHasher; -use common_hashtable::HashTableKeyable; - -#[derive(Clone, Copy)] -pub struct KeysRef { - pub length: usize, - pub address: usize, -} - -impl KeysRef { - pub fn create(address: usize, length: usize) -> KeysRef { - KeysRef { length, address } - } -} - -impl Eq for KeysRef {} - -impl PartialEq for KeysRef { - fn eq(&self, other: &Self) -> bool { - if self.length != other.length { - return false; - } - - unsafe { - let self_value = std::slice::from_raw_parts(self.address as *const u8, self.length); - let other_value = std::slice::from_raw_parts(other.address as *const u8, other.length); - self_value == other_value - } - } -} - -impl HashTableKeyable for KeysRef { - const BEFORE_EQ_HASH: bool = true; - - fn is_zero(&self) -> bool { - self.length == 0 - } - - fn fast_hash(&self) -> u64 { - unsafe { - // TODO(Winter) We need more efficient hash algorithm - let value = std::slice::from_raw_parts(self.address as *const u8, self.length); - - let mut hasher = AHasher::default(); - hasher.write(value); - hasher.finish() - } - } - - fn set_key(&mut self, new_value: &Self) { - self.length = new_value.length; - self.address = new_value.address; - } -} diff --git a/src/query/service/src/pipelines/processors/transforms/group_by/mod.rs b/src/query/service/src/pipelines/processors/transforms/group_by/mod.rs index 06c8f70db89b7..81b1142078401 100644 --- a/src/query/service/src/pipelines/processors/transforms/group_by/mod.rs +++ b/src/query/service/src/pipelines/processors/transforms/group_by/mod.rs @@ -19,11 +19,11 @@ mod aggregator_polymorphic_keys; mod aggregator_state; mod aggregator_state_entity; mod aggregator_state_iterator; -pub(crate) mod keys_ref; pub use aggregator_groups_builder::GroupColumnsBuilder; pub use aggregator_keys_builder::KeysColumnBuilder; pub use aggregator_keys_iter::KeysColumnIter; pub use aggregator_polymorphic_keys::PolymorphicKeysHelper; pub use aggregator_state::AggregatorState; -pub use aggregator_state_entity::StateEntity; +pub use aggregator_state_entity::StateEntityMutRef; +pub use aggregator_state_entity::StateEntityRef; diff --git a/src/query/service/src/pipelines/processors/transforms/hash_join/common.rs b/src/query/service/src/pipelines/processors/transforms/hash_join/common.rs index 284597114b390..04b9e7bbed798 100644 --- a/src/query/service/src/pipelines/processors/transforms/hash_join/common.rs +++ b/src/query/service/src/pipelines/processors/transforms/hash_join/common.rs @@ -32,9 +32,7 @@ use common_datavalues::NullableType; use common_datavalues::Series; use common_exception::ErrorCode; use common_exception::Result; -use common_hashtable::HashMap; -use common_hashtable::HashTableKeyable; -use common_hashtable::KeyValueEntity; +use common_hashtable::HashtableLike; use crate::pipelines::processors::transforms::hash_join::desc::MarkerKind; use crate::pipelines::processors::transforms::hash_join::row::RowPtr; @@ -62,15 +60,15 @@ impl JoinHashTable { } #[inline] - pub(crate) fn probe_key( + pub(crate) fn probe_key<'a, H: HashtableLike>>( &self, - hash_table: &HashMap>, - key: Key, + hash_table: &'a H, + key: H::KeyRef<'_>, valids: &Option, i: usize, - ) -> Option<*mut KeyValueEntity>> { + ) -> Option> { if valids.as_ref().map_or(true, |v| v.get_bit(i)) { - return hash_table.find_key(&key); + return hash_table.entry(key); } None } diff --git a/src/query/service/src/pipelines/processors/transforms/hash_join/hash_join_state_impl.rs b/src/query/service/src/pipelines/processors/transforms/hash_join/hash_join_state_impl.rs index 5a933692940ad..53fd6986eab65 100644 --- a/src/query/service/src/pipelines/processors/transforms/hash_join/hash_join_state_impl.rs +++ b/src/query/service/src/pipelines/processors/transforms/hash_join/hash_join_state_impl.rs @@ -26,7 +26,6 @@ use common_exception::ErrorCode; use common_exception::Result; use super::ProbeState; -use crate::pipelines::processors::transforms::group_by::keys_ref::KeysRef; use crate::pipelines::processors::transforms::hash_join::row::RowPtr; use crate::pipelines::processors::HashJoinState; use crate::pipelines::processors::HashTable; @@ -99,7 +98,6 @@ impl HashJoinState for JoinHashTable { let build_keys_iter = $method.build_keys_iter(&keys_state)?; for (row_index, key) in build_keys_iter.enumerate().take($chunk.num_rows()) { - let mut inserted = true; let ptr = RowPtr { chunk_index: $chunk_index, row_index, @@ -109,11 +107,13 @@ impl HashJoinState for JoinHashTable { let mut self_row_ptrs = self.row_ptrs.write(); self_row_ptrs.push(ptr); } - let entity = $table.insert_key(&key, &mut inserted); - if inserted { - entity.set_value(vec![ptr]); - } else { - entity.get_mut_value().push(ptr); + match unsafe { $table.insert(key) } { + Ok(entity) => { + entity.write(vec![ptr]); + } + Err(entity) => { + entity.push(ptr); + } } } }}; @@ -170,7 +170,6 @@ impl HashJoinState for JoinHashTable { .hash_method .build_keys_iter(chunk.keys_state.as_ref().unwrap())?; for (row_index, key) in build_keys_iter.enumerate().take(chunk.num_rows()) { - let mut inserted = true; let ptr = RowPtr { chunk_index, row_index, @@ -180,12 +179,13 @@ impl HashJoinState for JoinHashTable { let mut self_row_ptrs = self.row_ptrs.write(); self_row_ptrs.push(ptr); } - let keys_ref = KeysRef::create(key.as_ptr() as usize, key.len()); - let entity = table.hash_table.insert_key(&keys_ref, &mut inserted); - if inserted { - entity.set_value(vec![ptr]); - } else { - entity.get_mut_value().push(ptr); + match unsafe { table.hash_table.insert_borrowing(key) } { + Ok(entity) => { + entity.write(vec![ptr]); + } + Err(entity) => { + entity.push(ptr); + } } } } diff --git a/src/query/service/src/pipelines/processors/transforms/hash_join/join_hash_table.rs b/src/query/service/src/pipelines/processors/transforms/hash_join/join_hash_table.rs index 8930991b95982..8ec6a05c3570a 100644 --- a/src/query/service/src/pipelines/processors/transforms/hash_join/join_hash_table.rs +++ b/src/query/service/src/pipelines/processors/transforms/hash_join/join_hash_table.rs @@ -29,12 +29,13 @@ use common_datavalues::DataSchemaRef; use common_datavalues::DataTypeImpl; use common_exception::Result; use common_hashtable::HashMap; +use common_hashtable::HashtableKeyable; +use common_hashtable::UnsizedHashMap; use parking_lot::RwLock; use primitive_types::U256; use primitive_types::U512; use super::ProbeState; -use crate::pipelines::processors::transforms::group_by::keys_ref::KeysRef; use crate::pipelines::processors::transforms::hash_join::desc::HashJoinDesc; use crate::pipelines::processors::transforms::hash_join::row::RowPtr; use crate::pipelines::processors::transforms::hash_join::row::RowSpace; @@ -46,54 +47,24 @@ use crate::sql::executor::PhysicalScalar; use crate::sql::planner::plans::JoinType; pub struct SerializerHashTable { - pub(crate) hash_table: HashMap>, + pub(crate) hash_table: UnsizedHashMap<[u8], Vec>, pub(crate) hash_method: HashMethodSerializer, } -pub struct KeyU8HashTable { - pub(crate) hash_table: HashMap>, - pub(crate) hash_method: HashMethodFixedKeys, -} - -pub struct KeyU16HashTable { - pub(crate) hash_table: HashMap>, - pub(crate) hash_method: HashMethodFixedKeys, -} - -pub struct KeyU32HashTable { - pub(crate) hash_table: HashMap>, - pub(crate) hash_method: HashMethodFixedKeys, -} - -pub struct KeyU64HashTable { - pub(crate) hash_table: HashMap>, - pub(crate) hash_method: HashMethodFixedKeys, -} - -pub struct KeyU128HashTable { - pub(crate) hash_table: HashMap>, - pub(crate) hash_method: HashMethodFixedKeys, -} - -pub struct KeyU256HashTable { - pub(crate) hash_table: HashMap>, - pub(crate) hash_method: HashMethodFixedKeys, -} - -pub struct KeyU512HashTable { - pub(crate) hash_table: HashMap>, - pub(crate) hash_method: HashMethodFixedKeys, +pub struct FixedKeyHashTable { + pub(crate) hash_table: HashMap>, + pub(crate) hash_method: HashMethodFixedKeys, } pub enum HashTable { SerializerHashTable(SerializerHashTable), - KeyU8HashTable(KeyU8HashTable), - KeyU16HashTable(KeyU16HashTable), - KeyU32HashTable(KeyU32HashTable), - KeyU64HashTable(KeyU64HashTable), - KeyU128HashTable(KeyU128HashTable), - KeyU256HashTable(KeyU256HashTable), - KeyU512HashTable(KeyU512HashTable), + KeyU8HashTable(FixedKeyHashTable), + KeyU16HashTable(FixedKeyHashTable), + KeyU32HashTable(FixedKeyHashTable), + KeyU64HashTable(FixedKeyHashTable), + KeyU128HashTable(FixedKeyHashTable), + KeyU256HashTable(FixedKeyHashTable), + KeyU512HashTable(FixedKeyHashTable), } pub struct JoinHashTable { @@ -126,7 +97,7 @@ impl JoinHashTable { HashMethodKind::Serializer(_) => Arc::new(JoinHashTable::try_create( ctx, HashTable::SerializerHashTable(SerializerHashTable { - hash_table: HashMap::>::create(), + hash_table: UnsizedHashMap::<[u8], Vec>::new(), hash_method: HashMethodSerializer::default(), }), build_schema, @@ -135,8 +106,8 @@ impl JoinHashTable { )?), HashMethodKind::KeysU8(hash_method) => Arc::new(JoinHashTable::try_create( ctx, - HashTable::KeyU8HashTable(KeyU8HashTable { - hash_table: HashMap::>::create(), + HashTable::KeyU8HashTable(FixedKeyHashTable { + hash_table: HashMap::>::new(), hash_method, }), build_schema, @@ -145,8 +116,8 @@ impl JoinHashTable { )?), HashMethodKind::KeysU16(hash_method) => Arc::new(JoinHashTable::try_create( ctx, - HashTable::KeyU16HashTable(KeyU16HashTable { - hash_table: HashMap::>::create(), + HashTable::KeyU16HashTable(FixedKeyHashTable { + hash_table: HashMap::>::new(), hash_method, }), build_schema, @@ -155,8 +126,8 @@ impl JoinHashTable { )?), HashMethodKind::KeysU32(hash_method) => Arc::new(JoinHashTable::try_create( ctx, - HashTable::KeyU32HashTable(KeyU32HashTable { - hash_table: HashMap::>::create(), + HashTable::KeyU32HashTable(FixedKeyHashTable { + hash_table: HashMap::>::new(), hash_method, }), build_schema, @@ -165,8 +136,8 @@ impl JoinHashTable { )?), HashMethodKind::KeysU64(hash_method) => Arc::new(JoinHashTable::try_create( ctx, - HashTable::KeyU64HashTable(KeyU64HashTable { - hash_table: HashMap::>::create(), + HashTable::KeyU64HashTable(FixedKeyHashTable { + hash_table: HashMap::>::new(), hash_method, }), build_schema, @@ -175,8 +146,8 @@ impl JoinHashTable { )?), HashMethodKind::KeysU128(hash_method) => Arc::new(JoinHashTable::try_create( ctx, - HashTable::KeyU128HashTable(KeyU128HashTable { - hash_table: HashMap::>::create(), + HashTable::KeyU128HashTable(FixedKeyHashTable { + hash_table: HashMap::>::new(), hash_method, }), build_schema, @@ -185,8 +156,8 @@ impl JoinHashTable { )?), HashMethodKind::KeysU256(hash_method) => Arc::new(JoinHashTable::try_create( ctx, - HashTable::KeyU256HashTable(KeyU256HashTable { - hash_table: HashMap::>::create(), + HashTable::KeyU256HashTable(FixedKeyHashTable { + hash_table: HashMap::>::new(), hash_method, }), build_schema, @@ -195,8 +166,8 @@ impl JoinHashTable { )?), HashMethodKind::KeysU512(hash_method) => Arc::new(JoinHashTable::try_create( ctx, - HashTable::KeyU512HashTable(KeyU512HashTable { - hash_table: HashMap::>::create(), + HashTable::KeyU512HashTable(FixedKeyHashTable { + hash_table: HashMap::>::new(), hash_method, }), build_schema, @@ -279,10 +250,8 @@ impl JoinHashTable { .hash_method .build_keys_state(&probe_keys, input.num_rows())?; let keys_iter = table.hash_method.build_keys_iter(&keys_state)?; - let keys_ref = - keys_iter.map(|key| KeysRef::create(key.as_ptr() as usize, key.len())); - self.result_blocks(&table.hash_table, probe_state, keys_ref, input) + self.result_blocks(&table.hash_table, probe_state, keys_iter, input) } HashTable::KeyU8HashTable(table) => { let keys_state = table diff --git a/src/query/service/src/pipelines/processors/transforms/hash_join/mod.rs b/src/query/service/src/pipelines/processors/transforms/hash_join/mod.rs index 9b055b9c63540..19a50895d8dc4 100644 --- a/src/query/service/src/pipelines/processors/transforms/hash_join/mod.rs +++ b/src/query/service/src/pipelines/processors/transforms/hash_join/mod.rs @@ -25,15 +25,9 @@ mod util; pub use desc::HashJoinDesc; pub use hash_join_state::HashJoinState; +pub use join_hash_table::FixedKeyHashTable; pub use join_hash_table::HashTable; pub use join_hash_table::JoinHashTable; -pub use join_hash_table::KeyU128HashTable; -pub use join_hash_table::KeyU16HashTable; -pub use join_hash_table::KeyU256HashTable; -pub use join_hash_table::KeyU32HashTable; -pub use join_hash_table::KeyU512HashTable; -pub use join_hash_table::KeyU64HashTable; -pub use join_hash_table::KeyU8HashTable; pub use join_hash_table::SerializerHashTable; pub use probe_state::ProbeState; pub use result_blocks::*; diff --git a/src/query/service/src/pipelines/processors/transforms/hash_join/probe_join/inner_join.rs b/src/query/service/src/pipelines/processors/transforms/hash_join/probe_join/inner_join.rs index 12b04d9a9b601..a055f5d9618b3 100644 --- a/src/query/service/src/pipelines/processors/transforms/hash_join/probe_join/inner_join.rs +++ b/src/query/service/src/pipelines/processors/transforms/hash_join/probe_join/inner_join.rs @@ -20,24 +20,24 @@ use common_catalog::table_context::TableContext; use common_datablocks::DataBlock; use common_exception::ErrorCode; use common_exception::Result; -use common_hashtable::HashMap; -use common_hashtable::HashTableKeyable; +use common_hashtable::HashtableEntryRefLike; +use common_hashtable::HashtableLike; use crate::pipelines::processors::transforms::hash_join::row::RowPtr; use crate::pipelines::processors::transforms::hash_join::ProbeState; use crate::pipelines::processors::JoinHashTable; impl JoinHashTable { - pub(crate) fn probe_inner_join( + pub(crate) fn probe_inner_join<'a, H: HashtableLike>, IT>( &self, - hash_table: &HashMap>, + hash_table: &H, probe_state: &mut ProbeState, keys_iter: IT, input: &DataBlock, ) -> Result> where - Key: HashTableKeyable + Clone + 'static, - IT: Iterator + TrustedLen, + IT: Iterator> + TrustedLen, + H::Key: 'a, { let valids = &probe_state.valids; let block_size = self @@ -54,13 +54,13 @@ impl JoinHashTable { for (i, key) in keys_iter.enumerate() { // If the join is derived from correlated subquery, then null equality is safe. let probe_result_ptr = if self.hash_join_desc.from_correlated_subquery { - hash_table.find_key(&key) + hash_table.entry(key) } else { self.probe_key(hash_table, key, valids, i) }; if let Some(v) = probe_result_ptr { - let probed_rows = v.get_value(); + let probed_rows = v.get(); if probe_indexes.len() + probed_rows.len() < probe_indexes.capacity() { build_indexes.extend_from_slice(probed_rows); diff --git a/src/query/service/src/pipelines/processors/transforms/hash_join/probe_join/left_join.rs b/src/query/service/src/pipelines/processors/transforms/hash_join/probe_join/left_join.rs index 5187c39acf6d9..1e857fc937598 100644 --- a/src/query/service/src/pipelines/processors/transforms/hash_join/probe_join/left_join.rs +++ b/src/query/service/src/pipelines/processors/transforms/hash_join/probe_join/left_join.rs @@ -28,8 +28,8 @@ use common_datavalues::DataType; use common_datavalues::DataValue; use common_exception::ErrorCode; use common_exception::Result; -use common_hashtable::HashMap; -use common_hashtable::HashTableKeyable; +use common_hashtable::HashtableEntryRefLike; +use common_hashtable::HashtableLike; use crate::pipelines::processors::transforms::hash_join::desc::MarkerKind; use crate::pipelines::processors::transforms::hash_join::row::RowPtr; @@ -38,16 +38,21 @@ use crate::pipelines::processors::JoinHashTable; use crate::sql::plans::JoinType; impl JoinHashTable { - pub(crate) fn probe_left_join( + pub(crate) fn probe_left_join< + 'a, + const WITH_OTHER_CONJUNCT: bool, + H: HashtableLike>, + IT, + >( &self, - hash_table: &HashMap>, + hash_table: &H, probe_state: &mut ProbeState, keys_iter: IT, input: &DataBlock, ) -> Result> where - Key: HashTableKeyable + Clone + 'static, - IT: Iterator + TrustedLen, + IT: Iterator> + TrustedLen, + H::Key: 'a, { let valids = &probe_state.valids; let block_size = self.ctx.get_settings().get_max_block_size()? as usize; @@ -74,13 +79,13 @@ impl JoinHashTable { // Start to probe hash table for (i, key) in keys_iter.enumerate() { let probe_result_ptr = match self.hash_join_desc.from_correlated_subquery { - true => hash_table.find_key(&key), + true => hash_table.entry(key), false => self.probe_key(hash_table, key, valids, i), }; let (validity_value, probed_rows) = match probe_result_ptr { None => (false, &dummy_probed_rows), - Some(v) => (true, v.get_value()), + Some(v) => (true, v.get()), }; if self.hash_join_desc.join_type == JoinType::Full { diff --git a/src/query/service/src/pipelines/processors/transforms/hash_join/probe_join/left_mark.rs b/src/query/service/src/pipelines/processors/transforms/hash_join/probe_join/left_mark.rs index 3fb76c3706721..e3a673d56e2a7 100644 --- a/src/query/service/src/pipelines/processors/transforms/hash_join/probe_join/left_mark.rs +++ b/src/query/service/src/pipelines/processors/transforms/hash_join/probe_join/left_mark.rs @@ -22,8 +22,8 @@ use common_datavalues::BooleanViewer; use common_datavalues::ScalarViewer; use common_exception::ErrorCode; use common_exception::Result; -use common_hashtable::HashMap; -use common_hashtable::HashTableKeyable; +use common_hashtable::HashtableEntryRefLike; +use common_hashtable::HashtableLike; use crate::pipelines::processors::transforms::hash_join::desc::MarkerKind; use crate::pipelines::processors::transforms::hash_join::row::RowPtr; @@ -31,16 +31,16 @@ use crate::pipelines::processors::transforms::hash_join::ProbeState; use crate::pipelines::processors::JoinHashTable; impl JoinHashTable { - pub(crate) fn probe_left_mark_join( + pub(crate) fn probe_left_mark_join<'a, H: HashtableLike>, IT>( &self, - hash_table: &HashMap>, + hash_table: &H, probe_state: &mut ProbeState, keys_iter: IT, input: &DataBlock, ) -> Result> where - Key: HashTableKeyable + Clone + 'static, - IT: Iterator + TrustedLen, + IT: Iterator> + TrustedLen, + H::Key: 'a, { let valids = &probe_state.valids; let mut block_size = self.ctx.get_settings().get_max_block_size()? as usize; @@ -69,12 +69,12 @@ impl JoinHashTable { } let probe_result_ptr = match self.hash_join_desc.from_correlated_subquery { - true => hash_table.find_key(&key), + true => hash_table.entry(key), false => self.probe_key(hash_table, key, valids, i), }; if let Some(v) = probe_result_ptr { - let probed_rows = v.get_value(); + let probed_rows = v.get(); for probed_row in probed_rows { if let Some(p) = self_row_ptrs.iter_mut().find(|p| (*p).eq(&probed_row)) { @@ -86,16 +86,20 @@ impl JoinHashTable { Ok(vec![DataBlock::empty()]) } - pub(crate) fn probe_left_mark_join_with_conjunct( + pub(crate) fn probe_left_mark_join_with_conjunct< + 'a, + H: HashtableLike>, + IT, + >( &self, - hash_table: &HashMap>, + hash_table: &H, probe_state: &mut ProbeState, keys_iter: IT, input: &DataBlock, ) -> Result> where - Key: HashTableKeyable + Clone + 'static, - IT: Iterator + TrustedLen, + IT: Iterator> + TrustedLen, + H::Key: 'a, { let valids = &probe_state.valids; let block_size = self.ctx.get_settings().get_max_block_size()? as usize; @@ -118,12 +122,12 @@ impl JoinHashTable { for (i, key) in keys_iter.enumerate() { let probe_result_ptr = match self.hash_join_desc.from_correlated_subquery { - true => hash_table.find_key(&key), + true => hash_table.entry(key), false => self.probe_key(hash_table, key, valids, i), }; if let Some(v) = probe_result_ptr { - let probed_rows = v.get_value(); + let probed_rows = v.get(); if probe_indexes.len() + probed_rows.len() < probe_indexes.capacity() { build_indexes.extend_from_slice(probed_rows); diff --git a/src/query/service/src/pipelines/processors/transforms/hash_join/probe_join/left_semi_join.rs b/src/query/service/src/pipelines/processors/transforms/hash_join/probe_join/left_semi_join.rs index 332fa4956ae89..bbecb7a3e2605 100644 --- a/src/query/service/src/pipelines/processors/transforms/hash_join/probe_join/left_semi_join.rs +++ b/src/query/service/src/pipelines/processors/transforms/hash_join/probe_join/left_semi_join.rs @@ -23,8 +23,8 @@ use common_datavalues::BooleanColumn; use common_datavalues::Column; use common_exception::ErrorCode; use common_exception::Result; -use common_hashtable::HashMap; -use common_hashtable::HashTableKeyable; +use common_hashtable::HashtableEntryRefLike; +use common_hashtable::HashtableLike; use crate::pipelines::processors::transforms::hash_join::row::RowPtr; use crate::pipelines::processors::transforms::hash_join::ProbeState; @@ -32,16 +32,16 @@ use crate::pipelines::processors::JoinHashTable; /// Semi join contain semi join and semi-anti join impl JoinHashTable { - pub(crate) fn probe_left_semi_join( + pub(crate) fn probe_left_semi_join<'a, H: HashtableLike>, IT>( &self, - hash_table: &HashMap>, + hash_table: &H, probe_state: &mut ProbeState, keys_iter: IT, input: &DataBlock, ) -> Result> where - Key: HashTableKeyable + Clone + 'static, - IT: Iterator + TrustedLen, + IT: Iterator> + TrustedLen, + H::Key: 'a, { match self.hash_join_desc.other_predicate.is_none() { true => { @@ -56,16 +56,16 @@ impl JoinHashTable { } } - pub(crate) fn probe_left_anti_semi_join( + pub(crate) fn probe_left_anti_semi_join<'a, H: HashtableLike>, IT>( &self, - hash_table: &HashMap>, + hash_table: &H, probe_state: &mut ProbeState, keys_iter: IT, input: &DataBlock, ) -> Result> where - Key: HashTableKeyable + Clone + 'static, - IT: Iterator + TrustedLen, + IT: Iterator> + TrustedLen, + H::Key: 'a, { match self.hash_join_desc.other_predicate.is_none() { true => { @@ -80,23 +80,23 @@ impl JoinHashTable { } } - fn left_semi_anti_join( + fn left_semi_anti_join<'a, const SEMI: bool, H: HashtableLike>, IT>( &self, - hash_table: &HashMap>, + hash_table: &H, probe_state: &mut ProbeState, keys_iter: IT, input: &DataBlock, ) -> Result> where - Key: HashTableKeyable + Clone + 'static, - IT: Iterator + TrustedLen, + IT: Iterator> + TrustedLen, + H::Key: 'a, { let valids = &probe_state.valids; let mut probe_indexes = Vec::with_capacity(keys_iter.size_hint().0); for (i, key) in keys_iter.enumerate() { let probe_result_ptr = if self.hash_join_desc.from_correlated_subquery { - hash_table.find_key(&key) + hash_table.entry(key) } else { self.probe_key(hash_table, key, valids, i) }; @@ -114,16 +114,21 @@ impl JoinHashTable { )?]) } - fn left_semi_anti_join_with_other_conjunct( + fn left_semi_anti_join_with_other_conjunct< + 'a, + const SEMI: bool, + H: HashtableLike>, + IT, + >( &self, - hash_table: &HashMap>, + hash_table: &H, probe_state: &mut ProbeState, keys_iter: IT, input: &DataBlock, ) -> Result> where - Key: HashTableKeyable + Clone + 'static, - IT: Iterator + TrustedLen, + IT: Iterator> + TrustedLen, + H::Key: 'a, { let valids = &probe_state.valids; let block_size = self.ctx.get_settings().get_max_block_size()? as usize; @@ -144,7 +149,7 @@ impl JoinHashTable { for (i, key) in keys_iter.enumerate() { let probe_result_ptr = match self.hash_join_desc.from_correlated_subquery { - true => hash_table.find_key(&key), + true => hash_table.entry(key), false => self.probe_key(hash_table, key, valids, i), }; @@ -153,7 +158,7 @@ impl JoinHashTable { continue; } None => &dummy_probed_rows, - Some(v) => v.get_value(), + Some(v) => v.get(), }; if probe_result_ptr.is_some() && !SEMI { diff --git a/src/query/service/src/pipelines/processors/transforms/hash_join/probe_join/right_join.rs b/src/query/service/src/pipelines/processors/transforms/hash_join/probe_join/right_join.rs index 5d01d969cbb85..3eb93d0d70747 100644 --- a/src/query/service/src/pipelines/processors/transforms/hash_join/probe_join/right_join.rs +++ b/src/query/service/src/pipelines/processors/transforms/hash_join/probe_join/right_join.rs @@ -21,8 +21,8 @@ use common_catalog::table_context::TableContext; use common_datablocks::DataBlock; use common_exception::ErrorCode; use common_exception::Result; -use common_hashtable::HashMap; -use common_hashtable::HashTableKeyable; +use common_hashtable::HashtableEntryRefLike; +use common_hashtable::HashtableLike; use crate::pipelines::processors::transforms::hash_join::row::RowPtr; use crate::pipelines::processors::transforms::hash_join::ProbeState; @@ -31,16 +31,16 @@ use crate::sql::plans::JoinType; impl JoinHashTable { /// Used by right join/right semi(anti) join - pub(crate) fn probe_right_join( + pub(crate) fn probe_right_join<'a, H: HashtableLike>, IT>( &self, - hash_table: &HashMap>, + hash_table: &H, probe_state: &mut ProbeState, keys_iter: IT, input: &DataBlock, ) -> Result> where - Key: HashTableKeyable + Clone + 'static, - IT: Iterator + TrustedLen, + IT: Iterator> + TrustedLen, + H::Key: 'a, { let valids = &probe_state.valids; let block_size = self.ctx.get_settings().get_max_block_size()? as usize; @@ -57,7 +57,7 @@ impl JoinHashTable { let probe_result_ptr = self.probe_key(hash_table, key, valids, i); if let Some(v) = probe_result_ptr { - let probed_rows = v.get_value(); + let probed_rows = v.get(); if probe_indexes.len() + probed_rows.len() < probe_indexes.capacity() { build_indexes.extend(probed_rows); diff --git a/src/query/service/src/pipelines/processors/transforms/hash_join/probe_join/right_mark.rs b/src/query/service/src/pipelines/processors/transforms/hash_join/probe_join/right_mark.rs index ec297bd378713..94394a64cd294 100644 --- a/src/query/service/src/pipelines/processors/transforms/hash_join/probe_join/right_mark.rs +++ b/src/query/service/src/pipelines/processors/transforms/hash_join/probe_join/right_mark.rs @@ -22,8 +22,8 @@ use common_datavalues::BooleanViewer; use common_datavalues::ScalarViewer; use common_exception::ErrorCode; use common_exception::Result; -use common_hashtable::HashMap; -use common_hashtable::HashTableKeyable; +use common_hashtable::HashtableEntryRefLike; +use common_hashtable::HashtableLike; use crate::pipelines::processors::transforms::hash_join::desc::MarkerKind; use crate::pipelines::processors::transforms::hash_join::row::RowPtr; @@ -31,23 +31,23 @@ use crate::pipelines::processors::transforms::hash_join::ProbeState; use crate::pipelines::processors::JoinHashTable; impl JoinHashTable { - pub(crate) fn probe_right_mark_join( + pub(crate) fn probe_right_mark_join<'a, H: HashtableLike>, IT>( &self, - hash_table: &HashMap>, + hash_table: &H, probe_state: &mut ProbeState, keys_iter: IT, input: &DataBlock, ) -> Result> where - Key: HashTableKeyable + Clone + 'static, - IT: Iterator + TrustedLen, + IT: Iterator> + TrustedLen, + H::Key: 'a, { let valids = &probe_state.valids; let has_null = *self.hash_join_desc.marker_join_desc.has_null.read(); let markers = probe_state.markers.as_mut().unwrap(); for (i, key) in keys_iter.enumerate() { let probe_result_ptr = match self.hash_join_desc.from_correlated_subquery { - true => hash_table.find_key(&key), + true => hash_table.entry(key), false => self.probe_key(hash_table, key, valids, i), }; @@ -62,16 +62,20 @@ impl JoinHashTable { )?]) } - pub(crate) fn probe_right_mark_join_with_conjunct( + pub(crate) fn probe_right_mark_join_with_conjunct< + 'a, + H: HashtableLike>, + IT, + >( &self, - hash_table: &HashMap>, + hash_table: &H, probe_state: &mut ProbeState, keys_iter: IT, input: &DataBlock, ) -> Result> where - Key: HashTableKeyable + Clone + 'static, - IT: Iterator + TrustedLen, + IT: Iterator> + TrustedLen, + H::Key: 'a, { let valids = &probe_state.valids; let block_size = self.ctx.get_settings().get_max_block_size()? as usize; @@ -87,12 +91,12 @@ impl JoinHashTable { for (i, key) in keys_iter.enumerate() { let probe_result_ptr = match self.hash_join_desc.from_correlated_subquery { - true => hash_table.find_key(&key), + true => hash_table.entry(key), false => self.probe_key(hash_table, key, valids, i), }; if let Some(v) = probe_result_ptr { - let probed_rows = v.get_value(); + let probed_rows = v.get(); if probe_indexes.len() + probed_rows.len() < probe_indexes.capacity() { build_indexes.extend_from_slice(probed_rows); diff --git a/src/query/service/src/pipelines/processors/transforms/hash_join/result_blocks.rs b/src/query/service/src/pipelines/processors/transforms/hash_join/result_blocks.rs index 5c33ca705a698..2cbaa6509826a 100644 --- a/src/query/service/src/pipelines/processors/transforms/hash_join/result_blocks.rs +++ b/src/query/service/src/pipelines/processors/transforms/hash_join/result_blocks.rs @@ -17,8 +17,7 @@ use std::iter::TrustedLen; use common_datablocks::DataBlock; use common_exception::ErrorCode; use common_exception::Result; -use common_hashtable::HashMap; -use common_hashtable::HashTableKeyable; +use common_hashtable::HashtableLike; use super::JoinHashTable; use super::ProbeState; @@ -26,16 +25,16 @@ use crate::pipelines::processors::transforms::hash_join::row::RowPtr; use crate::sql::planner::plans::JoinType; impl JoinHashTable { - pub(crate) fn result_blocks( + pub(crate) fn result_blocks<'a, H: HashtableLike>, IT>( &self, - hash_table: &HashMap>, + hash_table: &H, probe_state: &mut ProbeState, keys_iter: IT, input: &DataBlock, ) -> Result> where - Key: HashTableKeyable + Clone + 'static, - IT: Iterator + TrustedLen, + IT: Iterator> + TrustedLen, + H::Key: 'a, { match self.hash_join_desc.join_type { JoinType::Inner => self.probe_inner_join(hash_table, probe_state, keys_iter, input), diff --git a/src/query/service/src/pipelines/processors/transforms/mod.rs b/src/query/service/src/pipelines/processors/transforms/mod.rs index f881ff53a6655..5ef252b6d8502 100644 --- a/src/query/service/src/pipelines/processors/transforms/mod.rs +++ b/src/query/service/src/pipelines/processors/transforms/mod.rs @@ -40,17 +40,11 @@ use common_pipeline_transforms::processors::transforms::transform_expression; use common_pipeline_transforms::processors::transforms::transform_sort_merge; use common_pipeline_transforms::processors::transforms::transform_sort_partial; pub use common_pipeline_transforms::processors::ExpressionExecutor; +pub use hash_join::FixedKeyHashTable; pub use hash_join::HashJoinDesc; pub use hash_join::HashJoinState; pub use hash_join::HashTable; pub use hash_join::JoinHashTable; -pub use hash_join::KeyU128HashTable; -pub use hash_join::KeyU16HashTable; -pub use hash_join::KeyU256HashTable; -pub use hash_join::KeyU32HashTable; -pub use hash_join::KeyU512HashTable; -pub use hash_join::KeyU64HashTable; -pub use hash_join::KeyU8HashTable; pub use hash_join::SerializerHashTable; pub use transform_addon::TransformAddOn; pub use transform_aggregator::TransformAggregator;