| 962 | } |
| 963 | |
| 964 | ExtensionIdRegistry::SubstraitAggregateToArrow DecodeBasicAggregate( |
| 965 | const std::string& arrow_function_name) { |
| 966 | return [arrow_function_name](const SubstraitCall& call) -> Result<compute::Aggregate> { |
| 967 | std::string fixed_arrow_func; |
| 968 | if (call.is_hash()) { |
| 969 | fixed_arrow_func = "hash_"; |
| 970 | } |
| 971 | |
| 972 | switch (call.size()) { |
| 973 | case 0: { |
| 974 | if (call.id().name == "count") { |
| 975 | fixed_arrow_func += "count_all"; |
| 976 | return compute::Aggregate{std::move(fixed_arrow_func), ""}; |
| 977 | } |
| 978 | return Status::Invalid("Expected aggregate call ", call.id().uri, "#", |
| 979 | call.id().name, " to have at least one argument"); |
| 980 | } |
| 981 | default: { |
| 982 | // Handles all arity > 0 |
| 983 | |
| 984 | std::shared_ptr<compute::FunctionOptions> options = nullptr; |
| 985 | if (arrow_function_name == "stddev" || arrow_function_name == "variance") { |
| 986 | // See the following URL for the spec of stddev and variance: |
| 987 | // https://github.com/substrait-io/substrait/blob/ |
| 988 | // 73228b4112d79eb1011af0ebb41753ce23ca180c/ |
| 989 | // extensions/functions_arithmetic.yaml#L1240 |
| 990 | auto maybe_dist = call.GetOption("distribution"); |
| 991 | if (maybe_dist) { |
| 992 | auto& prefs = **maybe_dist; |
| 993 | if (prefs.size() != 1) { |
| 994 | return Status::Invalid("expected a single preference for ", |
| 995 | arrow_function_name, " but got ", prefs.size()); |
| 996 | } |
| 997 | int ddof; |
| 998 | if (prefs[0] == "POPULATION") { |
| 999 | ddof = 1; |
| 1000 | } else if (prefs[0] == "SAMPLE") { |
| 1001 | ddof = 0; |
| 1002 | } else { |
| 1003 | return Status::Invalid("unknown distribution preference ", prefs[0]); |
| 1004 | } |
| 1005 | options = std::make_shared<compute::VarianceOptions>(ddof); |
| 1006 | } |
| 1007 | } |
| 1008 | fixed_arrow_func += arrow_function_name; |
| 1009 | |
| 1010 | std::vector<FieldRef> target; |
| 1011 | for (int i = 0; i < call.size(); i++) { |
| 1012 | ARROW_ASSIGN_OR_RAISE(compute::Expression arg, call.GetValueArg(i)); |
| 1013 | const FieldRef* arg_ref = arg.field_ref(); |
| 1014 | if (!arg_ref) { |
| 1015 | return Status::Invalid("Expected an aggregate call ", call.id().uri, "#", |
| 1016 | call.id().name, " to have a direct reference"); |
| 1017 | } |
| 1018 | // Copy arg_ref here because field_ref() return const FieldRef* |
| 1019 | target.emplace_back(*arg_ref); |
| 1020 | } |
| 1021 | return compute::Aggregate{std::move(fixed_arrow_func), |