| 1111 | } |
| 1112 | |
| 1113 | Status KrpcDataStreamSender::Send(RuntimeState* state, RowBatch* batch) { |
| 1114 | SCOPED_TIMER(profile()->total_time_counter()); |
| 1115 | DCHECK(!closed_); |
| 1116 | DCHECK(!flushed_); |
| 1117 | |
| 1118 | if (batch->num_rows() == 0) return Status::OK(); |
| 1119 | if (partition_type_ == TPartitionType::UNPARTITIONED) { |
| 1120 | // Only skip compression if there is a single channel and destination is in the same |
| 1121 | // process. TODO: could be optimized to send the uncompressed buffer to the local |
| 1122 | // targets to avoid decompression cost at the receiver. |
| 1123 | bool is_local = channels_.size() == 1 && channels_[0]->IsLocal(); |
| 1124 | RETURN_IF_ERROR(SerializeBatch( |
| 1125 | batch, serialization_batch_.get(), !is_local, channels_.size())); |
| 1126 | // TransmitData() will block if there are still in-flight rpcs (and those will |
| 1127 | // reference the previously written in_flight_batch_). |
| 1128 | for (int i = 0; i < channels_.size(); ++i) { |
| 1129 | // Do not swap serialization_batch_ with the channel's outbound_batch_ to allow |
| 1130 | // multiple channels to use the data backing serialization_batch_ in parallel. |
| 1131 | RETURN_IF_ERROR( |
| 1132 | channels_[i]->TransmitData(&serialization_batch_, false /*swap_batch*/)); |
| 1133 | } |
| 1134 | // At this point no RPCs can still refer to the old in_flight_batch_. |
| 1135 | in_flight_batch_.swap(serialization_batch_); |
| 1136 | } else if (partition_type_ == TPartitionType::RANDOM || |
| 1137 | (channels_.size() == 1 && partition_type_ != TPartitionType::DIRECTED)) { |
| 1138 | // Round-robin batches among channels. Wait for the current channel to finish its |
| 1139 | // rpc before overwriting its batch. |
| 1140 | Channel* current_channel = channels_[current_channel_idx_].get(); |
| 1141 | RETURN_IF_ERROR(current_channel->SerializeAndSendBatch(batch)); |
| 1142 | current_channel_idx_ = (current_channel_idx_ + 1) % channels_.size(); |
| 1143 | } else if (partition_type_ == TPartitionType::KUDU) { |
| 1144 | DCHECK_EQ(partition_expr_evals_.size(), 1); |
| 1145 | int num_channels = channels_.size(); |
| 1146 | const int num_rows = batch->num_rows(); |
| 1147 | const int hash_batch_size = RowBatch::HASH_BATCH_SIZE; |
| 1148 | int channel_ids[hash_batch_size]; |
| 1149 | for (int batch_start = 0; batch_start < num_rows; batch_start += hash_batch_size) { |
| 1150 | int batch_window_size = min(num_rows - batch_start, hash_batch_size); |
| 1151 | for (int i = 0; i < batch_window_size; ++i) { |
| 1152 | TupleRow* row = batch->GetRow(i + batch_start); |
| 1153 | int32_t partition = |
| 1154 | *reinterpret_cast<int32_t*>(partition_expr_evals_[0]->GetValue(row)); |
| 1155 | if (partition < 0) { |
| 1156 | // This row doesn't correspond to a partition, |
| 1157 | // e.g. it's outside the given ranges. |
| 1158 | partition = next_unknown_partition_; |
| 1159 | ++next_unknown_partition_; |
| 1160 | } |
| 1161 | channel_ids[i] = partition % num_channels; |
| 1162 | } |
| 1163 | |
| 1164 | for (int i = 0; i < batch_window_size; ++i) { |
| 1165 | TupleRow* row = batch->GetRow(i + batch_start); |
| 1166 | int channel_id = channel_ids[i]; |
| 1167 | PartitionRowCollector& collector = partition_row_collectors_[channel_id]; |
| 1168 | RETURN_IF_ERROR(collector.AppendRow(row, row_desc_)); |
| 1169 | } |
| 1170 | } |