Skip to content

Commit

Permalink
feat(query): support interval function (#17093)
Browse files Browse the repository at this point in the history
* feat(query): support interval function

* use register_simple_domain_type_cmp & MonthDayNano -> MonthDayMicros
  • Loading branch information
TCeason authored Dec 24, 2024
1 parent ff6cb3e commit b72f3e1
Show file tree
Hide file tree
Showing 12 changed files with 751 additions and 12 deletions.
4 changes: 2 additions & 2 deletions src/common/column/src/types/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -85,8 +85,8 @@ pub enum PrimitiveType {
Float64,
/// Two i32 representing days and ms
DaysMs,
/// months_days_ns(i32, i32, i64)
MonthDayNano,
/// months_days_micros(i32, i32, i64)
MonthDayMicros,
}

mod private {
Expand Down
42 changes: 37 additions & 5 deletions src/common/column/src/types/native.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,10 @@
// See the License for the specific language governing permissions and
// limitations under the License.

use std::cmp::Ordering;
use std::convert::TryFrom;
use std::hash::Hash;
use std::hash::Hasher;
use std::ops::Neg;
use std::panic::RefUnwindSafe;

Expand Down Expand Up @@ -252,11 +255,7 @@ impl NativeType for days_ms {
Copy,
Clone,
Default,
PartialEq,
PartialOrd,
Ord,
Eq,
Hash,
Zeroable,
Pod,
Serialize,
Expand All @@ -268,6 +267,33 @@ impl NativeType for days_ms {
#[repr(C)]
pub struct months_days_micros(pub i128);

const MICROS_PER_DAY: i64 = 24 * 3600 * 1_000_000;
const MICROS_PER_MONTH: i64 = 30 * MICROS_PER_DAY;

impl Hash for months_days_micros {
fn hash<H: Hasher>(&self, state: &mut H) {
self.total_micros().hash(state)
}
}
impl PartialEq for months_days_micros {
fn eq(&self, other: &Self) -> bool {
self.total_micros() == other.total_micros()
}
}
impl PartialOrd for months_days_micros {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}

impl Ord for months_days_micros {
fn cmp(&self, other: &Self) -> Ordering {
let total_micros = self.total_micros();
let other_micros = other.total_micros();
total_micros.cmp(&other_micros)
}
}

impl months_days_micros {
pub fn new(months: i32, days: i32, microseconds: i64) -> Self {
let months_bits = (months as i128) << 96;
Expand All @@ -291,10 +317,16 @@ impl months_days_micros {
pub fn microseconds(&self) -> i64 {
(self.0 & 0xFFFFFFFFFFFFFFFF) as i64
}

pub fn total_micros(&self) -> i64 {
(self.months() as i64 * MICROS_PER_MONTH)
+ (self.days() as i64 * MICROS_PER_DAY)
+ self.microseconds()
}
}

impl NativeType for months_days_micros {
const PRIMITIVE: PrimitiveType = PrimitiveType::MonthDayNano;
const PRIMITIVE: PrimitiveType = PrimitiveType::MonthDayMicros;
type Bytes = [u8; 16];
#[inline]
fn to_le_bytes(&self) -> Self::Bytes {
Expand Down
10 changes: 9 additions & 1 deletion src/common/io/tests/it/interval.rs
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ fn test_interval_from_string() {
}

#[test]
fn test_string_to_interval() {
fn test_interval_to_string() {
let tests = vec![
(
Interval {
Expand Down Expand Up @@ -142,6 +142,14 @@ fn test_string_to_interval() {
},
"-1 day 1:23:45",
),
(
Interval {
months: 0,
days: -225,
micros: -52550000000,
},
"-225 days -14:35:50",
),
(
Interval {
months: 0,
Expand Down
2 changes: 1 addition & 1 deletion src/common/native/src/write/primitive.rs
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,7 @@ pub(crate) fn write_primitive<T: NativeType, W: Write>(

PrimitiveType::Float16 => unimplemented!(),
PrimitiveType::DaysMs => unimplemented!(),
PrimitiveType::MonthDayNano => unimplemented!(),
PrimitiveType::MonthDayMicros => unimplemented!(),
PrimitiveType::UInt128 => unimplemented!(),
}
w.write_all(scratch.as_slice())?;
Expand Down
6 changes: 6 additions & 0 deletions src/query/functions/src/scalars/comparison.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ use databend_common_expression::types::DataType;
use databend_common_expression::types::DateType;
use databend_common_expression::types::EmptyArrayType;
use databend_common_expression::types::GenericType;
use databend_common_expression::types::IntervalType;
use databend_common_expression::types::MutableBitmap;
use databend_common_expression::types::NumberClass;
use databend_common_expression::types::NumberType;
Expand Down Expand Up @@ -65,6 +66,7 @@ pub fn register(registry: &mut FunctionRegistry) {
register_array_cmp(registry);
register_tuple_cmp(registry);
register_like(registry);
register_interval_cmp(registry);
}

pub const ALL_COMP_FUNC_NAMES: &[&str] = &["eq", "noteq", "lt", "lte", "gt", "gte", "contains"];
Expand Down Expand Up @@ -226,6 +228,10 @@ fn register_timestamp_cmp(registry: &mut FunctionRegistry) {
register_simple_domain_type_cmp!(registry, TimestampType);
}

fn register_interval_cmp(registry: &mut FunctionRegistry) {
register_simple_domain_type_cmp!(registry, IntervalType);
}

fn register_boolean_cmp(registry: &mut FunctionRegistry) {
registry.register_comparison_2_arg::<BooleanType, BooleanType, _, _>(
"eq",
Expand Down
Loading

0 comments on commit b72f3e1

Please sign in to comment.