Skip to content

Commit

Permalink
chore: refactoring some functions (#16265)
Browse files Browse the repository at this point in the history 
chore: improve some functions
  • Loading branch information
@b41sh
b41sh authored Aug 16, 2024
1 parent 1593b2a commit 695e169
Show file tree
Hide file tree
Showing 4 changed files with 288 additions and 194 deletions.
153 changes: 124 additions & 29 deletions src/query/functions/src/scalars/bitmap.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,12 @@ pub fn register(registry: &mut FunctionRegistry) {
"to_bitmap",
|_, _| FunctionDomain::MayThrow,
vectorize_with_builder_1_arg::<StringType, BitmapType>(|s, builder, ctx| {
if let Some(validity) = &ctx.validity {
if !validity.get_bit(builder.len()) {
builder.commit_row();
return;
}
}
match parse_bitmap(s.as_bytes()) {
Ok(rb) => {
rb.serialize_into(&mut builder.data).unwrap();
Expand All @@ -61,7 +67,13 @@ pub fn register(registry: &mut FunctionRegistry) {
registry.register_passthrough_nullable_1_arg::<UInt64Type, BitmapType, _, _>(
"to_bitmap",
|_, _| FunctionDomain::Full,
vectorize_with_builder_1_arg::<UInt64Type, BitmapType>(|arg, builder, _| {
vectorize_with_builder_1_arg::<UInt64Type, BitmapType>(|arg, builder, ctx| {
if let Some(validity) = &ctx.validity {
if !validity.get_bit(builder.len()) {
builder.commit_row();
return;
}
}
let mut rb = RoaringTreemap::new();
rb.insert(arg);

Expand All @@ -76,7 +88,13 @@ pub fn register(registry: &mut FunctionRegistry) {
registry.register_passthrough_nullable_1_arg::<ArrayType<NullableType<NumberType<NUM_TYPE>>>, BitmapType, _, _>(
"build_bitmap",
|_, _| FunctionDomain::Full,
vectorize_with_builder_1_arg::<ArrayType<NullableType<NumberType<NUM_TYPE>>>, BitmapType>(|arg, builder, _ctx| {
vectorize_with_builder_1_arg::<ArrayType<NullableType<NumberType<NUM_TYPE>>>, BitmapType>(|arg, builder, ctx| {
if let Some(validity) = &ctx.validity {
if !validity.get_bit(builder.len()) {
builder.commit_row();
return;
}
}
let mut rb = RoaringTreemap::new();
for a in arg.iter() {
if let Some(a) = a {
Expand All @@ -100,6 +118,12 @@ pub fn register(registry: &mut FunctionRegistry) {
"build_bitmap",
|_, _| FunctionDomain::MayThrow,
vectorize_with_builder_1_arg::<ArrayType<NullableType<NumberType<NUM_TYPE>>>, BitmapType>(|arg, builder, ctx| {
if let Some(validity) = &ctx.validity {
if !validity.get_bit(builder.len()) {
builder.commit_row();
return;
}
}
let mut rb = RoaringTreemap::new();
for a in arg.iter() {
if let Some(a) = a {
Expand All @@ -124,13 +148,19 @@ pub fn register(registry: &mut FunctionRegistry) {
"bitmap_count",
|_, _| FunctionDomain::Full,
vectorize_with_builder_1_arg::<BitmapType, UInt64Type>(|arg, builder, ctx| {
if let Some(validity) = &ctx.validity {
if !validity.get_bit(builder.len()) {
builder.push(0_u64);
return;
}
}
match deserialize_bitmap(arg) {
Ok(rb) => {
builder.push(rb.len());
}
Err(e) => {
builder.push(0_u64);
ctx.set_error(builder.len(), e.to_string());
builder.push(0_u64);
}
}
}),
Expand All @@ -142,6 +172,12 @@ pub fn register(registry: &mut FunctionRegistry) {
"to_string",
|_, _| FunctionDomain::MayThrow,
vectorize_with_builder_1_arg::<BitmapType, StringType>(|b, builder, ctx| {
if let Some(validity) = &ctx.validity {
if !validity.get_bit(builder.len()) {
builder.commit_row();
return;
}
}
match deserialize_bitmap(b) {
Ok(rb) => {
let raw = rb.into_iter().collect::<Vec<_>>();
Expand All @@ -161,13 +197,21 @@ pub fn register(registry: &mut FunctionRegistry) {
"bitmap_contains",
|_, _, _| FunctionDomain::Full,
vectorize_with_builder_2_arg::<BitmapType, UInt64Type, BooleanType>(
|b, item, builder, ctx| match deserialize_bitmap(b) {
Ok(rb) => {
builder.push(rb.contains(item));
|b, item, builder, ctx| {
if let Some(validity) = &ctx.validity {
if !validity.get_bit(builder.len()) {
builder.push(false);
return;
}
}
Err(e) => {
builder.push(false);
ctx.set_error(builder.len(), e.to_string());
match deserialize_bitmap(b) {
Ok(rb) => {
builder.push(rb.contains(item));
}
Err(e) => {
ctx.set_error(builder.len(), e.to_string());
builder.push(false);
}
}
},
),
Expand All @@ -177,35 +221,51 @@ pub fn register(registry: &mut FunctionRegistry) {
"bitmap_subset_limit",
|_, _, _, _| FunctionDomain::MayThrow,
vectorize_with_builder_3_arg::<BitmapType, UInt64Type, UInt64Type, BitmapType>(
|b, range_start, limit, builder, ctx| match deserialize_bitmap(b) {
Ok(rb) => {
let collection = rb.iter().filter(|x| x >= &range_start).take(limit as usize);
let subset_bitmap = RoaringTreemap::from_iter(collection);
subset_bitmap.serialize_into(&mut builder.data).unwrap();
builder.commit_row();
|b, range_start, limit, builder, ctx| {
if let Some(validity) = &ctx.validity {
if !validity.get_bit(builder.len()) {
builder.commit_row();
return;
}
}
Err(e) => {
ctx.set_error(builder.len(), e.to_string());
match deserialize_bitmap(b) {
Ok(rb) => {
let collection = rb.iter().filter(|x| x >= &range_start).take(limit as usize);
let subset_bitmap = RoaringTreemap::from_iter(collection);
subset_bitmap.serialize_into(&mut builder.data).unwrap();
}
Err(e) => {
ctx.set_error(builder.len(), e.to_string());
}
}
},
builder.commit_row();
}
),
);

registry.register_passthrough_nullable_3_arg::<BitmapType, UInt64Type, UInt64Type, BitmapType, _, _>(
"bitmap_subset_in_range",
|_, _, _, _| FunctionDomain::MayThrow,
vectorize_with_builder_3_arg::<BitmapType, UInt64Type, UInt64Type, BitmapType>(
|b, start, end, builder, ctx| match deserialize_bitmap(b) {
Ok(rb) => {
let collection = rb.iter().filter(|x| x >= &start && x < &end);
let subset_bitmap = RoaringTreemap::from_iter(collection);
subset_bitmap.serialize_into(&mut builder.data).unwrap();
builder.commit_row();
|b, start, end, builder, ctx| {
if let Some(validity) = &ctx.validity {
if !validity.get_bit(builder.len()) {
builder.commit_row();
return;
}
}
Err(e) => {
ctx.set_error(builder.len(), e.to_string());
match deserialize_bitmap(b) {
Ok(rb) => {
let collection = rb.iter().filter(|x| x >= &start && x < &end);
let subset_bitmap = RoaringTreemap::from_iter(collection);
subset_bitmap.serialize_into(&mut builder.data).unwrap();
}
Err(e) => {
ctx.set_error(builder.len(), e.to_string());
}
}
},
builder.commit_row();
}
),
);

Expand All @@ -214,26 +274,31 @@ pub fn register(registry: &mut FunctionRegistry) {
|_, _, _, _| FunctionDomain::MayThrow,
vectorize_with_builder_3_arg::<BitmapType, UInt64Type, UInt64Type, BitmapType>(
|b, offset, length, builder, ctx| {
if let Some(validity) = &ctx.validity {
if !validity.get_bit(builder.len()) {
builder.commit_row();
return;
}
}
match deserialize_bitmap(b) {
Ok(rb) => {
let subset_start = offset;
let subset_length = length;
if subset_start >= b.len() as u64 {
let rb = RoaringTreemap::new();
rb.serialize_into(&mut builder.data).unwrap();
builder.commit_row();
} else {
let adjusted_length = (subset_start + subset_length).min(b.len() as u64) - subset_start;
let subset_bitmap = &rb.into_iter().collect::<Vec<_>>()[subset_start as usize..(subset_start + adjusted_length) as usize];
let rb = RoaringTreemap::from_iter(subset_bitmap.iter());
rb.serialize_into(&mut builder.data).unwrap();
builder.commit_row();
}
}
Err(e) => {
ctx.set_error(builder.len(), e.to_string());
}
}
builder.commit_row();
},
),
);
Expand All @@ -243,6 +308,12 @@ pub fn register(registry: &mut FunctionRegistry) {
|_, _, _| FunctionDomain::Full,
vectorize_with_builder_2_arg::<BitmapType, BitmapType, BooleanType>(
|b, items, builder, ctx| {
if let Some(validity) = &ctx.validity {
if !validity.get_bit(builder.len()) {
builder.push(false);
return;
}
}
let rb = match deserialize_bitmap(b) {
Ok(rb) => rb,
Err(e) => {
Expand All @@ -269,6 +340,12 @@ pub fn register(registry: &mut FunctionRegistry) {
|_, _, _| FunctionDomain::Full,
vectorize_with_builder_2_arg::<BitmapType, BitmapType, BooleanType>(
|b, items, builder, ctx| {
if let Some(validity) = &ctx.validity {
if !validity.get_bit(builder.len()) {
builder.push(false);
return;
}
}
let rb = match deserialize_bitmap(b) {
Ok(rb) => rb,
Err(e) => {
Expand All @@ -294,6 +371,12 @@ pub fn register(registry: &mut FunctionRegistry) {
"bitmap_max",
|_, _| FunctionDomain::MayThrow,
vectorize_with_builder_1_arg::<BitmapType, UInt64Type>(|b, builder, ctx| {
if let Some(validity) = &ctx.validity {
if !validity.get_bit(builder.len()) {
builder.push(0);
return;
}
}
let val = match deserialize_bitmap(b) {
Ok(rb) => match rb.max() {
Some(val) => val,
Expand All @@ -315,6 +398,12 @@ pub fn register(registry: &mut FunctionRegistry) {
"bitmap_min",
|_, _| FunctionDomain::MayThrow,
vectorize_with_builder_1_arg::<BitmapType, UInt64Type>(|b, builder, ctx| {
if let Some(validity) = &ctx.validity {
if !validity.get_bit(builder.len()) {
builder.push(0);
return;
}
}
let val = match deserialize_bitmap(b) {
Ok(rb) => match rb.min() {
Some(val) => val,
Expand Down Expand Up @@ -382,6 +471,12 @@ fn bitmap_logic_operate(
ctx: &mut EvalContext,
op: LogicOp,
) {
if let Some(validity) = &ctx.validity {
if !validity.get_bit(builder.len()) {
builder.commit_row();
return;
}
}
let Some(rb1) = deserialize_bitmap(arg1)
.map_err(|e| {
ctx.set_error(builder.len(), e.to_string());
Expand Down
Loading

0 comments on commit 695e169

Please sign in to comment.