|
| 1 | +/* |
| 2 | + * Copyright 2024-present Alibaba Inc. |
| 3 | + * |
| 4 | + * Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | + * you may not use this file except in compliance with the License. |
| 6 | + * You may obtain a copy of the License at |
| 7 | + * |
| 8 | + * http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | + * |
| 10 | + * Unless required by applicable law or agreed to in writing, software |
| 11 | + * distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | + * See the License for the specific language governing permissions and |
| 14 | + * limitations under the License. |
| 15 | + */ |
| 16 | + |
| 17 | +#include "paimon/common/reader/blob_fallback_batch_reader.h" |
| 18 | + |
| 19 | +#include <algorithm> |
| 20 | +#include <string_view> |
| 21 | +#include <utility> |
| 22 | + |
| 23 | +#include "arrow/api.h" |
| 24 | +#include "arrow/c/bridge.h" |
| 25 | +#include "fmt/format.h" |
| 26 | +#include "paimon/common/data/blob_defs.h" |
| 27 | +#include "paimon/common/data/blob_utils.h" |
| 28 | +#include "paimon/common/metrics/metrics_impl.h" |
| 29 | +#include "paimon/common/reader/reader_utils.h" |
| 30 | +#include "paimon/common/utils/arrow/mem_utils.h" |
| 31 | +#include "paimon/common/utils/arrow/status_utils.h" |
| 32 | + |
| 33 | +namespace paimon { |
| 34 | + |
| 35 | +Result<std::unique_ptr<BlobFallbackBatchReader>> BlobFallbackBatchReader::Create( |
| 36 | + std::vector<std::vector<Segment>>&& sequence_groups, |
| 37 | + const std::shared_ptr<arrow::Schema>& read_schema, int32_t read_batch_size, |
| 38 | + const std::shared_ptr<MemoryPool>& pool) { |
| 39 | + if (sequence_groups.size() < 2) { |
| 40 | + return Status::Invalid( |
| 41 | + "Blob fallback needs at least two sequence groups; a single group should be read " |
| 42 | + "sequentially."); |
| 43 | + } |
| 44 | + if (read_schema == nullptr) { |
| 45 | + return Status::Invalid("Blob fallback read schema cannot be nullptr."); |
| 46 | + } |
| 47 | + if (read_batch_size <= 0) { |
| 48 | + return Status::Invalid(fmt::format( |
| 49 | + "Blob fallback read batch size '{}' should be larger than zero", read_batch_size)); |
| 50 | + } |
| 51 | + int32_t blob_field_idx = -1; |
| 52 | + for (int32_t i = 0; i < read_schema->num_fields(); i++) { |
| 53 | + if (BlobUtils::IsBlobField(read_schema->field(i))) { |
| 54 | + if (blob_field_idx != -1) { |
| 55 | + return Status::Invalid( |
| 56 | + "Blob fallback read schema should contain exactly one blob field."); |
| 57 | + } |
| 58 | + blob_field_idx = i; |
| 59 | + } |
| 60 | + } |
| 61 | + if (blob_field_idx == -1) { |
| 62 | + return Status::Invalid("Blob fallback read schema should contain a blob field."); |
| 63 | + } |
| 64 | + std::vector<GroupCursor> groups; |
| 65 | + groups.reserve(sequence_groups.size()); |
| 66 | + for (auto& segments : sequence_groups) { |
| 67 | + if (segments.empty()) { |
| 68 | + return Status::Invalid("Blob fallback sequence group should not be empty."); |
| 69 | + } |
| 70 | + GroupCursor cursor; |
| 71 | + cursor.segments = std::move(segments); |
| 72 | + groups.push_back(std::move(cursor)); |
| 73 | + } |
| 74 | + return std::unique_ptr<BlobFallbackBatchReader>(new BlobFallbackBatchReader( |
| 75 | + std::move(groups), read_schema, blob_field_idx, read_batch_size, pool)); |
| 76 | +} |
| 77 | + |
| 78 | +BlobFallbackBatchReader::BlobFallbackBatchReader(std::vector<GroupCursor>&& groups, |
| 79 | + const std::shared_ptr<arrow::Schema>& read_schema, |
| 80 | + int32_t blob_field_idx, int32_t read_batch_size, |
| 81 | + const std::shared_ptr<MemoryPool>& pool) |
| 82 | + : groups_(std::move(groups)), |
| 83 | + read_schema_(read_schema), |
| 84 | + blob_field_idx_(blob_field_idx), |
| 85 | + read_batch_size_(read_batch_size), |
| 86 | + arrow_pool_(GetArrowPool(pool)) {} |
| 87 | + |
| 88 | +Result<int64_t> BlobFallbackBatchReader::FillWindow(size_t group_idx, int64_t want, |
| 89 | + std::vector<Chunk>* chunks) { |
| 90 | + GroupCursor& cursor = groups_[group_idx]; |
| 91 | + int64_t collected = 0; |
| 92 | + while (collected < want) { |
| 93 | + if (!cursor.pending.empty()) { |
| 94 | + const std::shared_ptr<arrow::StructArray>& front = cursor.pending.front(); |
| 95 | + int64_t available = front->length() - cursor.pending_pos; |
| 96 | + int64_t take = std::min(available, want - collected); |
| 97 | + chunks->push_back(Chunk{front, cursor.pending_pos, take}); |
| 98 | + cursor.pending_pos += take; |
| 99 | + collected += take; |
| 100 | + if (cursor.pending_pos == front->length()) { |
| 101 | + cursor.pending.pop_front(); |
| 102 | + cursor.pending_pos = 0; |
| 103 | + } |
| 104 | + continue; |
| 105 | + } |
| 106 | + if (cursor.segment_idx >= cursor.segments.size()) { |
| 107 | + // group exhausted; only the first group may define a shorter window |
| 108 | + break; |
| 109 | + } |
| 110 | + Segment& segment = cursor.segments[cursor.segment_idx]; |
| 111 | + if (segment.reader == nullptr) { |
| 112 | + // gap segment: all rows are placeholders |
| 113 | + if (cursor.gap_remaining < 0) { |
| 114 | + cursor.gap_remaining = segment.gap_selected_row_count; |
| 115 | + } |
| 116 | + if (cursor.gap_remaining == 0) { |
| 117 | + cursor.segment_idx++; |
| 118 | + cursor.gap_remaining = -1; |
| 119 | + continue; |
| 120 | + } |
| 121 | + int64_t take = std::min(cursor.gap_remaining, want - collected); |
| 122 | + chunks->push_back(Chunk{nullptr, 0, take}); |
| 123 | + cursor.gap_remaining -= take; |
| 124 | + collected += take; |
| 125 | + if (cursor.gap_remaining == 0) { |
| 126 | + cursor.segment_idx++; |
| 127 | + cursor.gap_remaining = -1; |
| 128 | + } |
| 129 | + continue; |
| 130 | + } |
| 131 | + PAIMON_ASSIGN_OR_RAISE(ReadBatchWithBitmap batch_with_bitmap, |
| 132 | + segment.reader->NextBatchWithBitmap()); |
| 133 | + if (BatchReader::IsEofBatch(batch_with_bitmap)) { |
| 134 | + cursor.segment_idx++; |
| 135 | + continue; |
| 136 | + } |
| 137 | + auto& [read_batch, bitmap] = batch_with_bitmap; |
| 138 | + auto& [c_array, c_schema] = read_batch; |
| 139 | + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> src_array, |
| 140 | + arrow::ImportArray(c_array.get(), c_schema.get())); |
| 141 | + PAIMON_ASSIGN_OR_RAISE(arrow::ArrayVector selected_array_vec, |
| 142 | + ReaderUtils::GenerateFilteredArrayVector(src_array, bitmap)); |
| 143 | + for (const auto& selected_array : selected_array_vec) { |
| 144 | + if (selected_array->length() == 0) { |
| 145 | + continue; |
| 146 | + } |
| 147 | + auto struct_array = std::dynamic_pointer_cast<arrow::StructArray>(selected_array); |
| 148 | + if (struct_array == nullptr) { |
| 149 | + return Status::Invalid("Blob fallback expects file readers to emit struct arrays."); |
| 150 | + } |
| 151 | + cursor.pending.push_back(std::move(struct_array)); |
| 152 | + } |
| 153 | + } |
| 154 | + return collected; |
| 155 | +} |
| 156 | + |
| 157 | +Result<std::vector<bool>> BlobFallbackBatchReader::ComputePlaceholderFlags( |
| 158 | + const std::vector<Chunk>& chunks, int64_t row_count) const { |
| 159 | + std::vector<bool> flags(row_count, false); |
| 160 | + int64_t pos = 0; |
| 161 | + for (const auto& chunk : chunks) { |
| 162 | + if (chunk.array == nullptr) { |
| 163 | + // gap rows stand for placeholders |
| 164 | + std::fill(flags.begin() + pos, flags.begin() + pos + chunk.length, true); |
| 165 | + } else { |
| 166 | + std::shared_ptr<arrow::Array> blob_col = chunk.array->field(blob_field_idx_); |
| 167 | + auto binary_col = std::dynamic_pointer_cast<arrow::LargeBinaryArray>(blob_col); |
| 168 | + if (binary_col == nullptr) { |
| 169 | + return Status::Invalid(fmt::format( |
| 170 | + "Blob fallback expects the blob column to be large binary, but got {}", |
| 171 | + blob_col->type()->ToString())); |
| 172 | + } |
| 173 | + for (int64_t k = 0; k < chunk.length; k++) { |
| 174 | + int64_t idx = chunk.offset + k; |
| 175 | + if (!binary_col->IsNull(idx)) { |
| 176 | + std::string_view value = binary_col->GetView(idx); |
| 177 | + flags[pos + k] = BlobDefs::IsPlaceholderSentinel(value.data(), value.size()); |
| 178 | + } |
| 179 | + } |
| 180 | + } |
| 181 | + pos += chunk.length; |
| 182 | + } |
| 183 | + return flags; |
| 184 | +} |
| 185 | + |
| 186 | +Result<std::shared_ptr<arrow::Array>> BlobFallbackBatchReader::AssembleColumn( |
| 187 | + int32_t field_idx, const std::vector<int32_t>& group_choice, |
| 188 | + const std::vector<std::vector<Chunk>>& group_chunks) const { |
| 189 | + const auto row_count = static_cast<int64_t>(group_choice.size()); |
| 190 | + arrow::ArrayVector pieces; |
| 191 | + int64_t run_start = 0; |
| 192 | + while (run_start < row_count) { |
| 193 | + const int32_t group = group_choice[run_start]; |
| 194 | + int64_t run_end = run_start + 1; |
| 195 | + while (run_end < row_count && group_choice[run_end] == group) { |
| 196 | + run_end++; |
| 197 | + } |
| 198 | + if (group < 0) { |
| 199 | + // placeholder in every layer: no value survives, the row degrades to null |
| 200 | + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( |
| 201 | + std::shared_ptr<arrow::Array> null_piece, |
| 202 | + arrow::MakeArrayOfNull(read_schema_->field(field_idx)->type(), run_end - run_start, |
| 203 | + arrow_pool_.get())); |
| 204 | + pieces.push_back(std::move(null_piece)); |
| 205 | + } else { |
| 206 | + int64_t pos = 0; |
| 207 | + for (const auto& chunk : group_chunks[group]) { |
| 208 | + int64_t overlap_start = std::max(run_start, pos); |
| 209 | + int64_t overlap_end = std::min(run_end, pos + chunk.length); |
| 210 | + if (overlap_start < overlap_end) { |
| 211 | + if (chunk.array == nullptr) { |
| 212 | + return Status::Invalid( |
| 213 | + "Unexpected: a gap row was chosen as a blob fallback result."); |
| 214 | + } |
| 215 | + std::shared_ptr<arrow::Array> column = chunk.array->field(field_idx); |
| 216 | + pieces.push_back(column->Slice(chunk.offset + (overlap_start - pos), |
| 217 | + overlap_end - overlap_start)); |
| 218 | + } |
| 219 | + pos += chunk.length; |
| 220 | + if (pos >= run_end) { |
| 221 | + break; |
| 222 | + } |
| 223 | + } |
| 224 | + } |
| 225 | + run_start = run_end; |
| 226 | + } |
| 227 | + if (pieces.size() == 1 && pieces[0]->offset() == 0) { |
| 228 | + return pieces[0]; |
| 229 | + } |
| 230 | + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> concat_array, |
| 231 | + arrow::Concatenate(pieces, arrow_pool_.get())); |
| 232 | + return concat_array; |
| 233 | +} |
| 234 | + |
| 235 | +Result<BatchReader::ReadBatch> BlobFallbackBatchReader::NextBatch() { |
| 236 | + if (closed_) { |
| 237 | + return Status::Invalid("blob fallback batch reader is closed"); |
| 238 | + } |
| 239 | + std::vector<std::vector<Chunk>> group_chunks(groups_.size()); |
| 240 | + // the first (newest) group defines the window; the others must step in lockstep |
| 241 | + PAIMON_ASSIGN_OR_RAISE(int64_t row_count, FillWindow(0, read_batch_size_, &group_chunks[0])); |
| 242 | + for (size_t g = 1; g < groups_.size(); g++) { |
| 243 | + int64_t expect = row_count == 0 ? 1 : row_count; |
| 244 | + PAIMON_ASSIGN_OR_RAISE(int64_t got, FillWindow(g, expect, &group_chunks[g])); |
| 245 | + int64_t expected = row_count == 0 ? 0 : row_count; |
| 246 | + if (got != expected) { |
| 247 | + return Status::Invalid(fmt::format( |
| 248 | + "All sequence groups of a blob fallback read should have the same number of " |
| 249 | + "rows: group {} yielded {} rows in a window of {}", |
| 250 | + g, got, expected)); |
| 251 | + } |
| 252 | + } |
| 253 | + if (row_count == 0) { |
| 254 | + return BatchReader::MakeEofBatch(); |
| 255 | + } |
| 256 | + |
| 257 | + std::vector<std::vector<bool>> placeholder_flags(groups_.size()); |
| 258 | + for (size_t g = 0; g < groups_.size(); g++) { |
| 259 | + PAIMON_ASSIGN_OR_RAISE(placeholder_flags[g], |
| 260 | + ComputePlaceholderFlags(group_chunks[g], row_count)); |
| 261 | + } |
| 262 | + // per row, the first group in max-sequence order with a real entry wins; -1 means the row is |
| 263 | + // a placeholder in every group |
| 264 | + std::vector<int32_t> group_choice(row_count, -1); |
| 265 | + for (int64_t r = 0; r < row_count; r++) { |
| 266 | + for (size_t g = 0; g < groups_.size(); g++) { |
| 267 | + if (!placeholder_flags[g][r]) { |
| 268 | + group_choice[r] = static_cast<int32_t>(g); |
| 269 | + break; |
| 270 | + } |
| 271 | + } |
| 272 | + } |
| 273 | + |
| 274 | + arrow::ArrayVector columns; |
| 275 | + columns.reserve(read_schema_->num_fields()); |
| 276 | + for (int32_t field_idx = 0; field_idx < read_schema_->num_fields(); field_idx++) { |
| 277 | + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> column, |
| 278 | + AssembleColumn(field_idx, group_choice, group_chunks)); |
| 279 | + columns.push_back(std::move(column)); |
| 280 | + } |
| 281 | + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::StructArray> target_array, |
| 282 | + arrow::StructArray::Make(columns, read_schema_->fields())); |
| 283 | + std::unique_ptr<ArrowArray> c_array = std::make_unique<ArrowArray>(); |
| 284 | + std::unique_ptr<ArrowSchema> c_schema = std::make_unique<ArrowSchema>(); |
| 285 | + PAIMON_RETURN_NOT_OK_FROM_ARROW( |
| 286 | + arrow::ExportArray(*target_array, c_array.get(), c_schema.get())); |
| 287 | + return std::make_pair(std::move(c_array), std::move(c_schema)); |
| 288 | +} |
| 289 | + |
| 290 | +Result<BatchReader::ReadBatchWithBitmap> BlobFallbackBatchReader::NextBatchWithBitmap() { |
| 291 | + PAIMON_ASSIGN_OR_RAISE(ReadBatch batch, NextBatch()); |
| 292 | + if (BatchReader::IsEofBatch(batch)) { |
| 293 | + return BatchReader::MakeEofBatchWithBitmap(); |
| 294 | + } |
| 295 | + return ReaderUtils::AddAllValidBitmap(std::move(batch)); |
| 296 | +} |
| 297 | + |
| 298 | +void BlobFallbackBatchReader::Close() { |
| 299 | + for (auto& group : groups_) { |
| 300 | + group.pending.clear(); |
| 301 | + for (auto& segment : group.segments) { |
| 302 | + if (segment.reader) { |
| 303 | + segment.reader->Close(); |
| 304 | + } |
| 305 | + } |
| 306 | + } |
| 307 | + closed_ = true; |
| 308 | +} |
| 309 | + |
| 310 | +std::shared_ptr<Metrics> BlobFallbackBatchReader::GetReaderMetrics() const { |
| 311 | + auto metrics = std::make_shared<MetricsImpl>(); |
| 312 | + for (const auto& group : groups_) { |
| 313 | + for (const auto& segment : group.segments) { |
| 314 | + if (segment.reader) { |
| 315 | + auto reader_metrics = segment.reader->GetReaderMetrics(); |
| 316 | + if (reader_metrics) { |
| 317 | + metrics->Merge(reader_metrics); |
| 318 | + } |
| 319 | + } |
| 320 | + } |
| 321 | + } |
| 322 | + return metrics; |
| 323 | +} |
| 324 | + |
| 325 | +} // namespace paimon |
0 commit comments