Skip to content

Commit 515a9bc

Browse files
DandandanDandandan
committed
Support list
1 parent 5ff1e73 commit 515a9bc

2 files changed

Lines changed: 77 additions & 242 deletions

File tree

‎datafusion/src/physical_plan/hash_aggregate.rs‎

Lines changed: 26 additions & 233 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ use crate::physical_plan::{
3535
use crate::scalar::ScalarValue;
3636

3737
use arrow::{
38-
array::{Array, Date64Array, UInt32Builder},
38+
array::{Array, UInt32Builder},
3939
error::{ArrowError, Result as ArrowResult},
4040
};
4141
use arrow::{
@@ -969,80 +969,61 @@ fn create_batch_from_map(
969969
if accumulators.is_empty() {
970970
return Ok(RecordBatch::new_empty(Arc::new(output_schema.to_owned())));
971971
}
972-
// 1. for each key
973-
// 2. create single-row ArrayRef with all group expressions
974-
// 3. create single-row ArrayRef with all aggregate states or values
975-
// 4. collect all in a vector per key of vec<ArrayRef>, vec[i][j]
976-
// 5. concatenate the arrays over the second index [j] into a single vec<ArrayRef>.
977-
let (_, (group_bys, accs, _)) = accumulators.iter().nth(0).unwrap();
978-
let group_by_data_types: Vec<DataType> = group_bys
979-
.iter()
980-
.map(|x| ScalarValue::from(x).get_datatype())
981-
.collect();
982-
let mut acc_data_types: Vec<Vec<DataType>> = vec![];
972+
let (_, (_, accs, _)) = accumulators.iter().nth(0).unwrap();
973+
let mut acc_data_types: Vec<usize> = vec![];
974+
975+
// Calculate number/shape of state arrays
983976
match mode {
984977
AggregateMode::Partial => {
985978
for acc in accs.iter() {
986979
let state = acc
987980
.state()
988981
.map_err(DataFusionError::into_arrow_external_error)?;
989-
acc_data_types
990-
.push(state.iter().map(ScalarValue::get_datatype).collect());
982+
acc_data_types.push(state.len());
991983
}
992984
}
993985
AggregateMode::Final | AggregateMode::FinalPartitioned => {
994-
for acc in accs {
995-
acc_data_types.push(vec![acc
996-
.evaluate()
997-
.map_err(DataFusionError::into_arrow_external_error)?
998-
.get_datatype()]);
986+
for _ in accs {
987+
acc_data_types.push(1);
999988
}
1000989
}
1001990
}
1002991

1003-
let mut arrays = (0..num_group_expr)
992+
let mut columns = (0..num_group_expr)
1004993
.map(|i| {
1005-
scalar_to_array(
1006-
group_by_data_types[i].clone(),
1007-
accumulators.iter().map(|(_, (group_by_values, _, _))| {
1008-
ScalarValue::from(&group_by_values[i])
1009-
}),
1010-
)
994+
ScalarValue::iter_to_array(accumulators.into_iter().map(
995+
|(_, (group_by_values, _, _))| ScalarValue::from(&group_by_values[i]),
996+
))
1011997
})
1012998
.collect::<Result<Vec<_>>>()
1013999
.map_err(|x| x.into_arrow_external_error())?;
10141000

10151001
// add state / evaluated arrays
1016-
for (x, state_dt) in acc_data_types.iter().enumerate() {
1017-
for y in 0..state_dt.len() {
1002+
for (x, &state_len) in acc_data_types.iter().enumerate() {
1003+
for y in 0..state_len {
10181004
match mode {
10191005
AggregateMode::Partial => {
1020-
let res = scalar_to_array(
1021-
state_dt[y].clone(),
1022-
accumulators.iter().map(|(_, (_, accumulator, _))| {
1023-
accumulator[x].state().unwrap()[y].clone()
1024-
}),
1025-
)
1006+
let res = ScalarValue::iter_to_array(accumulators.into_iter().map(
1007+
|(_, (_, accumulator, _))| {
1008+
let x = accumulator[x].state().unwrap();
1009+
x[y].clone()
1010+
},
1011+
))
10261012
.map_err(DataFusionError::into_arrow_external_error)?;
1027-
arrays.push(res);
1013+
1014+
columns.push(res);
10281015
}
10291016
AggregateMode::Final | AggregateMode::FinalPartitioned => {
1030-
let res = scalar_to_array(
1031-
state_dt[y].clone(),
1032-
accumulators.iter().map(|(_, (_, accumulator, _))| {
1033-
accumulator[x].evaluate().unwrap()
1034-
}),
1035-
)
1017+
let res = ScalarValue::iter_to_array(accumulators.into_iter().map(
1018+
|(_, (_, accumulator, _))| accumulator[x].evaluate().unwrap(),
1019+
))
10361020
.map_err(DataFusionError::into_arrow_external_error)?;
1037-
arrays.push(res);
1021+
columns.push(res);
10381022
}
10391023
}
10401024
}
10411025
}
10421026

1043-
// 5.
1044-
let columns = arrays;
1045-
10461027
// cast output if needed (e.g. for types like Dictionary where
10471028
// the intermediate GroupByScalar type was not the same as the
10481029
// output
@@ -1064,194 +1045,6 @@ fn create_accumulators(
10641045
.collect::<Result<Vec<_>>>()
10651046
}
10661047

1067-
fn scalar_to_array(
1068-
data_type: DataType,
1069-
iter: impl IntoIterator<Item = ScalarValue>,
1070-
) -> Result<ArrayRef> {
1071-
match data_type {
1072-
DataType::Int8 => Ok(Arc::new(
1073-
iter.into_iter()
1074-
.map(|x| match x {
1075-
ScalarValue::Int8(val) => val,
1076-
v => panic!("Unexpected type in scalar_to_array {:?}", v),
1077-
})
1078-
.collect::<Int8Array>(),
1079-
)),
1080-
DataType::Int16 => Ok(Arc::new(
1081-
iter.into_iter()
1082-
.map(|x| match x {
1083-
ScalarValue::Int16(val) => val,
1084-
v => panic!("Unexpected type in scalar_to_array {:?}", v),
1085-
})
1086-
.collect::<Int16Array>(),
1087-
)),
1088-
DataType::Int32 => Ok(Arc::new(
1089-
iter.into_iter()
1090-
.map(|x| match x {
1091-
ScalarValue::Int32(val) => val,
1092-
v => panic!("Unexpected type in scalar_to_array {:?}", v),
1093-
})
1094-
.collect::<Int32Array>(),
1095-
)),
1096-
DataType::Int64 => Ok(Arc::new(
1097-
iter.into_iter()
1098-
.map(|x| match x {
1099-
ScalarValue::Int64(val) => val,
1100-
v => panic!("Unexpected type in scalar_to_array {:?}", v),
1101-
})
1102-
.collect::<Int64Array>(),
1103-
)),
1104-
DataType::Boolean => Ok(Arc::new(
1105-
iter.into_iter()
1106-
.map(|x| match x {
1107-
ScalarValue::Boolean(val) => val,
1108-
v => panic!("Unexpected type in scalar_to_array {:?}", v),
1109-
})
1110-
.collect::<BooleanArray>(),
1111-
)),
1112-
DataType::UInt8 => Ok(Arc::new(
1113-
iter.into_iter()
1114-
.map(|x| match x {
1115-
ScalarValue::UInt8(val) => val,
1116-
v => panic!("Unexpected type in scalar_to_array {:?}", v),
1117-
})
1118-
.collect::<UInt8Array>(),
1119-
)),
1120-
DataType::UInt16 => Ok(Arc::new(
1121-
iter.into_iter()
1122-
.map(|x| match x {
1123-
ScalarValue::UInt16(val) => val,
1124-
v => panic!("Unexpected type in scalar_to_array {:?}", v),
1125-
})
1126-
.collect::<UInt16Array>(),
1127-
)),
1128-
DataType::UInt32 => Ok(Arc::new(
1129-
iter.into_iter()
1130-
.map(|x| match x {
1131-
ScalarValue::UInt32(val) => val,
1132-
v => panic!("Unexpected type in scalar_to_array {:?}", v),
1133-
})
1134-
.collect::<UInt32Array>(),
1135-
)),
1136-
DataType::UInt64 => Ok(Arc::new(
1137-
iter.into_iter()
1138-
.map(|x| match x {
1139-
ScalarValue::UInt64(val) => val,
1140-
v => panic!("Unexpected type in scalar_to_array {:?}", v),
1141-
})
1142-
.collect::<UInt64Array>(),
1143-
)),
1144-
DataType::Float32 => Ok(Arc::new(
1145-
iter.into_iter()
1146-
.map(|x| match x {
1147-
ScalarValue::Float32(val) => val,
1148-
v => panic!("Unexpected type in scalar_to_array {:?}", v),
1149-
})
1150-
.collect::<Float32Array>(),
1151-
)),
1152-
DataType::Float64 => Ok(Arc::new(
1153-
iter.into_iter()
1154-
.map(|x| match x {
1155-
ScalarValue::Float64(val) => val,
1156-
v => panic!("Unexpected type in scalar_to_array {:?}", v),
1157-
})
1158-
.collect::<Float64Array>(),
1159-
)),
1160-
DataType::Date32 => Ok(Arc::new(
1161-
iter.into_iter()
1162-
.map(|x| match x {
1163-
ScalarValue::Date32(val) => val,
1164-
v => panic!("Unexpected type in scalar_to_array {:?}", v),
1165-
})
1166-
.collect::<Date32Array>(),
1167-
)),
1168-
DataType::Date64 => Ok(Arc::new(
1169-
iter.into_iter()
1170-
.map(|x| match x {
1171-
ScalarValue::Date64(val) => val,
1172-
v => panic!("Unexpected type in scalar_to_array {:?}", v),
1173-
})
1174-
.collect::<Date64Array>(),
1175-
)),
1176-
DataType::Timestamp(TimeUnit::Millisecond, None) => Ok(Arc::new(
1177-
iter.into_iter()
1178-
.map(|x| match x {
1179-
ScalarValue::TimestampMillisecond(val) => val,
1180-
v => panic!("Unexpected type in scalar_to_array {:?}", v),
1181-
})
1182-
.collect::<TimestampMillisecondArray>(),
1183-
)),
1184-
// DataType::Time32(_) => {}
1185-
// DataType::Time64(_) => {}
1186-
// DataType::Duration(_) => {}
1187-
// DataType::Interval(_) => {}
1188-
// DataType::Binary => {}
1189-
// DataType::FixedSizeBinary(_) => {}
1190-
// DataType::LargeBinary => {}
1191-
DataType::Utf8 => Ok(Arc::new(
1192-
iter.into_iter()
1193-
.map(|x| match x {
1194-
ScalarValue::Utf8(val) => val,
1195-
v => panic!("Unexpected type in scalar_to_array {:?}", v),
1196-
})
1197-
.collect::<StringArray>(),
1198-
)),
1199-
DataType::LargeUtf8 => Ok(Arc::new(
1200-
iter.into_iter()
1201-
.map(|x| match x {
1202-
ScalarValue::LargeUtf8(val) => val,
1203-
v => panic!("Unexpected type in scalar_to_array {:?}", v),
1204-
})
1205-
.collect::<LargeStringArray>(),
1206-
)),
1207-
DataType::Dictionary(key_type, val_type) => {
1208-
// Construct array first and then the dictionary array
1209-
// TODO (perf): construct dictionary array in one go
1210-
let x = scalar_to_array(*val_type, iter)?;
1211-
return match *key_type {
1212-
DataType::Int8 => Ok(Arc::new(DictionaryArray::<Int8Type>::from(
1213-
x.data().clone(),
1214-
))),
1215-
DataType::Int16 => Ok(Arc::new(DictionaryArray::<Int16Type>::from(
1216-
x.data().clone(),
1217-
))),
1218-
DataType::Int32 => Ok(Arc::new(DictionaryArray::<Int32Type>::from(
1219-
x.data().clone(),
1220-
))),
1221-
DataType::Int64 => Ok(Arc::new(DictionaryArray::<Int64Type>::from(
1222-
x.data().clone(),
1223-
))),
1224-
DataType::UInt8 => Ok(Arc::new(DictionaryArray::<UInt8Type>::from(
1225-
x.data().clone(),
1226-
))),
1227-
DataType::UInt16 => Ok(Arc::new(DictionaryArray::<UInt16Type>::from(
1228-
x.data().clone(),
1229-
))),
1230-
DataType::UInt32 => Ok(Arc::new(DictionaryArray::<UInt32Type>::from(
1231-
x.data().clone(),
1232-
))),
1233-
DataType::UInt64 => Ok(Arc::new(DictionaryArray::<UInt64Type>::from(
1234-
x.data().clone(),
1235-
))),
1236-
_ => Err(DataFusionError::NotImplemented(format!(
1237-
"Key type {:?} not supported in scalar_to_array",
1238-
key_type
1239-
))),
1240-
};
1241-
}
1242-
1243-
// DataType::FixedSizeList(_, _) => {}
1244-
// DataType::LargeList(_) => {}
1245-
// DataType::Struct(_) => {}
1246-
// DataType::Union(_) => {}
1247-
// DataType::Decimal(_, _) => {}
1248-
data_type => Err(DataFusionError::NotImplemented(format!(
1249-
"to_array not implemented for {}",
1250-
data_type
1251-
))),
1252-
}
1253-
}
1254-
12551048
/// returns a vector of ArrayRefs, where each entry corresponds to either the
12561049
/// final value (mode = Final) or states (mode = Partial)
12571050
fn finalize_aggregation(

0 commit comments

Comments
 (0)