1use proc_macro::TokenStream;
2
3use quote::{format_ident, quote};
4
5fn uc_first(s: &str) -> String {
6 let mut c = s.chars();
7 match c.next() {
8 None => String::new(),
9 Some(f) => f.to_uppercase().collect::<String>() + c.as_str(),
10 }
11}
12
13
14#[proc_macro]
15pub fn create(_input: TokenStream) -> TokenStream {
16 let mut output = quote!();
17 for bit_size in vec!["32", "64"] {
18 for nullsafe in vec![true, false] {
19 let struct_name = format_ident!("Roaring{}{}Create", bit_size, if nullsafe {"Nullsafe"} else {""});
20 let map_type_name = format_ident!("Map{}", bit_size);
21 let return_type = quote!(&'a [u8]);
22 let return_type = if nullsafe { quote!(#return_type) } else { quote!(Option<#return_type>) };
23 let maybe_null_cnf = if nullsafe { quote!(_cfg) } else { quote! {cfg} };
24 let maybe_null = if nullsafe { quote!() } else { quote! {cfg.set_maybe_null(true);} };
25 let operation = if bit_size == "64" {
26 quote! {
27 let value = value as u64;
28 map.insert(value);
29 }
30 } else {
31 quote! {
32 if let Some(value) = value.to_u32() {
33 map.insert(value);
34 } else {
35 return Err(ProcessError);
36 }
37 }
38 };
39 let null_result = if nullsafe {
40 quote! {
41 let map = <#map_type_name>::new();
42 self.vec.reserve(map.serialized_size());
43 if let Ok(_) = map.serialize_into(&mut self.vec) {
44 Ok(&self.vec[0..self.vec.len()])
45 } else {
46 Err(ProcessError)
47 }
48 }
49 } else {
50 quote!(Ok(None))
51 };
52 let result = quote!(&self.vec[0..self.vec.len()]);
53 let result = if nullsafe { quote!(Ok(#result)) } else { quote!(Ok(Some(#result))) };
54 output.extend(quote! {
55 #[derive(Default)]
56 struct #struct_name {
57 vec: Vec<u8>
58 }
59
60 #[register]
61 impl BasicUdf for #struct_name {
62 type Returns<'a> = #return_type;
63
64 fn init(#maybe_null_cnf: &UdfCfg<Init>, args: &ArgList<Init>) -> Result<Self, String> {
65 #maybe_null
66 if args.len() < 1 {
67 return Err(format!("Expected one or more arguments; Got {} arguments.", args.len()));
68 }
69 for index in 0..args.len() {
70 if !args.get(index).unwrap().value().is_int() {
71 return Err(format!("{} argument mast be INTEGER or NULL.", index));
72 }
73 }
74 Ok(Self::default())
75 }
76
77 fn process<'a>(&'a mut self, _cfg: &UdfCfg<Process>, args: &ArgList<Process>, _error: Option<NonZeroU8>) -> Result<Self::Returns<'a>, ProcessError> {
78 let mut map = <#map_type_name>::new();
79 for arg in args.iter() {
80 if let Some(value) = arg.value().as_int() {
81 #operation
82 } else {
83 return {#null_result};
84 }
85 }
86 self.vec.reserve(map.serialized_size());
87 if let Ok(_) = map.serialize_into(&mut self.vec) {
88 #result
89 } else {
90 Err(ProcessError)
91 }
92 }
93 }
94 });
95 }
96 }
97 output.into()
98}
99
100#[proc_macro]
101pub fn map_int_op(_input: TokenStream) -> TokenStream {
102 let mut output = quote!();
103 for bit_size in vec!["32", "64"] {
104 for nullsafe in vec![true, false] {
105 for op in vec!["insert", "remove"] {
106 let struct_name = format_ident!("Roaring{}{}{}", bit_size, if nullsafe {"Nullsafe"} else {""}, uc_first(op));
107 let map_type_name = format_ident!("Map{}", bit_size);
108 let op_fn_name = format_ident!("{}", op);
109 let return_type = quote!(&'a [u8]);
110 let return_type = if nullsafe { quote!(#return_type) } else { quote!(Option<#return_type>) };
111 let maybe_null_cnf = if nullsafe { quote!(_cfg) } else { quote! {cfg} };
112 let maybe_null = if nullsafe { quote!() } else { quote! {cfg.set_maybe_null(true);} };
113 let operation = quote! {map.#op_fn_name(value)};
114 let operation = if bit_size == "64" {
115 quote! {
116 let value = value as u64;
117 #operation;
118 }
119 } else {
120 quote! {
121 if let Some(value) = value.to_u32() {
122 #operation;
123 } else {
124 return Err(ProcessError);
125 }
126 }
127 };
128 let null_result = if nullsafe {
129 quote! {
130 let map = <#map_type_name>::new();
131 self.vec.reserve(map.serialized_size());
132 if let Ok(_) = map.serialize_into(&mut self.vec) {
133 Ok(&self.vec[0..self.vec.len()])
134 } else {
135 Err(ProcessError)
136 }
137 }
138 } else {
139 quote!(Ok(None))
140 };
141 let result = quote!(&self.vec[0..self.vec.len()]);
142 let result = if nullsafe { quote!(Ok(#result)) } else { quote!(Ok(Some(#result))) };
143 output.extend(quote! {
144 #[derive(Default)]
145 struct #struct_name {
146 vec: Vec<u8>
147 }
148
149 #[register]
150 impl BasicUdf for #struct_name {
151 type Returns<'a> = #return_type;
152
153 fn init(#maybe_null_cnf: &UdfCfg<Init>, args: &ArgList<Init>) -> Result<Self, String> {
154 #maybe_null
155 if args.len() < 2 {
156 return Err(format!("Expected two or more arguments; Got {} arguments.", args.len()));
157 }
158 if !args.get(0).unwrap().value().is_string() {
159 return Err(format!("First argument mast be CLOB (bitmap) or NULL."));
160 }
161 for index in 1..args.len() {
162 if !args.get(index).unwrap().value().is_int() {
163 return Err(format!("{} argument mast be INTEGER or NULL.", index));
164 }
165 }
166 Ok(Self::default())
167 }
168
169 fn process<'a>(&'a mut self, _cfg: &UdfCfg<Process>, args: &ArgList<Process>, _error: Option<NonZeroU8>) -> Result<Self::Returns<'a>, ProcessError> {
170 let arg = args.get(0).unwrap().value();
171 let bytes = arg.as_bytes();
172 if let Some(bytes) = bytes {
173 if let Ok(mut map) = <#map_type_name>::deserialize_from(bytes) {
174 for index in 1..args.len() {
175 let arg = args.get(index).unwrap();
176 if let Some(value) = arg.value().as_int() {
177 {#operation};
178 } else {
179 return {#null_result};
180 }
181 }
182 self.vec.reserve(map.serialized_size());
183 if let Ok(_) = map.serialize_into(&mut self.vec) {
184 return {#result};
185 } else {
186 return Err(ProcessError);
187 }
188 } else {
189 return Err(ProcessError);
190 }
191 } else {
192 return {#null_result};
193 }
194 }
195 }
196 });
197 }
198 }
199 }
200 output.into()
201}
202
203#[proc_macro]
204pub fn map_op(_input: TokenStream) -> TokenStream {
205 let mut output = quote!();
206 for bit_size in vec!["32", "64"] {
207 for nullsafe in vec![true, false] {
208 for count in vec![true, false] {
209 for op in vec!["or", "xor", "and"] {
210 let struct_name = format_ident!("Roaring{}{}{}{}", bit_size, if nullsafe {"Nullsafe"} else {""}, uc_first(op), if count {"Count"} else {""});
211 let struct_fields = if count { quote!() } else { quote!(vec: Vec<u8>) };
212 let map_type_name = format_ident!("Map{}", bit_size);
213 let op_fn_name = format_ident!("bit{}_assign", op);
214 let return_type = if count { quote!(i64) } else { quote!(&'a [u8]) };
215 let return_type = if nullsafe { quote!(#return_type) } else { quote!(Option<#return_type>) };
216 let maybe_null_cnf = if nullsafe { quote!(_cfg) } else { quote! {cfg} };
217 let maybe_null = if nullsafe { quote!() } else { quote! {cfg.set_maybe_null(true);} };
218 let on_null_arg = if nullsafe { quote!(()) } else { quote!(return Ok(None)) };
219 let on_null_agg = if nullsafe { quote!(<#map_type_name>::new()) } else { quote!(return Ok(None)) };
220 let bytes_result = quote!(&self.vec[0..self.vec.len()]);
221 let bytes_result = if nullsafe { quote!(#bytes_result) } else { quote!(Some(#bytes_result)) };
222 let count_result = quote!(agg.len() as i64);
223 let count_result = if nullsafe { quote!(#count_result) } else { quote!(Some(#count_result)) };
224 let count_null_result = if nullsafe { quote!(0i64) } else { quote!(None) };
225 let result = if count {
226 quote! {
227 match maybe_agg {
228 Some(agg) => Ok(#count_result),
229 None => Ok(#count_null_result)
230 }
231 }
232 } else {
233 quote! {
234 let agg = match maybe_agg {
235 Some(agg) => agg,
236 None => #on_null_agg,
237 };
238 self.vec.reserve(agg.serialized_size());
239 if let Ok(_) = agg.serialize_into(&mut self.vec) {
240 Ok(#bytes_result)
241 } else {
242 Err(ProcessError)
243 }
244 }
245 };
246 output.extend(quote! {
249 #[derive(Default)]
250 struct #struct_name {
251 #struct_fields
252 }
253
254 #[register]
255 impl BasicUdf for #struct_name {
256 type Returns<'a> = #return_type;
257
258 fn init(#maybe_null_cnf: &UdfCfg<Init>, args: &ArgList<Init>) -> Result<Self, String> {
259 #maybe_null
260 if args.len() < 2 {
261 return Err(format!("Expected two or more CLOB (bitmap) or NULL arguments; Got {} arguments.", args.len()));
262 }
263 for arg in args.iter() {
264 if let SqlResult::String(_) = arg.value() {} else {
265 return Err(format!("Only CLOB (bitmap) or NULL arguments are allowed."));
266 }
267 }
268 Ok(Self::default())
269 }
270
271 fn process<'a>(&'a mut self, _cfg: &UdfCfg<Process>, args: &ArgList<Process>, _error: Option<NonZeroU8>) -> Result<Self::Returns<'a>, ProcessError> {
272 let mut maybe_agg: Option<#map_type_name> = None;
273 for arg in args.iter() {
274 match arg.value().as_bytes() {
275 None => #on_null_arg,
276 Some(bytes) => {
277 if let Ok(map) = <#map_type_name>::deserialize_from(bytes) {
278 match maybe_agg {
279 None => maybe_agg = Some(map),
280 Some(ref mut agg) => agg.#op_fn_name(map),
281 }
282 } else {
283 return Err(ProcessError);
284 }
285 }
286 }
287 }
288 #result
289 }
290 }
291 })
292 }
293 }
294 }
295 }
296 output.into()
297}
298
299#[proc_macro]
300pub fn contains(_input: TokenStream) -> TokenStream {
301 let mut output = quote!();
302 for bit_size in vec!["32", "64"] {
303 for nullsafe in vec![true, false] {
304 let struct_name = format_ident!("Roaring{}{}Contains", bit_size, if nullsafe {"Nullsafe"} else {""});
305 let cast_function_name = format_ident!("to_u{}", bit_size);
306 let map_type_name = format_ident!("Map{}", bit_size);
307 let return_type = if nullsafe { quote!(i64) } else { quote!(Option<i64>) };
308 let null_result = if nullsafe { quote!(0i64) } else { quote!(None) };
309 let result = quote!(map.contains(value).into());
310 let result = if nullsafe { quote!(#result) } else { quote!(Some(#result)) };
311 let maybe_null_cnf = if nullsafe { quote!(_cfg) } else { quote! {cfg} };
312 let maybe_null = if nullsafe { quote!() } else { quote! {cfg.set_maybe_null(true);} };
313 output.extend(quote! {
314 #[derive(Default)]
315 struct #struct_name;
316
317 #[register]
318 impl BasicUdf for #struct_name {
319 type Returns<'a> = #return_type;
320
321 fn init(#maybe_null_cnf: &UdfCfg<Init>, args: &ArgList<Init>) -> Result<Self, String> {
322 if args.len() < 2 {
323 return Err(format!("Expected two arguments; Got {} arguments.", args.len()));
324 }
325 if !args.get(0).unwrap().value().is_string() {
326 return Err(format!("First argument mast be CLOB (bitmap) or NULL."));
327 }
328 if !args.get(1).unwrap().value().is_int() {
329 return Err(format!("Second argument mast be INT or NULL."));
330 }
331 #maybe_null
332 Ok(Self::default())
333 }
334
335 fn process<'a>(&'a mut self, _cfg: &UdfCfg<Process>, args: &ArgList<Process>, _error: Option<NonZeroU8>) -> Result<Self::Returns<'a>, ProcessError> {
336 let arg0 = args.get(0).unwrap().value();
337 let arg1 = args.get(1).unwrap().value();
338 let bytes = arg0.as_bytes();
339 let value = arg1.as_int();
340 if let (Some(bytes), Some(value)) = (bytes, value) {
341 if let (Ok(map), Some(value)) = (#map_type_name::deserialize_from(bytes), value.#cast_function_name()) {
342 Ok(#result)
343 } else {
344 Err(ProcessError)
345 }
346 } else {
347 Ok(#null_result)
348 }
349 }
350 }
351 });
352 }
353 }
354 output.into()
355}
356
357#[proc_macro]
358pub fn count(_input: TokenStream) -> TokenStream {
359 let mut output = quote!();
360 for bit_size in vec!["32", "64"] {
361 for nullsafe in vec![true, false] {
362 let struct_name = format_ident!("Roaring{}{}Count", bit_size, if nullsafe {"Nullsafe"} else {""});
363 let map_type_name = format_ident!("Map{}", bit_size);
364 let return_type = if nullsafe { quote!(i64) } else { quote!(Option<i64>) };
365 let null_result = if nullsafe { quote!(0i64) } else { quote!(None) };
366 let result = quote!(map.len() as i64);
367 let result = if nullsafe { quote!(#result) } else { quote!(Some(#result)) };
368 let maybe_null_cnf = if nullsafe { quote!(_cfg) } else { quote! {cfg} };
369 let maybe_null = if nullsafe { quote!() } else { quote! {cfg.set_maybe_null(true);} };
370 output.extend(quote! {
371 #[derive(Default)]
372 struct #struct_name;
373
374 #[register]
375 impl BasicUdf for #struct_name {
376 type Returns<'a> = #return_type;
377
378 fn init(#maybe_null_cnf: &UdfCfg<Init>, args: &ArgList<Init>) -> Result<Self, String> {
379 #maybe_null
380 if args.len() != 1 {
381 return Err(format!("Expected one arguments; Got {} arguments.", args.len()));
382 }
383 if !args.get(0).unwrap().value().is_string() {
384 return Err(format!("First argument mast be CLOB (bitmap) or NULL."));
385 }
386 Ok(Self::default())
387 }
388
389 fn process<'a>(&'a mut self, _cfg: &UdfCfg<Process>, args: &ArgList<Process>, _error: Option<NonZeroU8>) -> Result<Self::Returns<'a>, ProcessError> {
390 let arg = args.get(0).unwrap().value();
391 let bytes = arg.as_bytes();
392 if let Some(bytes) = bytes {
393 if let Ok(map) = <#map_type_name>::deserialize_from(bytes) {
394 Ok(#result)
395 } else {
396 Err(ProcessError)
397 }
398 } else {
399 Ok(#null_result)
400 }
401 }
402 }
403 });
404 }
405 }
406 output.into()
407}
408
409#[proc_macro]
410pub fn json(_input: TokenStream) -> TokenStream {
411 let mut output = quote!();
412 for bit_size in vec!["32", "64"] {
413 for nullsafe in vec![true, false] {
414 let struct_name = format_ident!("Roaring{}{}Json", bit_size, if nullsafe {"Nullsafe"} else {""});
415 let map_type_name = format_ident!("Map{}", bit_size);
416 let return_type = if nullsafe { quote!(&'a [u8]) } else { quote!(Option<&'a [u8]>) };
417 let null_result = if nullsafe {
418 quote! {
419 self.vec.extend(b"[]");
420 Ok(&self.vec[0..self.vec.len()])
421 }
422 } else {
423 quote!(Ok(None))
424 };
425 let result = quote!(&self.vec[0..self.vec.len()]);
426 let result = if nullsafe { quote!(#result) } else { quote!(Some(#result)) };
427 let maybe_null_cnf = if nullsafe { quote!(_cfg) } else { quote! {cfg} };
428 let maybe_null = if nullsafe { quote!() } else { quote! {cfg.set_maybe_null(true);} };
429 output.extend(quote! {
430 #[derive(Default)]
431 struct #struct_name {
432 vec: Vec<u8>
433 }
434
435 #[register]
436 impl BasicUdf for #struct_name {
437 type Returns<'a> = #return_type;
438
439 fn init(#maybe_null_cnf: &UdfCfg<Init>, args: &ArgList<Init>) -> Result<Self, String> {
440 #maybe_null
441 if args.len() != 1 {
442 return Err(format!("Expected one arguments; Got {} arguments.", args.len()));
443 }
444 if !args.get(0).unwrap().value().is_string() {
445 return Err(format!("First argument mast be CLOB (bitmap) or NULL."));
446 }
447 Ok(Self::default())
448 }
449
450 fn process<'a>(&'a mut self, _cfg: &UdfCfg<Process>, args: &ArgList<Process>, _error: Option<NonZeroU8>) -> Result<Self::Returns<'a>, ProcessError> {
451 let arg = args.get(0).unwrap().value();
452 let bytes = arg.as_bytes();
453 if let Some(bytes) = bytes {
454 if let Ok(map) = <#map_type_name>::deserialize_from(bytes) {
455 let mut count = map.len() as usize;
456 let capacity = ((10usize+1)*count)+2; self.vec.reserve(capacity);
458 const COMA: &[u8] = b",";
459 self.vec.extend(b"[");
460 for item in map.iter() {
461 self.vec.extend(item.to_string().as_bytes());
462 count-=1;
463 if count != 0 {
464 self.vec.extend(COMA);
465 }
466 }
467 self.vec.extend(b"]");
468 Ok(#result)
469 } else {
470 Err(ProcessError)
471 }
472 } else {
473 #null_result
474 }
475 }
476 }
477 });
478 }
479 }
480 output.into()
481}
482
483
484#[proc_macro]
485pub fn group_map_op(_input: TokenStream) -> TokenStream {
486 let mut output = quote!();
487 for bit_size in vec!["32", "64"] {
488 for op in vec!["or", "xor", "and"] {
489 let ucf_op = uc_first(op);
490 let struct_name = format_ident!("Roaring{}Group{}", bit_size, ucf_op);
491 let struct_nullsafe_name = format_ident!("Roaring{}NullsafeGroup{}", bit_size, ucf_op);
492 let struct_count_name = format_ident!("Roaring{}Group{}Count", bit_size, ucf_op);
493 let struct_nullsafe_count_name = format_ident!("Roaring{}NullsafeGroup{}Count", bit_size, ucf_op);
494 let map_type_name = format_ident!("Map{}", bit_size);
495 let op_fn_name = format_ident!("bit{}_assign", op);
496 let add = quote! {
497 fn add(&mut self, _cfg: &UdfCfg<Process>, args: &ArgList<Process>, _error: Option<NonZeroU8>) -> Result<(), NonZeroU8> {
498 if let Some(bytes) = args.get(0).unwrap().value().as_bytes() {
499 if let Ok(map) = <#map_type_name>::deserialize_from(bytes) {
500 if let Some(ref mut acc) = self.map {
501 acc.#op_fn_name(map);
502 Ok(())
503 } else {
504 self.map = Some(map);
505 Ok(())
506 }
507 } else {
508 Err(NonZeroU8::new(1u8).unwrap())
509 }
510 } else {
511 Ok(())
512 }
513 }
514 };
515 let init = quote! {
516 fn init(_cfg: &UdfCfg<Init>, args: &ArgList<Init>) -> Result<Self, String> {
517 if args.len() != 1 {
518 return Err(format!("Expected one arguments; Got {} arguments.", args.len()));
519 }
520 if !args.get(0).unwrap().value().is_string() {
521 return Err(format!("{} argument mast be BLOB or NULL.", 0));
522 }
523 Ok(Self::default())
524 }
525 };
526 let init_nullable = quote! {
527 fn init(cfg: &UdfCfg<Init>, args: &ArgList<Init>) -> Result<Self, String> {
528 cfg.set_maybe_null(true);
529 if args.len() != 1 {
530 return Err(format!("Expected one arguments; Got {} arguments.", args.len()));
531 }
532 if !args.get(0).unwrap().value().is_string() {
533 return Err(format!("{} argument mast be BLOB or NULL.", 0));
534 }
535 Ok(Self::default())
536 }
537 };
538 let clear = quote! {
539 fn clear(&mut self, _cfg: &UdfCfg<Process>, _error: Option<NonZeroU8>) -> Result<(), NonZeroU8> {
540 self.map = None;
541 self.vec.clear();
542 Ok(())
543 }
544 };
545 let clear_count = quote! {
546 fn clear(&mut self, _cfg: &UdfCfg<Process>, _error: Option<NonZeroU8>) -> Result<(), NonZeroU8> {
547 self.map = None;
548 Ok(())
549 }
550 };
551 let fields = quote! {
552 map: Option<#map_type_name>,
553 vec: Vec<u8>,
554 };
555 let fields_count = quote! {
556 map: Option<#map_type_name>,
557 };
558 output.extend(quote! {
559 #[derive(Default)]
560 struct #struct_name {
561 #fields
562 }
563
564 #[register]
565 impl AggregateUdf for #struct_name {
566 #clear
567 #add
568 }
569
570 #[register]
571 impl BasicUdf for #struct_name {
572 type Returns<'a> = Option<&'a [u8]>;
573
574 #init_nullable
575
576 fn process<'a>(&'a mut self, _cfg: &UdfCfg<Process>, _args: &ArgList<Process>, _error: Option<NonZeroU8>) -> Result<Self::Returns<'a>, ProcessError> {
577 if let Some(ref agg) = self.map {
578 let capacity = agg.serialized_size();
579 if self.vec.capacity() < capacity {
580 self.vec.reserve(capacity - self.vec.capacity());
581 }
582 if let Ok(_) = agg.serialize_into(&mut self.vec) {
583 Ok(Some(&self.vec[0..self.vec.len()]))
584 } else {
585 Err(ProcessError)
586 }
587 } else {
588 Ok(None)
589 }
590 }
591 }
592
593 #[derive(Default)]
594 struct #struct_count_name {
595 #fields_count
596 }
597
598 #[register]
599 impl AggregateUdf for #struct_count_name {
600 #clear_count
601 #add
602 }
603
604 #[register]
605 impl BasicUdf for #struct_count_name {
606 type Returns<'a> = Option<i64>;
607
608 #init_nullable
609
610 fn process<'a>(&'a mut self, _cfg: &UdfCfg<Process>, _args: &ArgList<Process>, _error: Option<NonZeroU8>) -> Result<Self::Returns<'a>, ProcessError> {
611 if let Some(ref agg) = self.map {
612 Ok(Some(agg.len() as i64))
613 } else {
614 Ok(None)
615 }
616 }
617 }
618
619 #[derive(Default)]
620 struct #struct_nullsafe_name {
621 #fields
622 }
623
624 #[register]
625 impl AggregateUdf for #struct_nullsafe_name {
626 #clear
627 #add
628 }
629
630 #[register]
631 impl BasicUdf for #struct_nullsafe_name {
632 type Returns<'a> = &'a [u8];
633
634 #init
635
636 fn process<'a>(&'a mut self, _cfg: &UdfCfg<Process>, _args: &ArgList<Process>, _error: Option<NonZeroU8>) -> Result<Self::Returns<'a>, ProcessError> {
637 let map;
638 let null;
639 if let Some(ref agg) = self.map {
640 map = agg;
641 } else {
642 null = <#map_type_name>::new();
643 map = &null;
644 }
645 let capacity = map.serialized_size();
646 if self.vec.capacity() < capacity {
647 self.vec.reserve(capacity - self.vec.capacity());
648 }
649 if let Ok(_) = map.serialize_into(&mut self.vec) {
650 Ok(&self.vec[0..self.vec.len()])
651 } else {
652 Err(ProcessError)
653 }
654 }
655 }
656
657 #[derive(Default)]
658 struct #struct_nullsafe_count_name {
659 #fields_count
660 }
661
662 #[register]
663 impl AggregateUdf for #struct_nullsafe_count_name {
664 #clear_count
665 #add
666 }
667
668 #[register]
669 impl BasicUdf for #struct_nullsafe_count_name {
670 type Returns<'a> = i64;
671
672 #init
673
674 fn process<'a>(&'a mut self, _cfg: &UdfCfg<Process>, _args: &ArgList<Process>, _error: Option<NonZeroU8>) -> Result<Self::Returns<'a>, ProcessError> {
675 if let Some(ref agg) = self.map {
676 Ok(agg.len() as i64)
677 } else {
678 Ok(0i64)
679 }
680 }
681 }
682 })
683 }
684 }
685 output.into()
686}
687
688#[proc_macro]
689pub fn group_create(_input: TokenStream) -> TokenStream {
690 let mut output = quote!();
691 for bit_size in vec!["32", "64"] {
692 for nullsafe in vec![true, false] {
693 let struct_name = format_ident!("Roaring{}{}GroupCreate", bit_size, if nullsafe {"Nullsafe"} else {""});
694 let map_type_name = format_ident!("Map{}", bit_size);
695 let return_type = quote!(&'a [u8]);
696 let return_type = if nullsafe { quote!(#return_type) } else { quote!(Option<#return_type>) };
697 let maybe_null_cnf = if nullsafe { quote!(_cfg) } else { quote! {cfg} };
698 let maybe_null = if nullsafe { quote!() } else { quote! {cfg.set_maybe_null(true);} };
699 let bytes_result = quote!(&self.vec[0..self.vec.len()]);
700 let bytes_result = if nullsafe { quote!(#bytes_result) } else { quote!(Some(#bytes_result)) };
701 let extend_capacity = quote! {
702 let capacity = self.map.serialized_size();
703 if self.vec.capacity() < capacity {
704 self.vec.reserve(capacity - self.vec.capacity());
705 }
706 };
707 let null_bytes_result = if nullsafe { quote!(&self.vec[0..self.vec.len()]) } else { quote!(None) };
708 let null_result = quote! {
709 #extend_capacity
710 if let Ok(_) = self.map.serialize_into(&mut self.vec) {
711 Ok(#null_bytes_result)
712 } else {
713 Err(ProcessError)
714 }
715 };
716 let result = quote! {
717 #extend_capacity
718 if let Ok(_) = self.map.serialize_into(&mut self.vec) {
719 Ok(#bytes_result)
720 } else {
721 Err(ProcessError)
722 }
723 };
724
725 let mut operations = quote!();
726 for (func, op) in vec![("add", "insert"), ("remove", "remove")].iter() {
727 let func = format_ident!("{}", func);
728 let op = format_ident!("{}", op);
729 let operation = if bit_size == "64" {
730 quote! {
731 self.map.#op(value as u64);
732 Ok(())
733 }
734 } else {
735 quote! {
736 if let Some(value) = value.to_u32() {
737 self.map.#op(value);
738 Ok(())
739 } else {
740 Err(NonZeroU8::new(1u8).unwrap())
741 }
742 }
743 };
744 operations.extend(quote! {
745 fn #func(&mut self, _cfg: &UdfCfg<Process>, args: &ArgList<Process>, _error: Option<NonZeroU8>) -> Result<(), NonZeroU8> {
746 if let Some(value) = args.get(0).unwrap().value().as_int() {
747 #operation
748 } else {
749 Ok(())
750 }
751 }
752 });
753 }
754 output.extend(quote! {
755 #[derive(Default)]
756 struct #struct_name {
757 map: #map_type_name,
758 vec: Vec<u8>,
759 }
760
761 #[register]
762 impl AggregateUdf for #struct_name {
763 fn clear(&mut self, _cfg: &UdfCfg<Process>, _error: Option<NonZeroU8>) -> Result<(), NonZeroU8> {
764 self.map.clear();
765 self.vec.clear();
766 Ok(())
767 }
768 #operations
769 }
770
771 #[register]
772 impl BasicUdf for #struct_name {
773 type Returns<'a> = #return_type;
774
775 fn init(#maybe_null_cnf: &UdfCfg<Init>, args: &ArgList<Init>) -> Result<Self, String> {
776 #maybe_null
777 if args.len() != 1 {
778 return Err(format!("Expected one arguments; Got {} arguments.", args.len()));
779 }
780 if !args.get(0).unwrap().value().is_int() {
781 return Err(format!("{} argument mast be INTEGER or NULL.", 0));
782 }
783 Ok(Self::default())
784 }
785
786 fn process<'a>(&'a mut self, _cfg: &UdfCfg<Process>, _args: &ArgList<Process>, _error: Option<NonZeroU8>) -> Result<Self::Returns<'a>, ProcessError> {
787 if(self.map.len() > 0) {
788 #result
789 } else {
790 #null_result
791 }
792 }
793 }
794 })
795 }
796 }
797 output.into()
798}