| // Licensed to the Apache Software Foundation (ASF) under one |
| // or more contributor license agreements. See the NOTICE file |
| // distributed with this work for additional information |
| // regarding copyright ownership. The ASF licenses this file |
| // to you under the Apache License, Version 2.0 (the |
| // "License"); you may not use this file except in compliance |
| // with the License. You may obtain a copy of the License at |
| // |
| // http://www.apache.org/licenses/LICENSE-2.0 |
| // |
| // Unless required by applicable law or agreed to in writing, |
| // software distributed under the License is distributed on an |
| // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| // KIND, either express or implied. See the License for the |
| // specific language governing permissions and limitations |
| // under the License. |
| |
| use std::sync::Arc; |
| |
| use arrow_array::ArrayRef; |
| use arrow_schema::{DataType, TimeUnit}; |
| |
| use super::TransformFunction; |
| use crate::spec::{Datum, PrimitiveLiteral, PrimitiveType}; |
| |
| const NANOSECONDS_PER_MICROSECOND: i64 = 1_000; |
| |
| #[derive(Debug)] |
| pub struct Bucket { |
| mod_n: u32, |
| } |
| |
| impl Bucket { |
| pub fn new(mod_n: u32) -> Self { |
| Self { mod_n } |
| } |
| } |
| |
| impl Bucket { |
| /// When switch the hash function, we only need to change this function. |
| #[inline] |
| fn hash_bytes(mut v: &[u8]) -> i32 { |
| murmur3::murmur3_32(&mut v, 0).unwrap() as i32 |
| } |
| |
| #[inline] |
| fn hash_int(v: i32) -> i32 { |
| Self::hash_long(v as i64) |
| } |
| |
| #[inline] |
| fn hash_long(v: i64) -> i32 { |
| Self::hash_bytes(v.to_le_bytes().as_slice()) |
| } |
| |
| /// v is days from unix epoch |
| #[inline] |
| fn hash_date(v: i32) -> i32 { |
| Self::hash_int(v) |
| } |
| |
| /// v is microseconds from midnight |
| #[inline] |
| fn hash_time(v: i64) -> i32 { |
| Self::hash_long(v) |
| } |
| |
| /// v is microseconds from unix epoch |
| #[inline] |
| fn hash_timestamp(v: i64) -> i32 { |
| Self::hash_long(v) |
| } |
| |
| #[inline] |
| fn hash_str(s: &str) -> i32 { |
| Self::hash_bytes(s.as_bytes()) |
| } |
| |
| /// Decimal values are hashed using the minimum number of bytes required to hold the unscaled value as a two’s complement big-endian |
| /// ref: https://iceberg.apache.org/spec/#appendix-b-32-bit-hash-requirements |
| #[inline] |
| fn hash_decimal(v: i128) -> i32 { |
| if v == 0 { |
| return Self::hash_bytes(&[0]); |
| } |
| |
| let bytes = v.to_be_bytes(); |
| let start = if v > 0 { |
| // Positive: skip 0x00 unless next byte would appear negative |
| bytes |
| .windows(2) |
| .position(|w| w[0] != 0x00 || w[1] & 0x80 != 0) |
| .unwrap_or(15) |
| } else { |
| // Negative: skip 0xFF only if next byte stays negative |
| bytes |
| .windows(2) |
| .position(|w| w[0] != 0xFF || w[1] & 0x80 == 0) |
| .unwrap_or(15) |
| }; |
| |
| Self::hash_bytes(&bytes[start..]) |
| } |
| |
| /// def bucket_N(x) = (murmur3_x86_32_hash(x) & Integer.MAX_VALUE) % N |
| /// ref: https://iceberg.apache.org/spec/#partitioning |
| #[inline] |
| fn bucket_n(&self, v: i32) -> i32 { |
| (v & i32::MAX) % (self.mod_n as i32) |
| } |
| |
| #[inline] |
| fn bucket_int(&self, v: i32) -> i32 { |
| self.bucket_n(Self::hash_int(v)) |
| } |
| |
| #[inline] |
| fn bucket_long(&self, v: i64) -> i32 { |
| self.bucket_n(Self::hash_long(v)) |
| } |
| |
| #[inline] |
| fn bucket_decimal(&self, v: i128) -> i32 { |
| self.bucket_n(Self::hash_decimal(v)) |
| } |
| |
| #[inline] |
| fn bucket_date(&self, v: i32) -> i32 { |
| self.bucket_n(Self::hash_date(v)) |
| } |
| |
| #[inline] |
| fn bucket_time(&self, v: i64) -> i32 { |
| self.bucket_n(Self::hash_time(v)) |
| } |
| |
| #[inline] |
| fn bucket_timestamp(&self, v: i64) -> i32 { |
| self.bucket_n(Self::hash_timestamp(v)) |
| } |
| |
| #[inline] |
| fn bucket_str(&self, v: &str) -> i32 { |
| self.bucket_n(Self::hash_str(v)) |
| } |
| |
| #[inline] |
| fn bucket_bytes(&self, v: &[u8]) -> i32 { |
| self.bucket_n(Self::hash_bytes(v)) |
| } |
| |
| #[inline] |
| fn nanos_to_micros(v: i64) -> i64 { |
| v.div_euclid(NANOSECONDS_PER_MICROSECOND) |
| } |
| } |
| |
| impl TransformFunction for Bucket { |
| fn transform(&self, input: ArrayRef) -> crate::Result<ArrayRef> { |
| let res: arrow_array::Int32Array = match input.data_type() { |
| DataType::Int32 => input |
| .as_any() |
| .downcast_ref::<arrow_array::Int32Array>() |
| .unwrap() |
| .unary(|v| self.bucket_int(v)), |
| DataType::Int64 => input |
| .as_any() |
| .downcast_ref::<arrow_array::Int64Array>() |
| .unwrap() |
| .unary(|v| self.bucket_long(v)), |
| DataType::Decimal128(_, _) => input |
| .as_any() |
| .downcast_ref::<arrow_array::Decimal128Array>() |
| .unwrap() |
| .unary(|v| self.bucket_decimal(v)), |
| DataType::Date32 => input |
| .as_any() |
| .downcast_ref::<arrow_array::Date32Array>() |
| .unwrap() |
| .unary(|v| self.bucket_date(v)), |
| DataType::Time64(TimeUnit::Microsecond) => input |
| .as_any() |
| .downcast_ref::<arrow_array::Time64MicrosecondArray>() |
| .unwrap() |
| .unary(|v| self.bucket_time(v)), |
| DataType::Timestamp(TimeUnit::Microsecond, _) => input |
| .as_any() |
| .downcast_ref::<arrow_array::TimestampMicrosecondArray>() |
| .unwrap() |
| .unary(|v| self.bucket_timestamp(v)), |
| DataType::Time64(TimeUnit::Nanosecond) => input |
| .as_any() |
| .downcast_ref::<arrow_array::Time64NanosecondArray>() |
| .unwrap() |
| .unary(|v| self.bucket_time(Self::nanos_to_micros(v))), |
| DataType::Timestamp(TimeUnit::Nanosecond, _) => input |
| .as_any() |
| .downcast_ref::<arrow_array::TimestampNanosecondArray>() |
| .unwrap() |
| .unary(|v| self.bucket_timestamp(Self::nanos_to_micros(v))), |
| DataType::Utf8 => arrow_array::Int32Array::from_iter( |
| input |
| .as_any() |
| .downcast_ref::<arrow_array::StringArray>() |
| .unwrap() |
| .iter() |
| .map(|v| v.map(|v| self.bucket_str(v))), |
| ), |
| DataType::LargeUtf8 => arrow_array::Int32Array::from_iter( |
| input |
| .as_any() |
| .downcast_ref::<arrow_array::LargeStringArray>() |
| .unwrap() |
| .iter() |
| .map(|v| v.map(|v| self.bucket_str(v))), |
| ), |
| DataType::Binary => arrow_array::Int32Array::from_iter( |
| input |
| .as_any() |
| .downcast_ref::<arrow_array::BinaryArray>() |
| .unwrap() |
| .iter() |
| .map(|v| v.map(|v| self.bucket_bytes(v))), |
| ), |
| DataType::LargeBinary => arrow_array::Int32Array::from_iter( |
| input |
| .as_any() |
| .downcast_ref::<arrow_array::LargeBinaryArray>() |
| .unwrap() |
| .iter() |
| .map(|v| v.map(|v| self.bucket_bytes(v))), |
| ), |
| DataType::FixedSizeBinary(_) => arrow_array::Int32Array::from_iter( |
| input |
| .as_any() |
| .downcast_ref::<arrow_array::FixedSizeBinaryArray>() |
| .unwrap() |
| .iter() |
| .map(|v| v.map(|v| self.bucket_bytes(v))), |
| ), |
| _ => { |
| return Err(crate::Error::new( |
| crate::ErrorKind::FeatureUnsupported, |
| format!( |
| "Unsupported data type for bucket transform: {:?}", |
| input.data_type() |
| ), |
| )); |
| } |
| }; |
| Ok(Arc::new(res)) |
| } |
| |
| fn transform_literal(&self, input: &Datum) -> crate::Result<Option<Datum>> { |
| let val = match (input.data_type(), input.literal()) { |
| (PrimitiveType::Int, PrimitiveLiteral::Int(v)) => self.bucket_int(*v), |
| (PrimitiveType::Long, PrimitiveLiteral::Long(v)) => self.bucket_long(*v), |
| (PrimitiveType::Decimal { .. }, PrimitiveLiteral::Int128(v)) => self.bucket_decimal(*v), |
| (PrimitiveType::Date, PrimitiveLiteral::Int(v)) => self.bucket_date(*v), |
| (PrimitiveType::Time, PrimitiveLiteral::Long(v)) => self.bucket_time(*v), |
| (PrimitiveType::Timestamp, PrimitiveLiteral::Long(v)) => self.bucket_timestamp(*v), |
| (PrimitiveType::Timestamptz, PrimitiveLiteral::Long(v)) => self.bucket_timestamp(*v), |
| (PrimitiveType::TimestampNs, PrimitiveLiteral::Long(v)) => { |
| self.bucket_timestamp(Self::nanos_to_micros(*v)) |
| } |
| (PrimitiveType::TimestamptzNs, PrimitiveLiteral::Long(v)) => { |
| self.bucket_timestamp(Self::nanos_to_micros(*v)) |
| } |
| (PrimitiveType::String, PrimitiveLiteral::String(v)) => self.bucket_str(v.as_str()), |
| (PrimitiveType::Uuid, PrimitiveLiteral::UInt128(v)) => { |
| self.bucket_bytes(uuid::Uuid::from_u128(*v).as_ref()) |
| } |
| (PrimitiveType::Binary, PrimitiveLiteral::Binary(v)) => self.bucket_bytes(v.as_ref()), |
| (PrimitiveType::Fixed(_), PrimitiveLiteral::Binary(v)) => self.bucket_bytes(v.as_ref()), |
| _ => { |
| return Err(crate::Error::new( |
| crate::ErrorKind::FeatureUnsupported, |
| format!( |
| "Unsupported data type for bucket transform: {:?}", |
| input.data_type() |
| ), |
| )); |
| } |
| }; |
| Ok(Some(Datum::int(val))) |
| } |
| } |
| |
| #[cfg(test)] |
| mod test { |
| use std::sync::Arc; |
| |
| use arrow_array::{ArrayRef, Int32Array, TimestampMicrosecondArray, TimestampNanosecondArray}; |
| use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime}; |
| |
| use super::Bucket; |
| use crate::Result; |
| use crate::expr::PredicateOperator; |
| use crate::spec::PrimitiveType::{ |
| Binary, Date, Decimal, Fixed, Int, Long, String as StringType, Time, Timestamp, |
| TimestampNs, Timestamptz, TimestamptzNs, Uuid, |
| }; |
| use crate::spec::Type::{Primitive, Struct}; |
| use crate::spec::decimal_utils::decimal_new; |
| use crate::spec::{Datum, NestedField, PrimitiveType, StructType, Transform}; |
| use crate::transform::TransformFunction; |
| use crate::transform::test::{TestProjectionFixture, TestTransformFixture}; |
| |
| #[test] |
| fn test_bucket_transform() { |
| let trans = Transform::Bucket(8); |
| |
| let fixture = TestTransformFixture { |
| display: "bucket[8]".to_string(), |
| json: r#""bucket[8]""#.to_string(), |
| dedup_name: "bucket[8]".to_string(), |
| preserves_order: false, |
| satisfies_order_of: vec![ |
| (Transform::Bucket(8), true), |
| (Transform::Bucket(4), false), |
| (Transform::Void, false), |
| (Transform::Day, false), |
| ], |
| trans_types: vec![ |
| (Primitive(Binary), Some(Primitive(Int))), |
| (Primitive(Date), Some(Primitive(Int))), |
| ( |
| Primitive(Decimal { |
| precision: 8, |
| scale: 5, |
| }), |
| Some(Primitive(Int)), |
| ), |
| (Primitive(Fixed(8)), Some(Primitive(Int))), |
| (Primitive(Int), Some(Primitive(Int))), |
| (Primitive(Long), Some(Primitive(Int))), |
| (Primitive(StringType), Some(Primitive(Int))), |
| (Primitive(Uuid), Some(Primitive(Int))), |
| (Primitive(Time), Some(Primitive(Int))), |
| (Primitive(Timestamp), Some(Primitive(Int))), |
| (Primitive(Timestamptz), Some(Primitive(Int))), |
| (Primitive(TimestampNs), Some(Primitive(Int))), |
| (Primitive(TimestamptzNs), Some(Primitive(Int))), |
| ( |
| Struct(StructType::new(vec![ |
| NestedField::optional(1, "a", Primitive(Timestamp)).into(), |
| ])), |
| None, |
| ), |
| ], |
| }; |
| |
| fixture.assert_transform(trans); |
| } |
| |
| #[test] |
| fn test_projection_bucket_uuid() -> Result<()> { |
| let value = uuid::Uuid::from_u64_pair(123, 456); |
| let another = uuid::Uuid::from_u64_pair(456, 123); |
| |
| let fixture = TestProjectionFixture::new( |
| Transform::Bucket(10), |
| "name", |
| NestedField::required(1, "value", Primitive(Uuid)), |
| ); |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::Eq, Datum::uuid(value)), |
| Some("name = 4"), |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::NotEq, Datum::uuid(value)), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::LessThan, Datum::uuid(value)), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::LessThanOrEq, Datum::uuid(value)), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::GreaterThan, Datum::uuid(value)), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::GreaterThanOrEq, Datum::uuid(value)), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.set_predicate(PredicateOperator::In, vec![ |
| Datum::uuid(value), |
| Datum::uuid(another), |
| ]), |
| Some("name IN (4, 6)"), |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.set_predicate(PredicateOperator::NotIn, vec![ |
| Datum::uuid(value), |
| Datum::uuid(another), |
| ]), |
| None, |
| )?; |
| |
| Ok(()) |
| } |
| |
| #[test] |
| fn test_projection_bucket_fixed() -> Result<()> { |
| let value = "abcdefg".as_bytes().to_vec(); |
| let another = "abcdehij".as_bytes().to_vec(); |
| |
| let fixture = TestProjectionFixture::new( |
| Transform::Bucket(10), |
| "name", |
| NestedField::required(1, "value", Primitive(Fixed(value.len() as u64))), |
| ); |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::Eq, Datum::fixed(value.clone())), |
| Some("name = 4"), |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::NotEq, Datum::fixed(value.clone())), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::LessThan, Datum::fixed(value.clone())), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::LessThanOrEq, Datum::fixed(value.clone())), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::GreaterThan, Datum::fixed(value.clone())), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate( |
| PredicateOperator::GreaterThanOrEq, |
| Datum::fixed(value.clone()), |
| ), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.set_predicate(PredicateOperator::In, vec![ |
| Datum::fixed(value.clone()), |
| Datum::fixed(another.clone()), |
| ]), |
| Some("name IN (4, 6)"), |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.set_predicate(PredicateOperator::NotIn, vec![ |
| Datum::fixed(value.clone()), |
| Datum::fixed(another.clone()), |
| ]), |
| None, |
| )?; |
| |
| Ok(()) |
| } |
| |
| #[test] |
| fn test_projection_bucket_string() -> Result<()> { |
| let value = "abcdefg"; |
| let another = "abcdefgabc"; |
| |
| let fixture = TestProjectionFixture::new( |
| Transform::Bucket(10), |
| "name", |
| NestedField::required(1, "value", Primitive(PrimitiveType::String)), |
| ); |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::Eq, Datum::string(value)), |
| Some("name = 4"), |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::NotEq, Datum::string(value)), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::LessThan, Datum::string(value)), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::LessThanOrEq, Datum::string(value)), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::GreaterThan, Datum::string(value)), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::GreaterThanOrEq, Datum::string(value)), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.set_predicate(PredicateOperator::In, vec![ |
| Datum::string(value), |
| Datum::string(another), |
| ]), |
| Some("name IN (9, 4)"), |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.set_predicate(PredicateOperator::NotIn, vec![ |
| Datum::string(value), |
| Datum::string(another), |
| ]), |
| None, |
| )?; |
| |
| Ok(()) |
| } |
| |
| #[test] |
| fn test_projection_bucket_decimal() -> Result<()> { |
| let prev = "99.00"; |
| let curr = "100.00"; |
| let next = "101.00"; |
| |
| let fixture = TestProjectionFixture::new( |
| Transform::Bucket(10), |
| "name", |
| NestedField::required( |
| 1, |
| "value", |
| Primitive(Decimal { |
| precision: 9, |
| scale: 2, |
| }), |
| ), |
| ); |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::Eq, Datum::decimal_from_str(curr)?), |
| Some("name = 2"), |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::NotEq, Datum::decimal_from_str(curr)?), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::LessThan, Datum::decimal_from_str(curr)?), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate( |
| PredicateOperator::LessThanOrEq, |
| Datum::decimal_from_str(curr)?, |
| ), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate( |
| PredicateOperator::GreaterThan, |
| Datum::decimal_from_str(curr)?, |
| ), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate( |
| PredicateOperator::GreaterThanOrEq, |
| Datum::decimal_from_str(curr)?, |
| ), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.set_predicate(PredicateOperator::In, vec![ |
| Datum::decimal_from_str(next)?, |
| Datum::decimal_from_str(curr)?, |
| Datum::decimal_from_str(prev)?, |
| ]), |
| Some("name IN (2, 6)"), |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.set_predicate(PredicateOperator::NotIn, vec![ |
| Datum::decimal_from_str(curr)?, |
| Datum::decimal_from_str(next)?, |
| ]), |
| None, |
| )?; |
| |
| Ok(()) |
| } |
| |
| #[test] |
| fn test_projection_bucket_long() -> Result<()> { |
| let value = 100; |
| let fixture = TestProjectionFixture::new( |
| Transform::Bucket(10), |
| "name", |
| NestedField::required(1, "value", Primitive(Long)), |
| ); |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::Eq, Datum::long(value)), |
| Some("name = 6"), |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::NotEq, Datum::long(value)), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::LessThan, Datum::long(value)), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::LessThanOrEq, Datum::long(value)), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::GreaterThan, Datum::long(value)), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::GreaterThanOrEq, Datum::long(value)), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.set_predicate(PredicateOperator::In, vec![ |
| Datum::long(value - 1), |
| Datum::long(value), |
| Datum::long(value + 1), |
| ]), |
| Some("name IN (8, 7, 6)"), |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.set_predicate(PredicateOperator::NotIn, vec![ |
| Datum::long(value), |
| Datum::long(value + 1), |
| ]), |
| None, |
| )?; |
| |
| Ok(()) |
| } |
| |
| #[test] |
| fn test_projection_bucket_integer() -> Result<()> { |
| let value = 100; |
| |
| let fixture = TestProjectionFixture::new( |
| Transform::Bucket(10), |
| "name", |
| NestedField::required(1, "value", Primitive(Int)), |
| ); |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::Eq, Datum::int(value)), |
| Some("name = 6"), |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::NotEq, Datum::int(value)), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::LessThan, Datum::int(value)), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::LessThanOrEq, Datum::int(value)), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::GreaterThan, Datum::int(value)), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.binary_predicate(PredicateOperator::GreaterThanOrEq, Datum::int(value)), |
| None, |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.set_predicate(PredicateOperator::In, vec![ |
| Datum::int(value - 1), |
| Datum::int(value), |
| Datum::int(value + 1), |
| ]), |
| Some("name IN (8, 7, 6)"), |
| )?; |
| |
| fixture.assert_projection( |
| &fixture.set_predicate(PredicateOperator::NotIn, vec![ |
| Datum::int(value), |
| Datum::int(value + 1), |
| ]), |
| None, |
| )?; |
| |
| Ok(()) |
| } |
| |
| #[test] |
| fn test_hash() { |
| // test int |
| assert_eq!(Bucket::hash_int(34), 2017239379); |
| // test long |
| assert_eq!(Bucket::hash_long(34), 2017239379); |
| // test decimal |
| assert_eq!(Bucket::hash_decimal(1420), -500754589); |
| // test date |
| let date = NaiveDate::from_ymd_opt(2017, 11, 16).unwrap(); |
| assert_eq!( |
| Bucket::hash_date( |
| date.signed_duration_since(NaiveDate::from_ymd_opt(1970, 1, 1).unwrap()) |
| .num_days() as i32 |
| ), |
| -653330422 |
| ); |
| // test time |
| let time = NaiveTime::from_hms_opt(22, 31, 8).unwrap(); |
| assert_eq!( |
| Bucket::hash_time( |
| time.signed_duration_since(NaiveTime::from_hms_opt(0, 0, 0).unwrap()) |
| .num_microseconds() |
| .unwrap() |
| ), |
| -662762989 |
| ); |
| // test timestamp |
| let timestamp = |
| NaiveDateTime::parse_from_str("2017-11-16 22:31:08", "%Y-%m-%d %H:%M:%S").unwrap(); |
| assert_eq!( |
| Bucket::hash_timestamp( |
| timestamp |
| .signed_duration_since( |
| NaiveDateTime::parse_from_str("1970-01-01 00:00:00", "%Y-%m-%d %H:%M:%S") |
| .unwrap() |
| ) |
| .num_microseconds() |
| .unwrap() |
| ), |
| -2047944441 |
| ); |
| // test timestamp with tz |
| let timestamp = DateTime::parse_from_rfc3339("2017-11-16T14:31:08-08:00").unwrap(); |
| assert_eq!( |
| Bucket::hash_timestamp( |
| timestamp |
| .signed_duration_since( |
| DateTime::parse_from_rfc3339("1970-01-01T00:00:00-00:00").unwrap() |
| ) |
| .num_microseconds() |
| .unwrap() |
| ), |
| -2047944441 |
| ); |
| // test str |
| assert_eq!(Bucket::hash_str("iceberg"), 1210000089); |
| // test uuid |
| assert_eq!( |
| Bucket::hash_bytes( |
| [ |
| 0xF7, 0x9C, 0x3E, 0x09, 0x67, 0x7C, 0x4B, 0xBD, 0xA4, 0x79, 0x3F, 0x34, 0x9C, |
| 0xB7, 0x85, 0xE7 |
| ] |
| .as_ref() |
| ), |
| 1488055340 |
| ); |
| // test fixed and binary |
| assert_eq!( |
| Bucket::hash_bytes([0x00, 0x01, 0x02, 0x03].as_ref()), |
| -188683207 |
| ); |
| } |
| |
| #[test] |
| fn test_hash_decimal_with_negative_value() { |
| // Test cases from GitHub issue #1981 |
| assert_eq!(Bucket::hash_decimal(1), -463810133); |
| assert_eq!(Bucket::hash_decimal(-1), -43192051); |
| |
| // Additional test cases for edge case values |
| assert_eq!(Bucket::hash_decimal(0), Bucket::hash_decimal(0)); |
| assert_eq!(Bucket::hash_decimal(127), Bucket::hash_decimal(127)); |
| assert_eq!(Bucket::hash_decimal(-128), Bucket::hash_decimal(-128)); |
| |
| // Test minimum representation is used |
| // -1 should hash as [0xFF] not [0xFF, 0xFF, ..., 0xFF] |
| // 128 should hash as [0x00, 0x80] not [0x00, 0x00, ..., 0x80] |
| assert_eq!(Bucket::hash_decimal(128), Bucket::hash_bytes(&[0x00, 0x80])); |
| assert_eq!( |
| Bucket::hash_decimal(-129), |
| Bucket::hash_bytes(&[0xFF, 0x7F]) |
| ); |
| } |
| |
| #[test] |
| fn test_int_literal() { |
| let bucket = Bucket::new(10); |
| assert_eq!( |
| bucket.transform_literal(&Datum::int(34)).unwrap().unwrap(), |
| Datum::int(9) |
| ); |
| } |
| |
| #[test] |
| fn test_long_literal() { |
| let bucket = Bucket::new(10); |
| assert_eq!( |
| bucket.transform_literal(&Datum::long(34)).unwrap().unwrap(), |
| Datum::int(9) |
| ); |
| } |
| |
| #[test] |
| fn test_decimal_literal() { |
| let bucket = Bucket::new(10); |
| assert_eq!( |
| bucket |
| .transform_literal(&Datum::decimal(decimal_new(1420, 0)).unwrap()) |
| .unwrap() |
| .unwrap(), |
| Datum::int(9) |
| ); |
| } |
| |
| #[test] |
| fn test_date_literal() { |
| let bucket = Bucket::new(100); |
| assert_eq!( |
| bucket |
| .transform_literal(&Datum::date(17486)) |
| .unwrap() |
| .unwrap(), |
| Datum::int(26) |
| ); |
| } |
| |
| #[test] |
| fn test_time_literal() { |
| let bucket = Bucket::new(100); |
| assert_eq!( |
| bucket |
| .transform_literal(&Datum::time_micros(81068000000).unwrap()) |
| .unwrap() |
| .unwrap(), |
| Datum::int(59) |
| ); |
| } |
| |
| #[test] |
| fn test_timestamp_literal() { |
| let bucket = Bucket::new(100); |
| assert_eq!( |
| bucket |
| .transform_literal(&Datum::timestamp_micros(1510871468000000)) |
| .unwrap() |
| .unwrap(), |
| Datum::int(7) |
| ); |
| } |
| |
| #[test] |
| fn test_str_literal() { |
| let bucket = Bucket::new(100); |
| assert_eq!( |
| bucket |
| .transform_literal(&Datum::string("iceberg")) |
| .unwrap() |
| .unwrap(), |
| Datum::int(89) |
| ); |
| } |
| |
| #[test] |
| fn test_uuid_literal() { |
| let bucket = Bucket::new(100); |
| assert_eq!( |
| bucket |
| .transform_literal(&Datum::uuid( |
| "F79C3E09-677C-4BBD-A479-3F349CB785E7".parse().unwrap() |
| )) |
| .unwrap() |
| .unwrap(), |
| Datum::int(40) |
| ); |
| } |
| |
| #[test] |
| fn test_binary_literal() { |
| let bucket = Bucket::new(128); |
| assert_eq!( |
| bucket |
| .transform_literal(&Datum::binary(b"\x00\x01\x02\x03".to_vec())) |
| .unwrap() |
| .unwrap(), |
| Datum::int(57) |
| ); |
| } |
| |
| #[test] |
| fn test_fixed_literal() { |
| let bucket = Bucket::new(128); |
| assert_eq!( |
| bucket |
| .transform_literal(&Datum::fixed(b"foo".to_vec())) |
| .unwrap() |
| .unwrap(), |
| Datum::int(32) |
| ); |
| } |
| |
| #[test] |
| fn test_timestamptz_literal() { |
| let bucket = Bucket::new(100); |
| assert_eq!( |
| bucket |
| .transform_literal(&Datum::timestamptz_micros(1510871468000000)) |
| .unwrap() |
| .unwrap(), |
| Datum::int(7) |
| ); |
| } |
| |
| #[test] |
| fn test_timestamp_ns_literal() { |
| let bucket = Bucket::new(100); |
| let ns_value = 1510871468000000i64 * 1000; |
| assert_eq!( |
| bucket |
| .transform_literal(&Datum::timestamp_nanos(ns_value)) |
| .unwrap() |
| .unwrap(), |
| Datum::int(7) |
| ); |
| } |
| |
| #[test] |
| fn test_timestamptz_ns_literal() { |
| let bucket = Bucket::new(100); |
| let ns_value = 1510871468000000i64 * 1000; |
| assert_eq!( |
| bucket |
| .transform_literal(&Datum::timestamptz_nanos(ns_value)) |
| .unwrap() |
| .unwrap(), |
| Datum::int(7) |
| ); |
| } |
| |
| #[test] |
| fn test_transform_timestamp_nanos_and_micros_array_equivalence() { |
| let bucket = Bucket::new(100); |
| let micros_value = 1510871468000000; |
| let nanos_value = micros_value * 1000; |
| |
| let micro_array = TimestampMicrosecondArray::from_iter_values(vec![micros_value]); |
| let nano_array = TimestampNanosecondArray::from_iter_values(vec![nanos_value]); |
| |
| let transformed_micro: ArrayRef = bucket.transform(Arc::new(micro_array)).unwrap(); |
| let transformed_nano: ArrayRef = bucket.transform(Arc::new(nano_array)).unwrap(); |
| |
| let micro_result = transformed_micro |
| .as_any() |
| .downcast_ref::<Int32Array>() |
| .unwrap(); |
| let nano_result = transformed_nano |
| .as_any() |
| .downcast_ref::<Int32Array>() |
| .unwrap(); |
| |
| assert_eq!(micro_result.value(0), nano_result.value(0)); |
| } |
| |
| #[test] |
| fn test_nanos_to_micros_uses_floor_division() { |
| assert_eq!(Bucket::nanos_to_micros(i64::MIN), -9_223_372_036_854_776); |
| assert_eq!(Bucket::nanos_to_micros(1_001), 1); |
| assert_eq!(Bucket::nanos_to_micros(1_000), 1); |
| assert_eq!(Bucket::nanos_to_micros(0), 0); |
| assert_eq!(Bucket::nanos_to_micros(-1_000), -1); |
| assert_eq!(Bucket::nanos_to_micros(-1_001), -2); |
| assert_eq!(Bucket::nanos_to_micros(i64::MAX), 9_223_372_036_854_775); |
| } |
| |
| #[test] |
| fn test_negative_timestamp_nanos_round_down_to_micros() { |
| let bucket = Bucket::new(100); |
| let micros_value = -1_510_871_468_000_001; |
| let nanos_value = micros_value * 1000 + 999; |
| |
| let expected = bucket |
| .transform_literal(&Datum::timestamp_micros(micros_value)) |
| .unwrap(); |
| assert_eq!( |
| bucket |
| .transform_literal(&Datum::timestamp_nanos(nanos_value)) |
| .unwrap(), |
| expected |
| ); |
| assert_eq!( |
| bucket |
| .transform_literal(&Datum::timestamptz_nanos(nanos_value)) |
| .unwrap(), |
| expected |
| ); |
| |
| let micro_array = TimestampMicrosecondArray::from_iter_values([micros_value]); |
| let nano_array = TimestampNanosecondArray::from_iter_values([nanos_value]); |
| let micro_result = bucket.transform(Arc::new(micro_array)).unwrap(); |
| let nano_result = bucket.transform(Arc::new(nano_array)).unwrap(); |
| let micro_result = micro_result.as_any().downcast_ref::<Int32Array>().unwrap(); |
| let nano_result = nano_result.as_any().downcast_ref::<Int32Array>().unwrap(); |
| assert_eq!(micro_result, nano_result); |
| } |
| } |
