Skip to content

Commit 8f821d9

Browse files
jayzhan211jayzhan211
committed
first draft
Signed-off-by: jayzhan211 <jayzhan211@gmail.com>
1 parent f26c540 commit 8f821d9

2 files changed

Lines changed: 254 additions & 23 deletions

File tree

‎datafusion/core/tests/sqllogictests/test_files/array.slt‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -714,6 +714,14 @@ select make_array(f0) from fixed_size_list_array
714714
----
715715
[[1, 2], [3, 4]]
716716

717+
query ?
718+
select array_concat(column1, make_array(0)) from arrays_values_without_nulls;
719+
----
720+
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 0]
721+
[11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 0]
722+
[21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 0]
723+
[31, 32, 33, 34, 35, 26, 37, 38, 39, 40, 0]
724+
717725

718726
### Delete tables
719727

‎datafusion/physical-expr/src/array_expressions.rs‎

Lines changed: 246 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -373,35 +373,184 @@ pub fn array_prepend(args: &[ArrayRef]) -> Result<ArrayRef> {
373373
Ok(res)
374374
}
375375

376+
macro_rules! concat_internal {
377+
($args:expr, $DataType:ident, $ArrayType:ident) => {{
378+
let list_arrays =
379+
downcast_vec!($args, ListArray).collect::<Result<Vec<&ListArray>>>()?;
380+
381+
let column_len = $args[0].len();
382+
383+
let mut array_values =
384+
vec![
385+
downcast_arg!(new_empty_array(&DataType::$DataType), $ArrayType).clone();
386+
column_len
387+
];
388+
389+
let mut array_lens = vec![0; column_len];
390+
391+
for list_arr in list_arrays.iter() {
392+
for (idx, (value, arr)) in
393+
array_values.iter_mut().zip(list_arr.iter()).enumerate()
394+
{
395+
match arr {
396+
Some(child_array) => {
397+
// Skip null array
398+
if child_array.as_any().downcast_ref::<NullArray>().is_some() {
399+
continue;
400+
}
401+
402+
let child_array = downcast_arg!(child_array, $ArrayType);
403+
let concat_args = compute::concat(&[value, child_array])?;
404+
*value = downcast_arg!(concat_args, $ArrayType).clone();
405+
array_lens[idx] += child_array.len();
406+
}
407+
None => {
408+
// TODO: Not support error
409+
}
410+
}
411+
}
412+
}
413+
414+
let offsets: Vec<i32> = std::iter::once(0)
415+
.chain(array_lens.iter().scan(0, |state, &x| {
416+
*state += x;
417+
Some(*state as i32)
418+
}))
419+
.collect();
420+
421+
let values: Vec<&dyn Array> =
422+
array_values.iter().map(|arr| arr as &dyn Array).collect();
423+
let merged_values = compute::concat(values.as_slice())?;
424+
425+
let field = Arc::new(Field::new("item", DataType::$DataType, true));
426+
427+
Ok(Arc::new(ListArray::try_new(
428+
field,
429+
OffsetBuffer::new(offsets.into()),
430+
Arc::new(merged_values),
431+
None,
432+
)?))
433+
}};
434+
}
435+
436+
fn concat_(args: &[ArrayRef]) -> Result<ArrayRef> {
437+
// println!("args: {:?}", args);
438+
let list_arrays =
439+
downcast_vec!(args, ListArray).collect::<Result<Vec<&ListArray>>>()?;
440+
// println!("list_arrays: {:?}", list_arrays);
441+
442+
let column_len = args[0].len();
443+
444+
let mut array_values =
445+
vec![
446+
downcast_arg!(new_empty_array(&DataType::Int64), Int64Array).clone();
447+
column_len
448+
];
449+
450+
let mut array_lens = vec![0; column_len];
451+
452+
for list_arr in list_arrays.iter() {
453+
for (idx, (value, arr)) in
454+
array_values.iter_mut().zip(list_arr.iter()).enumerate()
455+
{
456+
match arr {
457+
Some(child_array) => {
458+
// Skip null array
459+
if child_array.as_any().downcast_ref::<NullArray>().is_some() {
460+
continue;
461+
}
462+
463+
let child_array = downcast_arg!(child_array, Int64Array);
464+
let concat_args = compute::concat(&[value, child_array])?;
465+
*value = downcast_arg!(concat_args, Int64Array).clone();
466+
array_lens[idx] += child_array.len();
467+
}
468+
None => {
469+
// TODO: Not support error
470+
}
471+
}
472+
}
473+
}
474+
475+
let offsets: Vec<i32> = std::iter::once(0)
476+
.chain(array_lens.iter().scan(0, |state, &x| {
477+
*state += x;
478+
Some(*state as i32)
479+
}))
480+
.collect();
481+
482+
let values: Vec<&dyn Array> =
483+
array_values.iter().map(|arr| arr as &dyn Array).collect();
484+
let merged_values = compute::concat(values.as_slice())?;
485+
486+
let field = Arc::new(Field::new("item", DataType::Int64, true));
487+
488+
Ok(Arc::new(ListArray::try_new(
489+
field,
490+
OffsetBuffer::new(offsets.into()),
491+
Arc::new(merged_values),
492+
None,
493+
)?))
494+
}
495+
496+
fn old_concat(args: &[ArrayRef]) -> Result<ArrayRef> {
497+
let list_arrays =
498+
downcast_vec!(args, ListArray).collect::<Result<Vec<&ListArray>>>()?;
499+
let len: usize = list_arrays.iter().map(|a| a.values().len()).sum();
500+
let capacity = Capacities::Array(list_arrays.iter().map(|a| a.len()).sum());
501+
let array_data: Vec<_> = list_arrays.iter().map(|a| a.to_data()).collect::<Vec<_>>();
502+
let array_data = array_data.iter().collect();
503+
let mut mutable = MutableArrayData::with_capacities(array_data, false, capacity);
504+
505+
for (i, a) in list_arrays.iter().enumerate() {
506+
mutable.extend(i, 0, a.len())
507+
}
508+
509+
let builder = mutable.into_builder();
510+
let list = builder
511+
.len(1)
512+
.buffers(vec![Buffer::from_slice_ref([0, len as i32])])
513+
.build()
514+
.unwrap();
515+
Ok(Arc::new(arrow::array::make_array(list)))
516+
517+
}
518+
376519
/// Array_concat/Array_cat SQL function
377520
pub fn array_concat(args: &[ArrayRef]) -> Result<ArrayRef> {
378521
match args[0].data_type() {
379522
DataType::List(field) => match field.data_type() {
380523
DataType::Null => array_concat(&args[1..]),
381-
_ => {
382-
let list_arrays = downcast_vec!(args, ListArray)
383-
.collect::<Result<Vec<&ListArray>>>()?;
384-
let len: usize = list_arrays.iter().map(|a| a.values().len()).sum();
385-
let capacity =
386-
Capacities::Array(list_arrays.iter().map(|a| a.len()).sum());
387-
let array_data: Vec<_> =
388-
list_arrays.iter().map(|a| a.to_data()).collect::<Vec<_>>();
389-
let array_data = array_data.iter().collect();
390-
let mut mutable =
391-
MutableArrayData::with_capacities(array_data, false, capacity);
392-
393-
for (i, a) in list_arrays.iter().enumerate() {
394-
mutable.extend(i, 0, a.len())
524+
DataType::List(_) => {
525+
old_concat(args)
526+
}
527+
data_type => {
528+
// println!("data_type: {:?}", data_type);
529+
// println!("args: {:?}", args);
530+
// concat_(args)
531+
// old_concat(args)
532+
533+
match data_type {
534+
DataType::Int64 => concat_internal!(args, Int64, Int64Array),
535+
DataType::Int32 => concat_internal!(args, Int32, Int32Array),
536+
DataType::Int16 => concat_internal!(args, Int16, Int16Array),
537+
DataType::Int8 => concat_internal!(args, Int8, Int8Array),
538+
DataType::UInt64 => concat_internal!(args, UInt64, UInt64Array),
539+
DataType::UInt32 => concat_internal!(args, UInt32, UInt32Array),
540+
DataType::UInt16 => concat_internal!(args, UInt16, UInt16Array),
541+
DataType::UInt8 => concat_internal!(args, UInt8, UInt8Array),
542+
DataType::Float64 => concat_internal!(args, Float64, Float64Array),
543+
DataType::Float32 => concat_internal!(args, Float32, Float32Array),
544+
DataType::Boolean => concat_internal!(args, Boolean, BooleanArray),
545+
DataType::Utf8 => concat_internal!(args, Utf8, StringArray),
546+
DataType::LargeUtf8 => {
547+
concat_internal!(args, LargeUtf8, LargeStringArray)
548+
}
549+
data_type => Err(DataFusionError::NotImplemented(format!(
550+
"Array_concat is not implemented for type '{:?}'.",
551+
data_type
552+
))),
395553
}
396-
397-
let builder = mutable.into_builder();
398-
let list = builder
399-
.len(1)
400-
.buffers(vec![Buffer::from_slice_ref([0, len as i32])])
401-
.build()
402-
.unwrap();
403-
404-
return Ok(Arc::new(arrow::array::make_array(list)));
405554
}
406555
},
407556
data_type => Err(DataFusionError::NotImplemented(format!(
@@ -410,6 +559,80 @@ pub fn array_concat(args: &[ArrayRef]) -> Result<ArrayRef> {
410559
}
411560
}
412561

562+
// [ListArray
563+
// [
564+
// PrimitiveArray<Int64>
565+
// [
566+
// 1,
567+
// 2,
568+
// 3,
569+
// 4,
570+
// 5,
571+
// 6,
572+
// 7,
573+
// 8,
574+
// 9,
575+
// 10,
576+
// ],
577+
// PrimitiveArray<Int64>
578+
// [
579+
// 11,
580+
// 12,
581+
// 13,
582+
// 14,
583+
// 15,
584+
// 16,
585+
// 17,
586+
// 18,
587+
// 19,
588+
// 20,
589+
// ],
590+
// PrimitiveArray<Int64>
591+
// [
592+
// 21,
593+
// 22,
594+
// 23,
595+
// 24,
596+
// 25,
597+
// 26,
598+
// 27,
599+
// 28,
600+
// 29,
601+
// 30,
602+
// ],
603+
// PrimitiveArray<Int64>
604+
// [
605+
// 31,
606+
// 32,
607+
// 33,
608+
// 34,
609+
// 35,
610+
// 26,
611+
// 37,
612+
// 38,
613+
// 39,
614+
// 40,
615+
// ],
616+
// ], ListArray
617+
// [
618+
// PrimitiveArray<Int64>
619+
// [
620+
// 0,
621+
// ],
622+
// PrimitiveArray<Int64>
623+
// [
624+
// 0,
625+
// ],
626+
// PrimitiveArray<Int64>
627+
// [
628+
// 0,
629+
// ],
630+
// PrimitiveArray<Int64>
631+
// [
632+
// 0,
633+
// ],
634+
// ]]
635+
413636
macro_rules! fill {
414637
($ARRAY:expr, $ELEMENT:expr, $ARRAY_TYPE:ident) => {{
415638
let arr = downcast_arg!($ARRAY, $ARRAY_TYPE);

0 commit comments

Comments
 (0)