mutates by copying from_key into to_key and changing from_key to zero
| 177 | |
| 178 | // mutates by copying from_key into to_key and changing from_key to zero |
| 179 | Result<BatchesWithSchema> MutateByKey(BatchesWithSchema& batches, std::string from_key, |
| 180 | std::string to_key, bool replace_key = false, |
| 181 | bool null_key = false, bool remove_key = false) { |
| 182 | int from_index = batches.schema->GetFieldIndex(from_key); |
| 183 | int n_fields = batches.schema->num_fields(); |
| 184 | auto fields = batches.schema->fields(); |
| 185 | BatchesWithSchema new_batches; |
| 186 | if (remove_key) { |
| 187 | ARROW_ASSIGN_OR_RAISE(new_batches.schema, batches.schema->RemoveField(from_index)); |
| 188 | } else { |
| 189 | auto new_field = batches.schema->field(from_index)->WithName(to_key); |
| 190 | ARROW_ASSIGN_OR_RAISE(new_batches.schema, |
| 191 | replace_key ? batches.schema->SetField(from_index, new_field) |
| 192 | : batches.schema->AddField(from_index, new_field)); |
| 193 | } |
| 194 | size_t batch_index = 0; |
| 195 | for (const ExecBatch& batch : batches.batches) { |
| 196 | std::vector<Datum> new_values; |
| 197 | for (int i = 0; i < n_fields; i++) { |
| 198 | const Datum& value = batch.values[i]; |
| 199 | if (i == from_index) { |
| 200 | if (remove_key) { |
| 201 | continue; |
| 202 | } |
| 203 | auto type = fields[i]->type(); |
| 204 | if (null_key) { |
| 205 | std::shared_ptr<Array> empty; |
| 206 | BuildNullArray(empty, type, batch.length); |
| 207 | new_values.push_back(empty); |
| 208 | } else if (is_primitive(type->id())) { |
| 209 | std::shared_ptr<Array> empty; |
| 210 | BuildZeroPrimitiveArray(empty, type, batch.length); |
| 211 | new_values.push_back(empty); |
| 212 | } else if (is_base_binary_like(type->id())) { |
| 213 | std::shared_ptr<Array> empty; |
| 214 | switch (type->id()) { |
| 215 | case Type::STRING: |
| 216 | BuildZeroBaseBinaryArray<StringBuilder>(empty, batch.length); |
| 217 | break; |
| 218 | case Type::LARGE_STRING: |
| 219 | BuildZeroBaseBinaryArray<LargeStringBuilder>(empty, batch.length); |
| 220 | break; |
| 221 | case Type::BINARY: |
| 222 | BuildZeroBaseBinaryArray<BinaryBuilder>(empty, batch.length); |
| 223 | break; |
| 224 | case Type::LARGE_BINARY: |
| 225 | BuildZeroBaseBinaryArray<LargeBinaryBuilder>(empty, batch.length); |
| 226 | break; |
| 227 | default: |
| 228 | DCHECK(false); |
| 229 | break; |
| 230 | } |
| 231 | new_values.push_back(empty); |
| 232 | } else { |
| 233 | ARROW_ASSIGN_OR_RAISE(auto sub, Subtract(value, value)); |
| 234 | new_values.push_back(sub); |
| 235 | } |
| 236 | if (replace_key) { |
no test coverage detected