|
| 1 | +// Copyright 2021 Datafuse Labs |
| 2 | +// |
| 3 | +// Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +// you may not use this file except in compliance with the License. |
| 5 | +// You may obtain a copy of the License at |
| 6 | +// |
| 7 | +// http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +// |
| 9 | +// Unless required by applicable law or agreed to in writing, software |
| 10 | +// distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +// See the License for the specific language governing permissions and |
| 13 | +// limitations under the License. |
| 14 | + |
| 15 | +use std::hash::Hash; |
| 16 | +use std::marker::PhantomData; |
| 17 | + |
| 18 | +use databend_common_expression::types::boolean::TrueIdxIter; |
| 19 | +use databend_common_expression::types::DataType; |
| 20 | +use databend_common_expression::types::DateType; |
| 21 | +use databend_common_expression::types::Decimal128Type; |
| 22 | +use databend_common_expression::types::Decimal256Type; |
| 23 | +use databend_common_expression::types::Decimal64Type; |
| 24 | +use databend_common_expression::types::NumberDataType; |
| 25 | +use databend_common_expression::types::NumberType; |
| 26 | +use databend_common_expression::types::StringType; |
| 27 | +use databend_common_expression::types::TimestampType; |
| 28 | +use databend_common_expression::types::ValueType; |
| 29 | +use databend_common_expression::with_number_mapped_type; |
| 30 | +use databend_common_expression::Column; |
| 31 | +use databend_common_expression::ScalarRef; |
| 32 | +use databend_common_expression::SELECTIVITY_THRESHOLD; |
| 33 | +use databend_storages_common_table_meta::meta::ColumnDistinctHLL; |
| 34 | + |
| 35 | +pub trait ColumnNDVEstimator: Send + Sync { |
| 36 | + fn update_column(&mut self, column: &Column); |
| 37 | + fn update_scalar(&mut self, scalar: &ScalarRef); |
| 38 | + fn finalize(&self) -> u64; |
| 39 | +} |
| 40 | + |
| 41 | +pub fn create_column_ndv_estimator(data_type: &DataType) -> Box<dyn ColumnNDVEstimator> { |
| 42 | + let inner_type = data_type.remove_nullable(); |
| 43 | + with_number_mapped_type!(|NUM_TYPE| match inner_type { |
| 44 | + DataType::Number(NumberDataType::NUM_TYPE) => { |
| 45 | + ColumnNDVEstimatorImpl::<NumberType<NUM_TYPE>>::create() |
| 46 | + } |
| 47 | + DataType::String => { |
| 48 | + ColumnNDVEstimatorImpl::<StringType>::create() |
| 49 | + } |
| 50 | + DataType::Date => { |
| 51 | + ColumnNDVEstimatorImpl::<DateType>::create() |
| 52 | + } |
| 53 | + DataType::Timestamp => { |
| 54 | + ColumnNDVEstimatorImpl::<TimestampType>::create() |
| 55 | + } |
| 56 | + DataType::Decimal(size) => { |
| 57 | + if size.can_carried_by_64() { |
| 58 | + ColumnNDVEstimatorImpl::<Decimal64Type>::create() |
| 59 | + } else if size.can_carried_by_128() { |
| 60 | + ColumnNDVEstimatorImpl::<Decimal128Type>::create() |
| 61 | + } else { |
| 62 | + ColumnNDVEstimatorImpl::<Decimal256Type>::create() |
| 63 | + } |
| 64 | + } |
| 65 | + _ => unreachable!("Unsupported data type: {:?}", data_type), |
| 66 | + }) |
| 67 | +} |
| 68 | + |
| 69 | +pub struct ColumnNDVEstimatorImpl<T> |
| 70 | +where |
| 71 | + T: ValueType + Send + Sync, |
| 72 | + for<'a> T::ScalarRef<'a>: Hash, |
| 73 | +{ |
| 74 | + hll: ColumnDistinctHLL, |
| 75 | + _phantom: PhantomData<T>, |
| 76 | +} |
| 77 | + |
| 78 | +impl<T> ColumnNDVEstimatorImpl<T> |
| 79 | +where |
| 80 | + T: ValueType + Send + Sync, |
| 81 | + for<'a> T::ScalarRef<'a>: Hash, |
| 82 | +{ |
| 83 | + pub fn create() -> Box<dyn ColumnNDVEstimator> { |
| 84 | + Box::new(Self { |
| 85 | + hll: ColumnDistinctHLL::new(), |
| 86 | + _phantom: Default::default(), |
| 87 | + }) |
| 88 | + } |
| 89 | +} |
| 90 | + |
| 91 | +impl<T> ColumnNDVEstimator for ColumnNDVEstimatorImpl<T> |
| 92 | +where |
| 93 | + T: ValueType + Send + Sync, |
| 94 | + for<'a> T::ScalarRef<'a>: Hash, |
| 95 | +{ |
| 96 | + fn update_column(&mut self, column: &Column) { |
| 97 | + let (column, validity) = if let Column::Nullable(box inner) = column { |
| 98 | + let validity = if inner.validity.null_count() == 0 { |
| 99 | + None |
| 100 | + } else { |
| 101 | + Some(&inner.validity) |
| 102 | + }; |
| 103 | + (&inner.column, validity) |
| 104 | + } else { |
| 105 | + (column, None) |
| 106 | + }; |
| 107 | + |
| 108 | + let column = T::try_downcast_column(column).unwrap(); |
| 109 | + if let Some(v) = validity { |
| 110 | + if v.true_count() as f64 / v.len() as f64 >= SELECTIVITY_THRESHOLD { |
| 111 | + for (data, valid) in T::iter_column(&column).zip(v.iter()) { |
| 112 | + if valid { |
| 113 | + self.hll.add_object(&data); |
| 114 | + } |
| 115 | + } |
| 116 | + } else { |
| 117 | + TrueIdxIter::new(v.len(), Some(v)).for_each(|idx| { |
| 118 | + let val = unsafe { T::index_column_unchecked(&column, idx) }; |
| 119 | + self.hll.add_object(&val); |
| 120 | + }) |
| 121 | + } |
| 122 | + } else { |
| 123 | + for value in T::iter_column(&column) { |
| 124 | + self.hll.add_object(&value); |
| 125 | + } |
| 126 | + } |
| 127 | + } |
| 128 | + |
| 129 | + fn update_scalar(&mut self, scalar: &ScalarRef) { |
| 130 | + if matches!(scalar, ScalarRef::Null) { |
| 131 | + return; |
| 132 | + } |
| 133 | + |
| 134 | + let val = T::try_downcast_scalar(scalar).unwrap(); |
| 135 | + self.hll.add_object(&val); |
| 136 | + } |
| 137 | + |
| 138 | + fn finalize(&self) -> u64 { |
| 139 | + self.hll.count() as u64 |
| 140 | + } |
| 141 | +} |
0 commit comments