| 1201 | mutable_, "Can only create RowPartitions once from a row container"); |
| 1202 | mutable_ = false; |
| 1203 | return std::make_unique<RowPartitions>(numRows_, pool); |
| 1204 | } |
| 1205 | |
| 1206 | int32_t RowContainer::listPartitionRows( |
| 1207 | RowContainerIterator& iter, |
| 1208 | uint8_t partition, |
| 1209 | int32_t maxRows, |
| 1210 | const RowPartitions& rowPartitions, |
| 1211 | char** result) { |
| 1212 | BOLT_CHECK( |
| 1213 | !mutable_, "Can't list partition rows from a mutable row container"); |
| 1214 | BOLT_CHECK_EQ( |
| 1215 | rowPartitions.size(), numRows_, "All rows must have a partition"); |
| 1216 | if (numRows_ == 0) { |
| 1217 | return 0; |
| 1218 | } |
| 1219 | const auto partitionNumberVector = |
| 1220 | xsimd::batch<uint8_t>::broadcast(partition); |
| 1221 | const auto& allocation = rowPartitions.allocation(); |
| 1222 | int32_t numResults = 0; |
| 1223 | while (numResults < maxRows && iter.rowNumber < numRows_) { |
| 1224 | constexpr int32_t kBatch = xsimd::batch<uint8_t>::size; |
| 1225 | // Start at multiple of kBatch. |
| 1226 | auto startRow = iter.rowNumber / kBatch * kBatch; |
| 1227 | // Ignore the possible hits at or below iter.rowNumber. |
| 1228 | uint32_t firstBatchMask = ~bits::lowMask(iter.rowNumber - startRow); |
| 1229 | int32_t runIndex; |
| 1230 | int32_t offsetInRun; |
| 1231 | BOLT_CHECK_LT(startRow, numRows_); |
| 1232 | allocation.findRun(startRow, &runIndex, &offsetInRun); |
| 1233 | auto run = allocation.runAt(runIndex); |
| 1234 | auto runEnd = run.numBytes(); |
| 1235 | auto runBytes = run.data<uint8_t>(); |
| 1236 | for (; offsetInRun < runEnd; offsetInRun += kBatch) { |
| 1237 | auto bits = |
| 1238 | simd::toBitMask( |
| 1239 | partitionNumberVector == |
| 1240 | xsimd::batch<uint8_t>::load_unaligned(runBytes + offsetInRun)) & |
| 1241 | firstBatchMask; |
| 1242 | firstBatchMask = ~0; |
| 1243 | bool atEnd = false; |
| 1244 | if (startRow + kBatch >= numRows_) { |
| 1245 | // Clear bits that are for rows past numRows_ - 1. |
| 1246 | bits &= bits::lowMask(numRows_ - startRow); |
| 1247 | atEnd = true; |
| 1248 | } |
| 1249 | while (bits) { |
| 1250 | const int32_t hit = __builtin_ctz(bits); |
| 1251 | const auto distance = hit + startRow - iter.rowNumber; |
| 1252 | skip(iter, distance); |
| 1253 | result[numResults++] = iter.currentRow(); |
| 1254 | if (numResults == maxRows) { |
| 1255 | skip(iter, 1); |
| 1256 | return numResults; |
| 1257 | } |
| 1258 | // Clear last set bit in 'bits'. |
| 1259 | bits &= bits - 1; |
| 1260 | } |