diff --git a/vortex-duckdb/cpp/include/table_function.hpp b/vortex-duckdb/cpp/include/table_function.hpp index 46b88e1c7ca..fc989353abc 100644 --- a/vortex-duckdb/cpp/include/table_function.hpp +++ b/vortex-duckdb/cpp/include/table_function.hpp @@ -64,6 +64,18 @@ struct VortexBindData final : FunctionData { } }; + struct DistributedFragment { + idx_t file_index; + idx_t row_start; + idx_t row_end; + idx_t estimated_bytes; + + bool operator==(const DistributedFragment &other) const { + return file_index == other.file_index && row_start == other.row_start && + row_end == other.row_end && estimated_bytes == other.estimated_bytes; + } + }; + struct PortableSnapshot { string portable_bind; vector distributed_files; @@ -78,7 +90,7 @@ struct VortexBindData final : FunctionData { bool explicit_split_mode = false; bool splits_applied = false; vector eligible_file_indexes; - vector assigned_file_indexes; + vector assigned_fragments; #endif }; diff --git a/vortex-duckdb/cpp/table_function.cpp b/vortex-duckdb/cpp/table_function.cpp index 022a86ea145..c5514b81552 100644 --- a/vortex-duckdb/cpp/table_function.cpp +++ b/vortex-duckdb/cpp/table_function.cpp @@ -51,7 +51,7 @@ unique_ptr VortexBindData::Copy() const { result->explicit_split_mode = explicit_split_mode; result->splits_applied = splits_applied; result->eligible_file_indexes = eligible_file_indexes; - result->assigned_file_indexes = assigned_file_indexes; + result->assigned_fragments = assigned_fragments; return result; #else const auto copied_ffi_data = duckdb_table_function_bind_data_clone(ffi_data->DataPtr()); @@ -71,7 +71,7 @@ bool VortexBindData::Equals(const FunctionData &other_base) const { portable_bind == other.portable_bind && distributed_files == other.distributed_files && aggregate_scan == other.aggregate_scan && explicit_split_mode == other.explicit_split_mode && splits_applied == other.splits_applied && eligible_file_indexes == other.eligible_file_indexes && - assigned_file_indexes == other.assigned_file_indexes; + assigned_fragments == other.assigned_fragments; #else // if "types" are different, "ffi_data" would also be different as it // contains types inside, so omit "types" from comparison. @@ -132,6 +132,57 @@ VortexBindData::PortableSnapshot VortexBindData::CreatePortableSnapshot() const result.distributed_files = std::move(files); return result; } + +static vector PlanVortexFragments(const string &portable_bind, + const vector &file_indexes, + idx_t target_fragment_count) { + if (!std::is_sorted(file_indexes.begin(), file_indexes.end()) || + std::adjacent_find(file_indexes.begin(), file_indexes.end()) != file_indexes.end()) { + throw InvalidInputException("Vortex fragment file indexes are not in canonical order"); + } + duckdb_vx_error error_out = nullptr; + auto fragment_plan_data = duckdb_table_function_distributed_plan_fragments( + reinterpret_cast(portable_bind.data()), + portable_bind.size(), + file_indexes.data(), + file_indexes.size(), + target_fragment_count, + &error_out); + if (error_out) { + throw InvalidInputException(IntoErrString(error_out)); + } + if (!fragment_plan_data) { + throw InvalidInputException("Vortex failed to plan distributed scan fragments"); + } + auto fragment_plan = unique_ptr(reinterpret_cast(fragment_plan_data)); + const auto fragment_count = duckdb_table_function_distributed_fragment_count(fragment_plan->DataPtr()); + vector fragments; + fragments.reserve(fragment_count); + for (idx_t fragment_index = 0; fragment_index < fragment_count; fragment_index++) { + VortexDistributedFragmentView view {}; + if (!duckdb_table_function_distributed_fragment_at(fragment_plan->DataPtr(), fragment_index, &view) || + !std::binary_search(file_indexes.begin(), file_indexes.end(), view.file_index) || + view.row_start > view.row_end || view.estimated_bytes == DConstants::INVALID_INDEX) { + throw InvalidInputException("Vortex produced an invalid distributed fragment at index %llu", + static_cast(fragment_index)); + } + VortexBindData::DistributedFragment fragment {view.file_index, + view.row_start, + view.row_end, + view.estimated_bytes}; + if (!fragments.empty()) { + const auto &previous = fragments.back(); + if (previous.file_index > fragment.file_index || + (previous.file_index == fragment.file_index && + (previous.row_start >= fragment.row_start || previous.row_end > fragment.row_start))) { + throw InvalidInputException("Vortex produced fragments outside canonical order at index %llu", + static_cast(fragment_index)); + } + } + fragments.push_back(fragment); + } + return fragments; +} #endif // This is a flaw of Duckdb API which doesn't allow passing non-const @@ -472,14 +523,23 @@ unique_ptr init_global(ClientContext &context, TableFu bool distributed = false; if (!bind_data.ffi_data) { const auto snapshot = bind_data.CreatePortableSnapshot(); - vector native_file_indexes; - const vector *runtime_file_indexes = &bind_data.assigned_file_indexes; + vector native_fragments; + const vector *runtime_fragments = &bind_data.assigned_fragments; if (!bind_data.explicit_split_mode) { + vector native_file_indexes; native_file_indexes.reserve(snapshot.distributed_files.size()); for (idx_t file_index = 0; file_index < snapshot.distributed_files.size(); file_index++) { native_file_indexes.push_back(file_index); } - runtime_file_indexes = &native_file_indexes; + native_fragments = + PlanVortexFragments(snapshot.portable_bind, native_file_indexes, native_file_indexes.size()); + runtime_fragments = &native_fragments; + } + vector runtime_fragment_views; + runtime_fragment_views.reserve(runtime_fragments->size()); + for (const auto &fragment : *runtime_fragments) { + runtime_fragment_views.push_back( + {fragment.file_index, fragment.row_start, fragment.row_end, fragment.estimated_bytes}); } // Optional filters (for example TopN's dynamic bound) are maintained // by an upstream operator that is absent from Vane's detached scan @@ -488,8 +548,8 @@ unique_ptr init_global(ClientContext &context, TableFu ffi_global_data = duckdb_table_function_init_global_distributed( reinterpret_cast(snapshot.portable_bind.data()), snapshot.portable_bind.size(), - runtime_file_indexes->data(), - runtime_file_indexes->size(), + runtime_fragment_views.data(), + runtime_fragment_views.size(), bind_data.explicit_split_mode, &ffi_input, &error_out); @@ -508,7 +568,7 @@ unique_ptr init_global(ClientContext &context, TableFu #ifdef VORTEX_VANE_DISTRIBUTED bool force_empty_output = false; force_empty_output = distributed && bind_data.explicit_split_mode && bind_data.splits_applied && - bind_data.assigned_file_indexes.empty(); + bind_data.assigned_fragments.empty(); return make_uniq(std::move(cdata), distributed, force_empty_output); #else return make_uniq(std::move(cdata)); @@ -671,9 +731,11 @@ InsertionOrderPreservingMap to_string(TableFunctionToStringInput &input) auto &bind_data = input.bind_data->Cast(); if (!bind_data.ffi_data) { result.insert("Function", "Vortex Scan"); - const auto file_count = bind_data.explicit_split_mode ? bind_data.assigned_file_indexes.size() - : bind_data.distributed_files.size(); - result.insert("Distributed files", std::to_string(file_count)); + if (bind_data.explicit_split_mode) { + result.insert("Assigned fragments", std::to_string(bind_data.assigned_fragments.size())); + } else { + result.insert("Distributed files", std::to_string(bind_data.distributed_files.size())); + } return result; } #endif @@ -688,7 +750,8 @@ namespace { static constexpr uint8_t VORTEX_SPLIT_PAYLOAD_VERSION = 1; static constexpr uint8_t VORTEX_BIND_SERDE_VERSION = 1; -static constexpr const char *VORTEX_SPLIT_CODEC = "vane.vortex-file-split"; +static constexpr idx_t VORTEX_MIN_FRAGMENT_PAYLOAD_BYTES = sizeof(uint64_t) * 7; +static constexpr const char *VORTEX_SPLIT_CODEC = "vane.vortex-file-fragment-split"; static bool IsCanonicalVortexScanId(const string &scan_split_set_id) { hugeint_t parsed; @@ -730,35 +793,44 @@ static void AppendSplitString(string &result, const string &value) { result.append(value); } -// Binary payload v1: -// "VXSP" | u8 version | string scan_split_set_id | u64 file_count | +// Binary fragment payload v1: +// "VXFR" | u8 version | string scan_split_set_id | u64 fragment_count | // repeated(u64 stable_file_index | string source_url | string path | -// u64 immutable_size) -// Normal scans encode one file. Aggregate-pushed scans encode their complete -// pruned file set so a single worker computes the final aggregate exactly once. -static string EncodeVortexSplit(const vector &file_indexes, +// u64 immutable_size | u64 row_start | u64 row_end | u64 estimated_bytes) +// Normal scans encode one independently assignable fragment. Aggregate-pushed +// scans encode one full-file fragment per selected file in one complete-set split. +static string EncodeVortexSplit(const vector &fragments, const vector &distributed_files, const string &scan_split_set_id) { - if (file_indexes.empty()) { + if (fragments.empty()) { throw InternalException("Cannot encode an empty distributed Vortex split"); } if (!IsCanonicalVortexScanId(scan_split_set_id)) { throw InternalException("Cannot encode a distributed Vortex split without a canonical scan identity"); } - string result("VXSP", 4); + string result("VXFR", 4); AppendSplitByte(result, VORTEX_SPLIT_PAYLOAD_VERSION); AppendSplitString(result, scan_split_set_id); - AppendSplitU64(result, file_indexes.size()); - for (auto file_index : file_indexes) { - if (file_index >= distributed_files.size()) { + AppendSplitU64(result, fragments.size()); + for (const auto &fragment : fragments) { + if (fragment.file_index >= distributed_files.size()) { throw InternalException("Cannot encode unknown distributed Vortex file index %llu", - static_cast(file_index)); + static_cast(fragment.file_index)); + } + if (fragment.row_start > fragment.row_end || fragment.estimated_bytes == DConstants::INVALID_INDEX) { + throw InternalException("Cannot encode an invalid distributed Vortex fragment range or estimate"); + } + const auto &file = distributed_files[fragment.file_index]; + if (fragment.estimated_bytes > file.size) { + throw InternalException("Cannot encode a distributed Vortex fragment larger than its file"); } - const auto &file = distributed_files[file_index]; - AppendSplitU64(result, file_index); + AppendSplitU64(result, fragment.file_index); AppendSplitString(result, file.source_url); AppendSplitString(result, file.path); AppendSplitU64(result, file.size); + AppendSplitU64(result, fragment.row_start); + AppendSplitU64(result, fragment.row_end); + AppendSplitU64(result, fragment.estimated_bytes); } return result; } @@ -799,23 +871,30 @@ class VortexSplitDecoder { } } + idx_t RemainingBytes() const { + return payload.size() - offset; + } + private: const string &payload; idx_t offset = 0; }; -struct DecodedVortexFile { +struct DecodedVortexFragment { idx_t file_index; VortexBindData::DistributedFile file; + idx_t row_start; + idx_t row_end; + idx_t estimated_bytes; }; struct DecodedVortexSplit { string scan_split_set_id; - vector files; + vector fragments; }; static DecodedVortexSplit DecodeVortexSplit(const string &payload) { - if (payload.size() < 5 || payload.compare(0, 4, "VXSP") != 0) { + if (payload.size() < 5 || payload.compare(0, 4, "VXFR") != 0) { throw InvalidInputException("Invalid distributed Vortex split payload magic"); } VortexSplitDecoder decoder(payload); @@ -831,75 +910,48 @@ static DecodedVortexSplit DecodeVortexSplit(const string &payload) { if (!IsCanonicalVortexScanId(result.scan_split_set_id)) { throw InvalidInputException("Distributed Vortex split contains an invalid scan identity"); } - auto file_count = decoder.ReadU64(); - if (file_count == 0 || file_count > payload.size()) { - throw InvalidInputException("Invalid file count in distributed Vortex split payload"); + auto fragment_count = decoder.ReadU64(); + if (fragment_count == 0 || + fragment_count > decoder.RemainingBytes() / VORTEX_MIN_FRAGMENT_PAYLOAD_BYTES) { + throw InvalidInputException("Invalid fragment count in distributed Vortex split payload"); } - result.files.reserve(file_count); - for (idx_t file_offset = 0; file_offset < file_count; file_offset++) { - DecodedVortexFile decoded; + result.fragments.reserve(fragment_count); + for (idx_t fragment_offset = 0; fragment_offset < fragment_count; fragment_offset++) { + DecodedVortexFragment decoded; decoded.file_index = decoder.ReadU64(); decoded.file.source_url = decoder.ReadString(); decoded.file.path = decoder.ReadString(); decoded.file.size = decoder.ReadU64(); + decoded.row_start = decoder.ReadU64(); + decoded.row_end = decoder.ReadU64(); + decoded.estimated_bytes = decoder.ReadU64(); if (decoded.file.size == DConstants::INVALID_INDEX) { throw InvalidInputException("Distributed Vortex split contains an invalid file size"); } + if (decoded.row_start > decoded.row_end || decoded.estimated_bytes == DConstants::INVALID_INDEX || + decoded.estimated_bytes > decoded.file.size) { + throw InvalidInputException("Distributed Vortex split contains an invalid fragment range"); + } if (decoded.file.source_url.empty() || !IsCanonicalVortexFilePath(decoded.file.path)) { throw InvalidInputException("Distributed Vortex split contains an invalid file identity"); } - result.files.push_back(std::move(decoded)); + result.fragments.push_back(std::move(decoded)); } decoder.Finish(); return result; } -static bool IsCanonicalVortexSplitId(const string &split_id) { - if (split_id.empty()) { - return false; - } - idx_t segment_start = 0; - optional_idx previous; - while (segment_start < split_id.size()) { - auto segment_end = split_id.find(',', segment_start); - if (segment_end == string::npos) { - segment_end = split_id.size(); - } - if (segment_end == segment_start || - (segment_end - segment_start > 1 && split_id[segment_start] == '0')) { - return false; - } - idx_t value = 0; - for (idx_t offset = segment_start; offset < segment_end; offset++) { - auto character = split_id[offset]; - if (character < '0' || character > '9') { - return false; - } - auto digit = static_cast(character - '0'); - if (value > (NumericLimits::Maximum() - digit) / 10) { - return false; - } - value = value * 10 + digit; - } - if (previous.IsValid() && previous.GetIndex() >= value) { - return false; - } - previous = optional_idx(value); - if (segment_end == split_id.size()) { - return true; - } - segment_start = segment_end + 1; - } - return false; -} - -static string CanonicalVortexSplitId(const vector &file_indexes) { +static string CanonicalVortexSplitId(const vector &fragments) { string result; - for (auto file_index : file_indexes) { + for (const auto &fragment : fragments) { if (!result.empty()) { result += ','; } - result += std::to_string(file_index); + result += std::to_string(fragment.file_index); + result += ':'; + result += std::to_string(fragment.row_start); + result += '-'; + result += std::to_string(fragment.row_end); } return result; } @@ -910,30 +962,31 @@ static idx_t SaturatingVortexSplitEstimate(idx_t left, idx_t right) { return right > maximum - left ? maximum : left + right; } -static idx_t ProportionalVortexSplitEstimate(idx_t total, idx_t numerator, idx_t denominator) { - D_ASSERT(denominator > 0 && numerator <= denominator); - if (numerator == 0) { - return 0; - } - if (numerator == denominator) { - return total; - } - const auto scaled = static_cast(total) * static_cast(numerator) / - static_cast(denominator); - // On platforms where long double is IEEE double, UINT64_MAX - 1 rounds to - // 2^64. Clamp in floating point before the integer conversion so an - // extreme estimate cannot invoke an out-of-range conversion. - if (scaled >= static_cast(total)) { - return total; - } - return static_cast(scaled); -} - static bool SameDistributedFile(const VortexBindData::DistributedFile &left, const VortexBindData::DistributedFile &right) { return left.source_url == right.source_url && left.path == right.path && left.size == right.size; } +static bool IsCompleteAggregateVortexAssignment(const vector &fragments, + const vector &eligible_file_indexes, + const vector &files) { + if (fragments.size() != eligible_file_indexes.size()) { + return false; + } + for (idx_t fragment_index = 0; fragment_index < fragments.size(); fragment_index++) { + const auto file_index = eligible_file_indexes[fragment_index]; + const auto &fragment = fragments[fragment_index]; + if (file_index >= files.size() || fragment.file_index != file_index || fragment.row_start != 0 || + fragment.estimated_bytes != files[file_index].size) { + return false; + } + } + // The worker checks row_end against the immutable file's actual row count when it opens the + // reader. Keeping that storage-dependent check there avoids reopening every aggregate file + // while applying or deserializing owned split state. + return true; +} + static void ValidatePortableVortexBind(const string &portable_bind, const vector &files, bool aggregate_scan, @@ -1022,9 +1075,26 @@ static void VortexScanSerialize(Serializer &serializer, serializer.WriteProperty(106, "scan_split_set_id", data.scan_split_set_id); serializer.WriteProperty(107, "explicit_split_mode", data.explicit_split_mode); serializer.WriteProperty(108, "splits_applied", data.splits_applied); - serializer.WriteProperty(109, "assigned_file_indexes", data.assigned_file_indexes); - serializer.WriteProperty(110, "aggregate_scan", snapshot.aggregate_scan); - serializer.WriteProperty(111, "eligible_file_indexes", data.eligible_file_indexes); + vector assigned_file_indexes; + vector assigned_row_starts; + vector assigned_row_ends; + vector assigned_estimated_bytes; + assigned_file_indexes.reserve(data.assigned_fragments.size()); + assigned_row_starts.reserve(data.assigned_fragments.size()); + assigned_row_ends.reserve(data.assigned_fragments.size()); + assigned_estimated_bytes.reserve(data.assigned_fragments.size()); + for (const auto &fragment : data.assigned_fragments) { + assigned_file_indexes.push_back(fragment.file_index); + assigned_row_starts.push_back(fragment.row_start); + assigned_row_ends.push_back(fragment.row_end); + assigned_estimated_bytes.push_back(fragment.estimated_bytes); + } + serializer.WriteProperty(109, "assigned_file_indexes", assigned_file_indexes); + serializer.WriteProperty(110, "assigned_row_starts", assigned_row_starts); + serializer.WriteProperty(111, "assigned_row_ends", assigned_row_ends); + serializer.WriteProperty(112, "assigned_estimated_bytes", assigned_estimated_bytes); + serializer.WriteProperty(113, "aggregate_scan", snapshot.aggregate_scan); + serializer.WriteProperty(114, "eligible_file_indexes", data.eligible_file_indexes); } static unique_ptr VortexScanDeserialize(Deserializer &deserializer, TableFunction &) { @@ -1042,8 +1112,11 @@ static unique_ptr VortexScanDeserialize(Deserializer &deserializer auto explicit_split_mode = deserializer.ReadProperty(107, "explicit_split_mode"); auto splits_applied = deserializer.ReadProperty(108, "splits_applied"); auto assigned_file_indexes = deserializer.ReadProperty>(109, "assigned_file_indexes"); - auto aggregate_scan = deserializer.ReadProperty(110, "aggregate_scan"); - auto eligible_file_indexes = deserializer.ReadProperty>(111, "eligible_file_indexes"); + auto assigned_row_starts = deserializer.ReadProperty>(110, "assigned_row_starts"); + auto assigned_row_ends = deserializer.ReadProperty>(111, "assigned_row_ends"); + auto assigned_estimated_bytes = deserializer.ReadProperty>(112, "assigned_estimated_bytes"); + auto aggregate_scan = deserializer.ReadProperty(113, "aggregate_scan"); + auto eligible_file_indexes = deserializer.ReadProperty>(114, "eligible_file_indexes"); if (types.size() != names.size() || portable_bind.empty() || !IsCanonicalVortexScanId(scan_split_set_id) || source_urls.size() != paths.size() || source_urls.size() != sizes.size()) { @@ -1074,27 +1147,46 @@ static unique_ptr VortexScanDeserialize(Deserializer &deserializer } previous_eligible = optional_idx(file_index); } - unordered_set assigned; - optional_idx previous_assigned; - for (auto file_index : assigned_file_indexes) { - if (file_index >= files.size() || !assigned.insert(file_index).second || - !eligible.count(file_index) || - (previous_assigned.IsValid() && previous_assigned.GetIndex() >= file_index)) { - throw SerializationException("Invalid assigned Vortex file index %llu", - static_cast(file_index)); + if (assigned_file_indexes.size() != assigned_row_starts.size() || + assigned_file_indexes.size() != assigned_row_ends.size() || + assigned_file_indexes.size() != assigned_estimated_bytes.size()) { + throw SerializationException("Serialized Vortex fragment vectors have different lengths"); + } + vector assigned_fragments; + assigned_fragments.reserve(assigned_file_indexes.size()); + bool has_previous_fragment = false; + idx_t previous_file_index = 0; + idx_t previous_row_start = 0; + idx_t previous_row_end = 0; + for (idx_t fragment_index = 0; fragment_index < assigned_file_indexes.size(); fragment_index++) { + const auto file_index = assigned_file_indexes[fragment_index]; + const auto row_start = assigned_row_starts[fragment_index]; + const auto row_end = assigned_row_ends[fragment_index]; + const auto estimated_bytes = assigned_estimated_bytes[fragment_index]; + const bool same_file = has_previous_fragment && previous_file_index == file_index; + if (file_index >= files.size() || !eligible.count(file_index) || row_start > row_end || + estimated_bytes == DConstants::INVALID_INDEX || estimated_bytes > files[file_index].size || + (has_previous_fragment && previous_file_index > file_index) || + (same_file && previous_row_start >= row_start) || (same_file && previous_row_end > row_start)) { + throw SerializationException("Invalid assigned Vortex fragment at index %llu", + static_cast(fragment_index)); } - previous_assigned = optional_idx(file_index); + assigned_fragments.push_back({file_index, row_start, row_end, estimated_bytes}); + has_previous_fragment = true; + previous_file_index = file_index; + previous_row_start = row_start; + previous_row_end = row_end; } if (!explicit_split_mode && - (splits_applied || !eligible_file_indexes.empty() || !assigned_file_indexes.empty())) { + (splits_applied || !eligible_file_indexes.empty() || !assigned_fragments.empty())) { throw SerializationException("Native Vortex bind contains distributed split state"); } - if (!splits_applied && !assigned_file_indexes.empty()) { + if (!splits_applied && !assigned_fragments.empty()) { throw SerializationException( - "Detached Vortex bind contains assigned files without an applied split batch"); + "Detached Vortex bind contains assigned fragments without an applied split batch"); } - if (aggregate_scan && splits_applied && !assigned_file_indexes.empty() && - assigned_file_indexes != eligible_file_indexes) { + if (aggregate_scan && splits_applied && !assigned_fragments.empty() && + !IsCompleteAggregateVortexAssignment(assigned_fragments, eligible_file_indexes, files)) { throw SerializationException( "Distributed aggregate Vortex bind contains an incomplete file assignment"); } @@ -1107,7 +1199,7 @@ static unique_ptr VortexScanDeserialize(Deserializer &deserializer result->explicit_split_mode = explicit_split_mode; result->splits_applied = splits_applied; result->eligible_file_indexes = std::move(eligible_file_indexes); - result->assigned_file_indexes = std::move(assigned_file_indexes); + result->assigned_fragments = std::move(assigned_fragments); return result; } @@ -1157,58 +1249,42 @@ VortexPlanDistributedScanSplits(const TableFunctionDistributedScanPlanningInput // physical scan. Such a bind is intentionally detached from the original // connection, but its owned portable state and immutable file identities // are sufficient for deterministic split planning. - // target_split_count is a granularity hint. The current stable Vortex - // reader API exposes complete files, but not independently reopenable - // fragments, so this callback deliberately keeps one split per file. auto snapshot = bind_data.CreatePortableSnapshot(); auto selected_file_indexes = SelectDistributedVortexFiles(input, snapshot.distributed_files.size()); vector result; if (selected_file_indexes.empty()) { return result; } - result.reserve(snapshot.aggregate_scan ? 1 : selected_file_indexes.size()); - idx_t total_bytes = 0; - for (auto file_index : selected_file_indexes) { - const auto &file = snapshot.distributed_files[file_index]; - total_bytes = SaturatingVortexSplitEstimate(total_bytes, file.size); + const auto fragments = + PlanVortexFragments(snapshot.portable_bind, + selected_file_indexes, + snapshot.aggregate_scan ? selected_file_indexes.size() + : MaxValue(input.target_split_count, 1)); + if (fragments.empty()) { + throw InvalidInputException("Vortex produced no fragments for a non-empty distributed scan"); } - const auto has_estimated_rows = input.estimated_cardinality != DConstants::INVALID_INDEX; - const auto estimated_rows = has_estimated_rows ? input.estimated_cardinality : 0; if (snapshot.aggregate_scan) { + idx_t total_bytes = 0; + for (const auto &fragment : fragments) { + total_bytes = SaturatingVortexSplitEstimate(total_bytes, fragment.estimated_bytes); + } DistributedScanSplit split; - split.split_id = CanonicalVortexSplitId(selected_file_indexes); - split.payload = - EncodeVortexSplit(selected_file_indexes, snapshot.distributed_files, bind_data.scan_split_set_id); + split.split_id = CanonicalVortexSplitId(fragments); + split.payload = EncodeVortexSplit(fragments, snapshot.distributed_files, bind_data.scan_split_set_id); split.estimated_bytes = optional_idx(total_bytes); split.estimated_cardinality = optional_idx(1); result.push_back(std::move(split)); return result; } - idx_t previous_cumulative_rows = 0; - uint64_t cumulative_bytes = 0; - for (idx_t selected_index = 0; selected_index < selected_file_indexes.size(); selected_index++) { - const auto file_index = selected_file_indexes[selected_index]; - const auto &file = snapshot.distributed_files[file_index]; - vector split_files {file_index}; + result.reserve(fragments.size()); + for (const auto &fragment : fragments) { + vector split_fragments {fragment}; DistributedScanSplit split; - split.split_id = CanonicalVortexSplitId(split_files); + split.split_id = CanonicalVortexSplitId(split_fragments); split.payload = - EncodeVortexSplit(split_files, snapshot.distributed_files, bind_data.scan_split_set_id); - split.estimated_bytes = optional_idx(file.size); - if (has_estimated_rows) { - idx_t cumulative_rows; - if (total_bytes > 0) { - cumulative_bytes = SaturatingVortexSplitEstimate(cumulative_bytes, file.size); - cumulative_rows = - ProportionalVortexSplitEstimate(estimated_rows, cumulative_bytes, total_bytes); - } else { - cumulative_rows = ProportionalVortexSplitEstimate(estimated_rows, - selected_index + 1, - selected_file_indexes.size()); - } - split.estimated_cardinality = optional_idx(cumulative_rows - previous_cumulative_rows); - previous_cumulative_rows = cumulative_rows; - } + EncodeVortexSplit(split_fragments, snapshot.distributed_files, bind_data.scan_split_set_id); + split.estimated_bytes = optional_idx(fragment.estimated_bytes); + split.estimated_cardinality = optional_idx(fragment.row_end - fragment.row_start); result.push_back(std::move(split)); } return result; @@ -1247,16 +1323,11 @@ static void VortexApplyDistributedSplits(optional_ptr worker_bind_ throw InvalidInputException("Distributed aggregate Vortex scans require one complete file-set split"); } unordered_set split_ids; - unordered_set file_indexes; unordered_set eligible_file_indexes(bind_data.eligible_file_indexes.begin(), bind_data.eligible_file_indexes.end()); - vector assigned; - assigned.reserve(bind_data.distributed_files.size()); + vector assigned; for (const auto &split : splits) { split.Validate(); - if (!IsCanonicalVortexSplitId(split.split_id)) { - throw InvalidInputException("Invalid distributed Vortex split id '%s'", split.split_id); - } if (split.payload.empty()) { throw InvalidInputException("Distributed Vortex split '%s' has an empty payload", split.split_id); } @@ -1268,59 +1339,97 @@ static void VortexApplyDistributedSplits(optional_ptr worker_bind_ throw InvalidInputException("Distributed Vortex split '%s' belongs to a different scan identity", split.split_id); } - if (!bind_data.aggregate_scan && decoded.files.size() != 1) { + if (!bind_data.aggregate_scan && decoded.fragments.size() != 1) { throw InvalidInputException( - "Non-aggregate distributed Vortex splits must reference exactly one file"); + "Non-aggregate distributed Vortex splits must reference exactly one fragment"); } - vector decoded_file_indexes; - decoded_file_indexes.reserve(decoded.files.size()); - for (const auto &decoded_file : decoded.files) { - if (decoded_file.file_index >= bind_data.distributed_files.size()) { + vector decoded_fragments; + decoded_fragments.reserve(decoded.fragments.size()); + for (const auto &decoded_fragment : decoded.fragments) { + if (decoded_fragment.file_index >= bind_data.distributed_files.size()) { throw InvalidInputException("Distributed Vortex split '%s' references an unknown file index", split.split_id); } - if (!eligible_file_indexes.count(decoded_file.file_index)) { + if (!eligible_file_indexes.count(decoded_fragment.file_index)) { throw InvalidInputException( "Distributed Vortex split '%s' references file index %llu outside the planned file set", split.split_id, - static_cast(decoded_file.file_index)); + static_cast(decoded_fragment.file_index)); } - if (!SameDistributedFile(decoded_file.file, - bind_data.distributed_files[decoded_file.file_index])) { + if (!SameDistributedFile(decoded_fragment.file, + bind_data.distributed_files[decoded_fragment.file_index])) { throw InvalidInputException( "Distributed Vortex split '%s' does not match the bound file identity", split.split_id); } - if (!file_indexes.insert(decoded_file.file_index).second) { - throw InvalidInputException( - "Distributed Vortex splits reference file index %llu more than once", - static_cast(decoded_file.file_index)); - } - decoded_file_indexes.push_back(decoded_file.file_index); - assigned.push_back(decoded_file.file_index); + decoded_fragments.push_back({decoded_fragment.file_index, + decoded_fragment.row_start, + decoded_fragment.row_end, + decoded_fragment.estimated_bytes}); } - if (split.split_id != CanonicalVortexSplitId(decoded_file_indexes)) { + if (split.split_id != CanonicalVortexSplitId(decoded_fragments)) { throw InvalidInputException( - "Distributed Vortex split id '%s' does not match its payload file indexes", + "Distributed Vortex split id '%s' does not match its fragment payload", split.split_id); } + if (!bind_data.aggregate_scan && + (!split.estimated_cardinality.IsValid() || !split.estimated_bytes.IsValid() || + split.estimated_cardinality.GetIndex() != + decoded_fragments[0].row_end - decoded_fragments[0].row_start || + split.estimated_bytes.GetIndex() != decoded_fragments[0].estimated_bytes)) { + throw InvalidInputException( + "Distributed Vortex split '%s' estimates do not match its fragment payload", + split.split_id); + } + assigned.insert(assigned.end(), decoded_fragments.begin(), decoded_fragments.end()); } - // A batch is a set of elementary scan splits. Canonicalize its file - // assignment so transport or retry code may reorder those splits without - // changing scan meaning or defeating idempotent re-application. - std::sort(assigned.begin(), assigned.end()); - if (bind_data.aggregate_scan && !splits.empty() && assigned != bind_data.eligible_file_indexes) { - throw InvalidInputException( - "Distributed aggregate Vortex split does not contain the complete planned file set"); + // A batch is a set of elementary scan fragments. Canonicalize its assignment + // so transport or retry code may reorder splits without changing scan meaning. + std::sort(assigned.begin(), assigned.end(), [](const auto &left, const auto &right) { + if (left.file_index != right.file_index) { + return left.file_index < right.file_index; + } + if (left.row_start != right.row_start) { + return left.row_start < right.row_start; + } + return left.row_end < right.row_end; + }); + for (idx_t fragment_index = 1; fragment_index < assigned.size(); fragment_index++) { + const auto &previous = assigned[fragment_index - 1]; + const auto ¤t = assigned[fragment_index]; + if (previous.file_index == current.file_index && + (previous.row_start >= current.row_start || previous.row_end > current.row_start)) { + throw InvalidInputException( + "Distributed Vortex fragment assignment overlaps within file index %llu", + static_cast(current.file_index)); + } + } + if (bind_data.aggregate_scan && !splits.empty()) { + if (!IsCompleteAggregateVortexAssignment(assigned, + bind_data.eligible_file_indexes, + bind_data.distributed_files)) { + throw InvalidInputException( + "Distributed aggregate Vortex split does not contain the complete planned file set"); + } + idx_t expected_bytes = 0; + for (const auto &fragment : assigned) { + expected_bytes = SaturatingVortexSplitEstimate(expected_bytes, fragment.estimated_bytes); + } + if (!splits[0].estimated_cardinality.IsValid() || !splits[0].estimated_bytes.IsValid() || + splits[0].estimated_cardinality.GetIndex() != 1 || + splits[0].estimated_bytes.GetIndex() != expected_bytes) { + throw InvalidInputException( + "Distributed aggregate Vortex split estimates do not match its fragment payload"); + } } if (bind_data.splits_applied) { - if (assigned != bind_data.assigned_file_indexes) { + if (assigned != bind_data.assigned_fragments) { throw InvalidInputException( "Distributed Vortex bind already has a different explicit split assignment"); } return; } - bind_data.assigned_file_indexes = std::move(assigned); + bind_data.assigned_fragments = std::move(assigned); bind_data.splits_applied = true; } diff --git a/vortex-duckdb/include/vortex.h b/vortex-duckdb/include/vortex.h index 0dbd299c450..95a8644e72f 100644 --- a/vortex-duckdb/include/vortex.h +++ b/vortex-duckdb/include/vortex.h @@ -13,6 +13,15 @@ #define COUNT_STAR_PROJ_IDX UINT64_MAX +#if defined(VORTEX_VANE_DISTRIBUTED) +typedef struct { + uint64_t file_index; + uint64_t row_start; + uint64_t row_end; + uint64_t estimated_bytes; +} VortexDistributedFragmentView; +#endif + #if defined(VORTEX_VANE_DISTRIBUTED) typedef struct { const uint8_t *name; @@ -119,6 +128,27 @@ const uint8_t *duckdb_table_function_distributed_bind_bytes(const void *portable size_t *size_out); #endif +#if defined(VORTEX_VANE_DISTRIBUTED) +extern +duckdb_vx_data duckdb_table_function_distributed_plan_fragments(const uint8_t *portable_bind, + size_t portable_bind_size, + const uint64_t *selected_file_indexes, + size_t selected_file_count, + size_t target_fragment_count, + duckdb_vx_error *error_out); +#endif + +#if defined(VORTEX_VANE_DISTRIBUTED) +extern size_t duckdb_table_function_distributed_fragment_count(const void *fragment_plan); +#endif + +#if defined(VORTEX_VANE_DISTRIBUTED) +extern +bool duckdb_table_function_distributed_fragment_at(const void *fragment_plan, + size_t index, + VortexDistributedFragmentView *fragment_out); +#endif + #if defined(VORTEX_VANE_DISTRIBUTED) extern size_t duckdb_table_function_distributed_file_count(const void *portable_bind); #endif @@ -158,8 +188,8 @@ bool duckdb_table_function_distributed_file_is_selected(duckdb_vx_table_filter_s extern duckdb_vx_data duckdb_table_function_init_global_distributed(const uint8_t *portable_bind, size_t portable_bind_size, - const uint64_t *assigned_file_indexes, - size_t assigned_file_count, + const VortexDistributedFragmentView *assigned_fragments, + size_t assigned_fragment_count, bool ignore_optional_filters, const duckdb_vx_tfunc_init_input *init_input, duckdb_vx_error *error_out); diff --git a/vortex-duckdb/src/distributed.rs b/vortex-duckdb/src/distributed.rs index 0809c9b4d26..554995e92e7 100644 --- a/vortex-duckdb/src/distributed.rs +++ b/vortex-duckdb/src/distributed.rs @@ -11,23 +11,31 @@ use std::sync::Arc; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; +use futures::StreamExt; +use futures::TryStreamExt; use prost::Message; use vortex::dtype::DType; use vortex::dtype::proto::dtype as pb_dtype; +use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_bail; use vortex::error::vortex_err; use vortex::expr::Expression; use vortex::expr::proto::ExprSerializeProtoExt; +use vortex::io::runtime::BlockingRuntime as _; use vortex::proto::expr as pb_expr; use vortex::scan::DataSource; use vortex_utils::aliases::hash_set::HashSet; +use vortex_utils::parallelism::get_available_parallelism; +use crate::RUNTIME; use crate::SESSION; use crate::convert::PushedAggregate; use crate::duckdb::AggregatePushdownInputRef; use crate::multi_file::BoundFile; use crate::multi_file::build_bound_file_scan; +use crate::multi_file::build_bound_fragment_scan; +use crate::multi_file::open_bound_file; use crate::multi_file::validate_bound_file; use crate::projection::DuckdbField; use crate::projection::extract_schema_from_dtype; @@ -98,6 +106,32 @@ pub struct DistributedRuntimeGlobal { pub global_data: TableFunctionGlobal, } +/// One independently reopenable row range within an immutable bound Vortex file. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct DistributedFragment { + /// Stable coordinator index of the bound file. + pub file_index: usize, + /// Inclusive root-coordinate row offset. + pub row_start: u64, + /// Exclusive root-coordinate row offset. + pub row_end: u64, + /// Proportional on-storage byte estimate for scheduling. + pub estimated_bytes: u64, +} + +/// Owned result of deterministic distributed fragment planning. +pub struct DistributedFragmentPlan { + /// Canonically ordered fragments grouped by file index. + pub fragments: Vec, +} + +struct NaturalFileFragments { + file_index: usize, + file_size: u64, + row_count: u64, + ranges: Vec>, +} + fn encode_expression(expression: &Expression) -> VortexResult> { Ok(expression.serialize_proto()?.encode_to_vec()) } @@ -173,6 +207,169 @@ fn decode_proto(bytes: &[u8]) -> VortexResult { Ok(proto) } +fn validate_natural_ranges( + path: &str, + row_count: u64, + ranges: &[std::ops::Range], +) -> VortexResult<()> { + if row_count == 0 { + if !ranges.is_empty() { + vortex_bail!("Empty Vortex file '{path}' produced non-empty scan fragments"); + } + return Ok(()); + } + if ranges.is_empty() + || ranges[0].start != 0 + || ranges.last().is_none_or(|range| range.end != row_count) + || ranges + .iter() + .any(|range| range.start >= range.end || range.end > row_count) + || ranges.windows(2).any(|pair| pair[0].end != pair[1].start) + { + vortex_bail!( + "Vortex file '{path}' produced scan fragments with a gap, overlap, or invalid bound" + ); + } + Ok(()) +} + +fn allocate_fragment_counts(files: &[NaturalFileFragments], target_count: usize) -> Vec { + if files.is_empty() { + return Vec::new(); + } + let capacities = files + .iter() + .map(|file| file.ranges.len().max(1)) + .collect::>(); + let maximum_count = capacities.iter().sum::(); + let desired_count = target_count.max(files.len()).min(maximum_count); + let remaining = desired_count - files.len(); + let total_extra_capacity = maximum_count - files.len(); + let mut counts = vec![1; files.len()]; + if remaining == 0 || total_extra_capacity == 0 { + return counts; + } + + let mut remainders = Vec::with_capacity(files.len()); + let mut allocated = 0; + for (file_index, &capacity) in capacities.iter().enumerate() { + let extra_capacity = capacity - 1; + let scaled = (remaining as u128) * (extra_capacity as u128); + let extra = usize::try_from(scaled / (total_extra_capacity as u128)) + .vortex_expect("proportional fragment allocation must fit in usize"); + counts[file_index] += extra; + allocated += extra; + remainders.push((scaled % (total_extra_capacity as u128), file_index)); + } + remainders + .sort_unstable_by(|left, right| right.0.cmp(&left.0).then_with(|| left.1.cmp(&right.1))); + for (_, file_index) in remainders.into_iter().take(remaining - allocated) { + counts[file_index] += 1; + } + counts +} + +fn coalesce_ranges( + ranges: &[std::ops::Range], + fragment_count: usize, +) -> Vec> { + if ranges.is_empty() { + // Keep one identity-bearing zero-row fragment so every eligible immutable file remains + // represented in normal and aggregate complete-set plans. + return vec![0..0]; + } + debug_assert!(fragment_count > 0 && fragment_count <= ranges.len()); + (0..fragment_count) + .map(|fragment_index| { + let start_index = usize::try_from( + (fragment_index as u128) * (ranges.len() as u128) / (fragment_count as u128), + ) + .vortex_expect("coalesced fragment start must fit in usize"); + let end_index = usize::try_from( + ((fragment_index + 1) as u128) * (ranges.len() as u128) / (fragment_count as u128), + ) + .vortex_expect("coalesced fragment end must fit in usize"); + ranges[start_index].start..ranges[end_index - 1].end + }) + .collect() +} + +fn estimate_fragment_bytes( + file_size: u64, + row_count: u64, + row_range: &std::ops::Range, +) -> u64 { + if row_count == 0 { + return file_size; + } + let scaled_start = u128::from(file_size) * u128::from(row_range.start) / u128::from(row_count); + let scaled_end = u128::from(file_size) * u128::from(row_range.end) / u128::from(row_count); + u64::try_from(scaled_end - scaled_start) + .vortex_expect("a proportional fragment estimate cannot exceed its u64 file size") +} + +/// Reopen selected immutable files and plan canonical row-range fragments. +pub fn plan_fragments( + bytes: &[u8], + selected_file_indexes: &[u64], + target_count: usize, +) -> VortexResult { + let decoded = decode_bind(bytes)?; + let mut previous_file_index = None; + let mut selected_files = Vec::with_capacity(selected_file_indexes.len()); + for &file_index in selected_file_indexes { + let file_index = usize::try_from(file_index)?; + if previous_file_index.is_some_and(|previous| previous >= file_index) { + vortex_bail!("Distributed Vortex fragment files are not in canonical order"); + } + previous_file_index = Some(file_index); + let file = decoded.files.get(file_index).ok_or_else(|| { + vortex_err!("Unknown distributed Vortex fragment file index: {file_index}") + })?; + selected_files.push((file_index, file.clone())); + } + let concurrency = get_available_parallelism() + .unwrap_or(1) + .min(selected_files.len().max(1)); + let files = RUNTIME.block_on(async move { + futures::stream::iter(selected_files) + .map(|(file_index, file)| async move { + let vortex_file = open_bound_file(&file).await?; + let row_count = vortex_file.row_count(); + let ranges = vortex_file.splits()?; + validate_natural_ranges(&file.path, row_count, &ranges)?; + Ok::<_, vortex::error::VortexError>(NaturalFileFragments { + file_index, + file_size: file.size, + row_count, + ranges, + }) + }) + // `buffered` opens files concurrently while preserving canonical input order. + .buffered(concurrency) + .try_collect::>() + .await + })?; + + let fragment_counts = allocate_fragment_counts(&files, target_count.max(1)); + let mut fragments = Vec::with_capacity(fragment_counts.iter().sum()); + for (file, fragment_count) in files.iter().zip(fragment_counts) { + for row_range in coalesce_ranges(&file.ranges, fragment_count) { + fragments.push(DistributedFragment { + file_index: file.file_index, + row_start: row_range.start, + row_end: row_range.end, + estimated_bytes: estimate_fragment_bytes( + file.file_size, + file.row_count, + &row_range, + ), + }); + } + } + Ok(DistributedFragmentPlan { fragments }) +} + pub fn serialize_bind(bind_data: &TableFunctionBind) -> VortexResult { let dtype = pb_dtype::DType::try_from(bind_data.data_source.dtype())?.encode_to_vec(); let projections = bind_data @@ -353,29 +550,64 @@ pub fn pushdown_serialized_projection_aggregates( pub fn deserialize_runtime_bind( bytes: &[u8], - assigned_file_indexes: &[u64], + assigned_fragments: &[DistributedFragment], ) -> VortexResult { let decoded = decode_bind(bytes)?; + let aggregate_scan = !decoded.aggregates.is_empty(); - let mut seen = HashSet::new(); - let mut selected_files = Vec::with_capacity(assigned_file_indexes.len()); - let mut file_indexes = Vec::with_capacity(assigned_file_indexes.len()); - for &file_index in assigned_file_indexes { - let index = usize::try_from(file_index)?; - if !seen.insert(index) { - vortex_bail!("Duplicate distributed Vortex file index: {index}"); + let mut selected_files = Vec::with_capacity(assigned_fragments.len()); + let mut row_ranges = Vec::with_capacity(assigned_fragments.len()); + let mut file_indexes = Vec::with_capacity(assigned_fragments.len()); + let mut previous_fragment: Option<&DistributedFragment> = None; + for fragment in assigned_fragments { + if fragment.row_start > fragment.row_end || fragment.estimated_bytes == u64::MAX { + vortex_bail!( + "Distributed Vortex fragment has an invalid row range or byte estimate: {}..{}", + fragment.row_start, + fragment.row_end + ); } - selected_files.push( - decoded - .files - .get(index) - .ok_or_else(|| vortex_err!("Unknown distributed Vortex file index: {index}"))? - .clone(), - ); + let index = fragment.file_index; + if let Some(previous) = previous_fragment + && (previous.file_index > index + || (previous.file_index == index && previous.row_start >= fragment.row_start)) + { + vortex_bail!("Distributed Vortex fragments are not in canonical order"); + } + if let Some(previous) = previous_fragment + && previous.file_index == index + && previous.row_end > fragment.row_start + { + vortex_bail!("Distributed Vortex fragments overlap within file index {index}"); + } + let file = decoded + .files + .get(index) + .ok_or_else(|| vortex_err!("Unknown distributed Vortex file index: {index}"))?; + if fragment.estimated_bytes > file.size { + vortex_bail!( + "Distributed Vortex fragment byte estimate {} exceeds file size {}", + fragment.estimated_bytes, + file.size + ); + } + if aggregate_scan && (fragment.row_start != 0 || fragment.estimated_bytes != file.size) { + vortex_bail!( + "Distributed aggregate Vortex fragment for file index {index} must start at row zero and estimate the complete file size" + ); + } + selected_files.push(file.clone()); + row_ranges.push(fragment.row_start..fragment.row_end); file_indexes.push(index); + previous_fragment = Some(fragment); } - let data_source = build_bound_file_scan(&selected_files, Some(decoded.dtype.clone()))?; + let data_source = build_bound_fragment_scan( + &selected_files, + &row_ranges, + aggregate_scan, + Some(decoded.dtype.clone()), + )?; if data_source.dtype() != &decoded.dtype { vortex_bail!( "Distributed Vortex file schema differs from the coordinator bind: expected {}, got {}", @@ -393,3 +625,194 @@ pub fn deserialize_runtime_bind( aggregates: decoded.aggregates, }) } + +#[cfg(test)] +mod tests { + use vortex::dtype::Nullability; + use vortex::dtype::PType; + use vortex::dtype::StructFields; + + use super::*; + + fn natural_file( + file_index: usize, + row_count: u64, + ranges: Vec>, + ) -> NaturalFileFragments { + NaturalFileFragments { + file_index, + file_size: row_count * 10, + row_count, + ranges, + } + } + + fn runtime_bind_bytes_with_aggregates( + file_sizes: &[u64], + aggregates: Vec, + ) -> VortexResult> { + let dtype = DType::Struct( + StructFields::from_iter([( + "value", + DType::Primitive(PType::I64, Nullability::NonNullable), + )]), + Nullability::NonNullable, + ); + Ok(PortableBindProto { + version: PORTABLE_BIND_VERSION, + dtype: pb_dtype::DType::try_from(&dtype)?.encode_to_vec(), + projections: Vec::new(), + filters: Vec::new(), + aggregates, + has_non_optional_filter: false, + files: file_sizes + .iter() + .enumerate() + .map(|(file_index, &size)| FileProto { + source_url: "file:///".to_string(), + path: format!("runtime-{file_index}.vortex"), + size: Some(size), + e_tag: Some(format!("etag-{file_index}")), + version: None, + }) + .collect(), + } + .encode_to_vec()) + } + + fn runtime_bind_bytes(file_sizes: &[u64]) -> VortexResult> { + runtime_bind_bytes_with_aggregates(file_sizes, Vec::new()) + } + + fn aggregate_runtime_bind_bytes(file_sizes: &[u64]) -> VortexResult> { + runtime_bind_bytes_with_aggregates( + file_sizes, + vec![AggregateProto { + projection_id: 0, + kind: 7, + }], + ) + } + + fn fragment( + file_index: usize, + row_start: u64, + row_end: u64, + estimated_bytes: u64, + ) -> DistributedFragment { + DistributedFragment { + file_index, + row_start, + row_end, + estimated_bytes, + } + } + + fn runtime_bind_error(bytes: &[u8], fragments: &[DistributedFragment]) -> String { + deserialize_runtime_bind(bytes, fragments) + .err() + .vortex_expect("invalid fragments must fail") + .to_string() + } + + #[test] + fn fragment_counts_honor_target_and_capacity() { + let files = vec![ + natural_file(0, 40, vec![0..10, 10..20, 20..30, 30..40]), + natural_file(1, 20, vec![0..10, 10..20]), + ]; + + assert_eq!(allocate_fragment_counts(&files, 1), vec![1, 1]); + assert_eq!(allocate_fragment_counts(&files, 4), vec![3, 1]); + assert_eq!(allocate_fragment_counts(&files, 20), vec![4, 2]); + } + + #[test] + fn coalesced_fragments_tile_natural_ranges() { + let ranges = vec![0..10, 10..20, 20..30, 30..40, 40..50]; + + assert_eq!(coalesce_ranges(&ranges, 2), vec![0..20, 20..50]); + assert_eq!(coalesce_ranges(&ranges, 3), vec![0..10, 10..30, 30..50]); + assert_eq!(coalesce_ranges(&[], 1), vec![0..0]); + } + + #[test] + fn proportional_byte_estimates_sum_to_file_size() { + let ranges = [0..3, 3..7, 7..10]; + let estimates = ranges + .iter() + .map(|range| estimate_fragment_bytes(101, 10, range)) + .collect::>(); + + assert_eq!(estimates, vec![30, 40, 31]); + assert_eq!(estimates.iter().sum::(), 101); + } + + #[test] + fn natural_fragment_validation_rejects_gaps_and_overlap() { + assert!(validate_natural_ranges("gap.vortex", 10, &[0..4, 5..10]).is_err()); + assert!(validate_natural_ranges("overlap.vortex", 10, &[0..6, 5..10]).is_err()); + assert!(validate_natural_ranges("reversed.vortex", 10, &[0..6, 6..5]).is_err()); + assert!(validate_natural_ranges("past-end.vortex", 10, &[0..11]).is_err()); + assert!(validate_natural_ranges("missing.vortex", 10, &[]).is_err()); + assert!(validate_natural_ranges("nonempty.vortex", 10, &[0..0, 0..10]).is_err()); + assert!(validate_natural_ranges("valid.vortex", 10, &[0..4, 4..10]).is_ok()); + assert!(validate_natural_ranges("empty.vortex", 0, &[]).is_ok()); + assert!(validate_natural_ranges("invalid-empty.vortex", 0, &[0..0]).is_err()); + } + + #[test] + fn runtime_bind_rejects_out_of_order_fragments() -> VortexResult<()> { + let bytes = runtime_bind_bytes(&[10, 10])?; + let error = runtime_bind_error(&bytes, &[fragment(1, 0, 10, 10), fragment(0, 0, 10, 10)]); + + assert!(error.contains("not in canonical order")); + Ok(()) + } + + #[test] + fn runtime_bind_rejects_overlapping_fragments() -> VortexResult<()> { + let bytes = runtime_bind_bytes(&[10])?; + let error = runtime_bind_error(&bytes, &[fragment(0, 0, 6, 6), fragment(0, 5, 10, 5)]); + + assert!(error.contains("overlap within file index 0")); + Ok(()) + } + + #[test] + fn runtime_bind_rejects_estimate_larger_than_file() -> VortexResult<()> { + let bytes = runtime_bind_bytes(&[10])?; + let error = runtime_bind_error(&bytes, &[fragment(0, 0, 10, 11)]); + + assert!(error.contains("byte estimate 11 exceeds file size 10")); + Ok(()) + } + + #[test] + fn runtime_bind_rejects_unknown_file_index() -> VortexResult<()> { + let bytes = runtime_bind_bytes(&[10])?; + let error = runtime_bind_error(&bytes, &[fragment(1, 0, 10, 10)]); + + assert!(error.contains("Unknown distributed Vortex file index: 1")); + Ok(()) + } + + #[test] + fn aggregate_runtime_bind_defers_complete_file_open() -> VortexResult<()> { + let bytes = aggregate_runtime_bind_bytes(&[10])?; + + assert!(deserialize_runtime_bind(&bytes, &[fragment(0, 0, 42, 10)]).is_ok()); + Ok(()) + } + + #[test] + fn aggregate_runtime_bind_rejects_partial_file_metadata() -> VortexResult<()> { + let bytes = aggregate_runtime_bind_bytes(&[10])?; + let start_error = runtime_bind_error(&bytes, &[fragment(0, 1, 42, 10)]); + let estimate_error = runtime_bind_error(&bytes, &[fragment(0, 0, 42, 9)]); + + assert!(start_error.contains("must start at row zero")); + assert!(estimate_error.contains("estimate the complete file size")); + Ok(()) + } +} diff --git a/vortex-duckdb/src/ffi.rs b/vortex-duckdb/src/ffi.rs index 5bf29bf0cf9..669427c261e 100644 --- a/vortex-duckdb/src/ffi.rs +++ b/vortex-duckdb/src/ffi.rs @@ -9,6 +9,8 @@ use std::ptr; use num_traits::AsPrimitive; use vortex::error::VortexExpect; #[cfg(vortex_vane_distributed)] +use vortex::error::VortexResult; +#[cfg(vortex_vane_distributed)] use vortex::error::vortex_err; use crate::convert::can_push_expression; @@ -20,6 +22,10 @@ use crate::copy::copy_to_initialize_global; use crate::copy::copy_to_sink; use crate::cpp; #[cfg(vortex_vane_distributed)] +use crate::distributed::DistributedFragment; +#[cfg(vortex_vane_distributed)] +use crate::distributed::DistributedFragmentPlan; +#[cfg(vortex_vane_distributed)] use crate::distributed::DistributedRuntimeGlobal; #[cfg(vortex_vane_distributed)] use crate::distributed::PortableDistributedBind; @@ -28,6 +34,8 @@ use crate::distributed::deserialize_bind; #[cfg(vortex_vane_distributed)] use crate::distributed::deserialize_runtime_bind; #[cfg(vortex_vane_distributed)] +use crate::distributed::plan_fragments; +#[cfg(vortex_vane_distributed)] use crate::distributed::pushdown_serialized_projection_aggregates; #[cfg(vortex_vane_distributed)] use crate::distributed::serialize_bind; @@ -82,6 +90,30 @@ pub struct VortexDistributedFieldView { pub logical_type: cpp::duckdb_logical_type, } +#[repr(C)] +#[derive(Clone, Copy)] +#[cfg(vortex_vane_distributed)] +pub struct VortexDistributedFragmentView { + pub file_index: u64, + pub row_start: u64, + pub row_end: u64, + pub estimated_bytes: u64, +} + +#[cfg(vortex_vane_distributed)] +impl TryFrom for DistributedFragment { + type Error = vortex::error::VortexError; + + fn try_from(fragment: VortexDistributedFragmentView) -> Result { + Ok(Self { + file_index: usize::try_from(fragment.file_index)?, + row_start: fragment.row_start, + row_end: fragment.row_end, + estimated_bytes: fragment.estimated_bytes, + }) + } +} + #[unsafe(no_mangle)] unsafe extern "C-unwind" fn duckdb_table_function_to_string( bind_data: *const c_void, @@ -368,6 +400,81 @@ pub unsafe extern "C-unwind" fn duckdb_table_function_distributed_bind_bytes( portable_bind.encoded.as_ptr() } +#[cfg(vortex_vane_distributed)] +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn duckdb_table_function_distributed_plan_fragments( + portable_bind: *const u8, + portable_bind_size: usize, + selected_file_indexes: *const u64, + selected_file_count: usize, + target_fragment_count: usize, + error_out: *mut cpp::duckdb_vx_error, +) -> cpp::duckdb_vx_data { + try_or_null(error_out, || { + if portable_bind.is_null() && portable_bind_size != 0 { + return Err(vortex_err!("Distributed Vortex bind bytes are null")); + } + if selected_file_indexes.is_null() && selected_file_count != 0 { + return Err(vortex_err!( + "Distributed Vortex fragment file indexes are null" + )); + } + let bytes = if portable_bind_size == 0 { + &[] + } else { + unsafe { std::slice::from_raw_parts(portable_bind, portable_bind_size) } + }; + let file_indexes = if selected_file_count == 0 { + &[] + } else { + unsafe { std::slice::from_raw_parts(selected_file_indexes, selected_file_count) } + }; + Ok(Data::from(Box::new(plan_fragments( + bytes, + file_indexes, + target_fragment_count, + )?)) + .as_ptr()) + }) +} + +#[cfg(vortex_vane_distributed)] +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn duckdb_table_function_distributed_fragment_count( + fragment_plan: *const c_void, +) -> usize { + let fragment_plan = unsafe { fragment_plan.cast::().as_ref() } + .vortex_expect("fragment_plan null pointer"); + fragment_plan.fragments.len() +} + +#[cfg(vortex_vane_distributed)] +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn duckdb_table_function_distributed_fragment_at( + fragment_plan: *const c_void, + index: usize, + fragment_out: *mut VortexDistributedFragmentView, +) -> bool { + if fragment_out.is_null() { + return false; + } + let fragment_plan = unsafe { fragment_plan.cast::().as_ref() } + .vortex_expect("fragment_plan null pointer"); + let Some(fragment) = fragment_plan.fragments.get(index) else { + return false; + }; + unsafe { + fragment_out.write(VortexDistributedFragmentView { + file_index: u64::try_from(fragment.file_index) + .vortex_expect("fragment file index must fit in u64"), + row_start: fragment.row_start, + row_end: fragment.row_end, + estimated_bytes: fragment.estimated_bytes, + }) + }; + true +} + #[cfg(vortex_vane_distributed)] #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn duckdb_table_function_distributed_file_count( @@ -480,8 +587,8 @@ pub unsafe extern "C-unwind" fn duckdb_table_function_distributed_file_is_select pub unsafe extern "C-unwind" fn duckdb_table_function_init_global_distributed( portable_bind: *const u8, portable_bind_size: usize, - assigned_file_indexes: *const u64, - assigned_file_count: usize, + assigned_fragments: *const VortexDistributedFragmentView, + assigned_fragment_count: usize, ignore_optional_filters: bool, init_input: *const cpp::duckdb_vx_tfunc_init_input, error_out: *mut cpp::duckdb_vx_error, @@ -490,20 +597,25 @@ pub unsafe extern "C-unwind" fn duckdb_table_function_init_global_distributed( if portable_bind.is_null() && portable_bind_size != 0 { return Err(vortex_err!("Distributed Vortex bind bytes are null")); } - if assigned_file_indexes.is_null() && assigned_file_count != 0 { - return Err(vortex_err!("Distributed Vortex file indexes are null")); + if assigned_fragments.is_null() && assigned_fragment_count != 0 { + return Err(vortex_err!("Distributed Vortex fragments are null")); } let bytes = if portable_bind_size == 0 { &[] } else { unsafe { std::slice::from_raw_parts(portable_bind, portable_bind_size) } }; - let indexes = if assigned_file_count == 0 { + let fragments = if assigned_fragment_count == 0 { &[] } else { - unsafe { std::slice::from_raw_parts(assigned_file_indexes, assigned_file_count) } + unsafe { std::slice::from_raw_parts(assigned_fragments, assigned_fragment_count) } }; - let bind_data = deserialize_runtime_bind(bytes, indexes)?; + let fragments = fragments + .iter() + .copied() + .map(DistributedFragment::try_from) + .collect::>>()?; + let bind_data = deserialize_runtime_bind(bytes, &fragments)?; let input = unsafe { init_input.as_ref() }.vortex_expect("init_input null pointer"); let runtime_input = cpp::duckdb_vx_tfunc_init_input { bind_data: (&raw const bind_data).cast(), diff --git a/vortex-duckdb/src/multi_file.rs b/vortex-duckdb/src/multi_file.rs index dfe330f9cb9..356536c0ea9 100644 --- a/vortex-duckdb/src/multi_file.rs +++ b/vortex-duckdb/src/multi_file.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +#[cfg(vortex_vane_distributed)] +use std::ops::Range; use std::sync::Arc; use std::sync::LazyLock; @@ -328,7 +330,7 @@ async fn verify_file( } #[cfg(vortex_vane_distributed)] -async fn open_bound_file(file: &BoundFile) -> VortexResult { +pub(crate) async fn open_bound_file(file: &BoundFile) -> VortexResult { let source_url = Url::parse(&file.source_url).map_err(|error| { vortex_err!( "Invalid bound Vortex source URL '{}': {error}", @@ -351,16 +353,71 @@ async fn open_bound_file(file: &BoundFile) -> VortexResult { #[cfg(vortex_vane_distributed)] struct BoundFileReaderFactory { file: BoundFile, + required_row_count: Option, } #[cfg(vortex_vane_distributed)] #[async_trait] impl LayoutReaderFactory for BoundFileReaderFactory { async fn open(&self) -> VortexResult> { - Ok(Some(open_bound_file(&self.file).await?.layout_reader()?)) + let vortex_file = open_bound_file(&self.file).await?; + if let Some(required_row_count) = self.required_row_count + && vortex_file.row_count() != required_row_count + { + vortex_bail!( + "Distributed aggregate Vortex fragment for '{}' does not cover the complete file: expected row count {}, got {}", + self.file.path, + required_row_count, + vortex_file.row_count() + ); + } + Ok(Some(vortex_file.layout_reader()?)) } } +/// Build a reader over immutable file fragments selected by a distributed worker assignment. +/// When `require_complete_files` is set, each range must start at zero and its end is checked +/// against the immutable file's actual row count when the deferred reader opens. +#[cfg(vortex_vane_distributed)] +pub fn build_bound_fragment_scan( + files: &[BoundFile], + row_ranges: &[Range], + require_complete_files: bool, + empty_dtype: Option, +) -> VortexResult { + if files.len() != row_ranges.len() { + vortex_bail!( + "Distributed Vortex fragment file count {} differs from row-range count {}", + files.len(), + row_ranges.len() + ); + } + if require_complete_files + && let Some(row_range) = row_ranges.iter().find(|row_range| row_range.start != 0) + { + vortex_bail!("Complete Vortex file range must start at row zero: {row_range:?}"); + } + let dtype = empty_dtype.ok_or_else(|| vortex_err!("Distributed fragment schema is missing"))?; + let factories = files + .iter() + .cloned() + .zip(row_ranges) + .map(|(file, row_range)| { + Arc::new(BoundFileReaderFactory { + file, + required_row_count: require_complete_files.then_some(row_range.end), + }) as Arc + }) + .collect(); + MultiLayoutDataSource::new_deferred_ranges( + dtype, + factories, + row_ranges.to_vec(), + Vec::new(), + &SESSION, + ) +} + /// Build a reader over an already selected file set. No glob is evaluated /// here, so an empty assignment stays empty and a worker cannot discover /// files that were not part of the coordinator bind. @@ -384,7 +441,12 @@ pub fn build_bound_file_scan( let remaining = files[1..] .iter() .cloned() - .map(|file| Arc::new(BoundFileReaderFactory { file }) as Arc) + .map(|file| { + Arc::new(BoundFileReaderFactory { + file, + required_row_count: None, + }) as Arc + }) .collect(); let byte_sizes = files.iter().map(|file| Some(file.size)).collect(); Ok(MultiLayoutDataSource::new_with_first( diff --git a/vortex-layout/src/scan/multi.rs b/vortex-layout/src/scan/multi.rs index d9de768f271..49bb8bd4a26 100644 --- a/vortex-layout/src/scan/multi.rs +++ b/vortex-layout/src/scan/multi.rs @@ -6,8 +6,8 @@ //! Readers may be pre-opened or deferred via [`LayoutReaderFactory`]. Deferred readers are opened //! concurrently during scanning using `buffer_unordered`: up to `concurrency` file opens run in //! parallel as spawned tasks on the session runtime. Once opened, each reader yields a single -//! partition covering its full row range; internal I/O pipelining and chunking are handled by -//! [`ScanBuilder`]. +//! partition covering either its full row range or an explicitly assigned child row range; +//! internal I/O pipelining and chunking are handled by [`ScanBuilder`]. //! //! # Schema Resolution //! @@ -82,8 +82,9 @@ pub trait LayoutReaderFactory: 'static + Send + Sync { /// Readers may be pre-opened or deferred via [`LayoutReaderFactory`]. Deferred readers are opened /// concurrently during scanning using `buffer_unordered`, mirroring the DuckDB scan pattern: up /// to `concurrency` file opens run in parallel as spawned tasks on the session runtime. Once -/// opened, each reader yields a single partition covering its full row range; internal I/O -/// pipelining and chunking are handled by [`ScanBuilder`]. +/// opened, each reader yields a single partition covering either its full row range or an +/// explicitly assigned child row range; internal I/O pipelining and chunking are handled by +/// [`ScanBuilder`]. #[derive(Clone)] pub struct MultiLayoutDataSource { dtype: DType, @@ -97,11 +98,15 @@ pub enum MultiLayoutChild { reader: LayoutReaderRef, /// On-storage file size in bytes, if known from the listing metadata. byte_size: Option, + /// Root-coordinate row range assigned to this child, or the complete reader when absent. + row_range: Option>, }, Deferred { factory: Arc, /// On-storage file size in bytes, if known from the listing metadata. byte_size: Option, + /// Root-coordinate row range assigned to this child, or the complete reader when absent. + row_range: Option>, }, } @@ -113,6 +118,24 @@ impl MultiLayoutChild { MultiLayoutChild::Deferred { byte_size, .. } => *byte_size, } } + + /// Root-coordinate row range assigned to this child, if any. + pub fn row_range(&self) -> Option<&Range> { + match self { + MultiLayoutChild::Opened { row_range, .. } + | MultiLayoutChild::Deferred { row_range, .. } => row_range.as_ref(), + } + } + + fn known_row_count(&self) -> Option { + if let Some(row_range) = self.row_range() { + return Some(row_range.end.saturating_sub(row_range.start)); + } + match self { + MultiLayoutChild::Opened { reader, .. } => Some(reader.row_count()), + MultiLayoutChild::Deferred { .. } => None, + } + } } impl MultiLayoutDataSource { @@ -148,12 +171,17 @@ impl MultiLayoutDataSource { children.push(MultiLayoutChild::Opened { reader: first, byte_size: first_size, + row_range: None, }); children.extend( remaining .into_iter() .zip_eq(sizes_iter) - .map(|(factory, byte_size)| MultiLayoutChild::Deferred { factory, byte_size }), + .map(|(factory, byte_size)| MultiLayoutChild::Deferred { + factory, + byte_size, + row_range: None, + }), ); Self { @@ -195,12 +223,77 @@ impl MultiLayoutDataSource { children: factories .into_iter() .zip_eq(sizes) - .map(|(factory, byte_size)| MultiLayoutChild::Deferred { factory, byte_size }) + .map(|(factory, byte_size)| MultiLayoutChild::Deferred { + factory, + byte_size, + row_range: None, + }) .collect(), concurrency, } } + /// Creates a multi-layout data source from deferred readers with independently assigned row + /// ranges. + /// + /// Each range uses the corresponding reader's root coordinate space. Ranges may be empty only + /// for an empty reader and must not be reversed. Reader bounds are checked after the deferred + /// reader is opened. + pub fn new_deferred_ranges( + dtype: DType, + factories: Vec>, + row_ranges: Vec>, + byte_sizes: Vec>, + session: &VortexSession, + ) -> VortexResult { + if factories.len() != row_ranges.len() { + vortex_bail!( + "row_ranges length {} must match the number of factories {}", + row_ranges.len(), + factories.len() + ); + } + if !byte_sizes.is_empty() && byte_sizes.len() != factories.len() { + vortex_bail!( + "byte_sizes length {} must match the number of factories {}", + byte_sizes.len(), + factories.len() + ); + } + if let Some(row_range) = row_ranges + .iter() + .find(|row_range| row_range.start > row_range.end) + { + vortex_bail!("Assigned child row range is reversed: {row_range:?}"); + } + + let concurrency = get_available_parallelism().unwrap_or(DEFAULT_CONCURRENCY); + let sizes = if byte_sizes.is_empty() { + vec![None; factories.len()] + } else { + byte_sizes + }; + let children = factories + .into_iter() + .zip_eq(row_ranges) + .zip_eq(sizes) + .map( + |((factory, row_range), byte_size)| MultiLayoutChild::Deferred { + factory, + byte_size, + row_range: Some(row_range), + }, + ) + .collect(); + + Ok(Self { + dtype, + session: session.clone(), + children, + concurrency, + }) + } + pub fn children(&self) -> &[MultiLayoutChild] { &self.children } @@ -223,30 +316,27 @@ impl DataSource for MultiLayoutDataSource { fn row_count(&self) -> Precision { let mut sum: u64 = 0; - let mut opened_count: u64 = 0; - let mut deferred_count: u64 = 0; + let mut known_count: u64 = 0; + let mut unknown_count: u64 = 0; for child in self.children.iter() { - match child { - MultiLayoutChild::Opened { reader, .. } => { - opened_count += 1; - sum = sum.saturating_add(reader.row_count()); - } - MultiLayoutChild::Deferred { .. } => { - deferred_count += 1; - } + if let Some(row_count) = child.known_row_count() { + known_count += 1; + sum = sum.saturating_add(row_count); + } else { + unknown_count += 1; } } - let total_count = opened_count + deferred_count; + let total_count = known_count + unknown_count; if total_count == 0 { return Precision::exact(0u64); } - if deferred_count == 0 { + if unknown_count == 0 { Precision::exact(sum) - } else if opened_count > 0 { - let avg = sum / opened_count; + } else if known_count > 0 { + let avg = sum / known_count; let extrapolated = avg.saturating_mul(total_count); Precision::inexact(extrapolated) } else { @@ -296,10 +386,12 @@ impl DataSource for MultiLayoutDataSource { for child in self.children.iter() { match child { - MultiLayoutChild::Opened { reader, .. } => ready.push_back(Arc::clone(reader)), - MultiLayoutChild::Deferred { factory, .. } => { - deferred.push_back(Arc::clone(factory)) - } + MultiLayoutChild::Opened { + reader, row_range, .. + } => ready.push_back((Arc::clone(reader), row_range.clone())), + MultiLayoutChild::Deferred { + factory, row_range, .. + } => deferred.push_back((Arc::clone(factory), row_range.clone())), } } @@ -364,13 +456,16 @@ impl BoundScanRequest { } } +type RangedLayoutReader = (LayoutReaderRef, Option>); +type RangedLayoutReaderFactory = (Arc, Option>); + struct MultiLayoutScan { session: VortexSession, source_dtype: DType, dtype: DType, request: BoundScanRequest, - ready: VecDeque, - deferred: VecDeque>, + ready: VecDeque, + deferred: VecDeque, handle: vortex_io::runtime::Handle, concurrency: usize, } @@ -409,12 +504,13 @@ impl DataSourceScan for MultiLayoutScan { // Deferred readers are opened concurrently via spawned tasks. // When ordered, we use `buffered` to preserve the original partition order. // When unordered, we use `buffer_unordered` to yield partitions as they open. - let spawned = stream::iter(deferred).map(move |factory| { + let spawned = stream::iter(deferred).map(move |(factory, row_range)| { handle.spawn(async move { - factory + let reader = factory .open() .instrument(tracing::info_span!("LayoutReaderFactory::open")) - .await + .await?; + Ok(reader.map(|reader| (reader, row_range))) }) }); @@ -449,23 +545,53 @@ impl DataSourceScan for MultiLayoutScan { .chain(deferred_stream) .enumerate() .flat_map(move |(i, reader_result)| match reader_result { - Ok(reader) => { - reader_partition(i, reader, session.clone(), &source_dtype, request.clone()) - } + Ok((reader, child_row_range)) => reader_partition( + i, + reader, + child_row_range, + session.clone(), + &source_dtype, + request.clone(), + ), Err(e) => stream::once(async move { Err(e) }).boxed(), }) .boxed() } } +fn assigned_reader_row_range( + child_row_range: Range, + request_row_range: Option<&Range>, + row_count: u64, +) -> VortexResult>> { + if child_row_range.start > child_row_range.end + || child_row_range.end > row_count + || (child_row_range.is_empty() && row_count != 0) + { + vortex_bail!( + "Assigned child row range {:?} is out of bounds for row count {}", + child_row_range, + row_count + ); + } + let row_range = if let Some(request_row_range) = request_row_range { + child_row_range.start.max(request_row_range.start) + ..child_row_range.end.min(request_row_range.end) + } else { + child_row_range + }; + Ok((!row_range.is_empty()).then_some(row_range)) +} + /// Generates a partition stream for a single layout reader. /// /// Checks file-level pruning first (via `pruning_evaluation`). If the filter proves no rows /// can match, returns an empty stream. Otherwise, yields a single partition covering the -/// reader's full row range. +/// reader's assigned row range. fn reader_partition( partition_idx: usize, reader: LayoutReaderRef, + child_row_range: Option>, session: VortexSession, source_dtype: &DType, request: BoundScanRequest, @@ -480,7 +606,18 @@ fn reader_partition( } let row_count = reader.row_count(); - let row_range = request.row_range.clone().unwrap_or(0..row_count); + let row_range = if let Some(child_row_range) = child_row_range { + match assigned_reader_row_range(child_row_range, request.row_range.as_ref(), row_count) { + Ok(Some(row_range)) => row_range, + Ok(None) => return stream::empty().boxed(), + Err(error) => return stream::once(async move { Err(error) }).boxed(), + } + } else { + request.row_range.clone().unwrap_or(0..row_count) + }; + if row_range.is_empty() { + return stream::empty().boxed(); + } let partition_idx_u64: u64 = partition_idx as u64; if let Some(range) = &request.partition_range @@ -632,6 +769,20 @@ mod tests { ) } + fn ranged_deferred_source(row_ranges: Vec>) -> VortexResult { + let factories: Vec> = row_ranges + .iter() + .map(|_| Arc::new(NeverOpened) as _) + .collect(); + MultiLayoutDataSource::new_deferred_ranges( + DType::Bool(Nullability::NonNullable), + factories, + row_ranges, + Vec::new(), + &new_session(), + ) + } + #[rstest] #[case::all_known(vec![Some(10), Some(20), Some(30)], Precision::exact(60u64))] #[case::some_known_extrapolates(vec![Some(10), None, Some(30)], Precision::inexact(60u64))] @@ -641,6 +792,48 @@ mod tests { assert_eq!(deferred_source(sizes).byte_size(), expected); } + #[test] + fn deferred_ranges_have_exact_row_count() -> VortexResult<()> { + let source = ranged_deferred_source(vec![0..4, 4..9, 20..20])?; + + assert_eq!(source.row_count(), Precision::exact(9u64)); + assert_eq!(source.children()[0].row_range(), Some(&(0..4))); + assert_eq!(source.children()[1].row_range(), Some(&(4..9))); + assert_eq!(source.children()[2].row_range(), Some(&(20..20))); + Ok(()) + } + + #[test] + fn deferred_ranges_reject_reversed_range() { + let error = ranged_deferred_source(std::iter::once(Range { start: 5, end: 4 }).collect()) + .err() + .map(|error| error.to_string()); + + assert!( + error.as_deref().is_some_and( + |message| message.contains("Assigned child row range is reversed: 5..4") + ) + ); + } + + #[test] + fn assigned_range_intersects_requested_root_range() -> VortexResult<()> { + assert_eq!( + assigned_reader_row_range(2..8, Some(&(5..12)), 10)?, + Some(5..8) + ); + assert_eq!(assigned_reader_row_range(2..8, Some(&(8..12)), 10)?, None); + assert_eq!(assigned_reader_row_range(0..0, None, 0)?, None); + Ok(()) + } + + #[test] + fn assigned_range_rejects_invalid_empty_and_bounds() { + assert!(assigned_reader_row_range(4..4, None, 10).is_err()); + assert!(assigned_reader_row_range(Range { start: 4, end: 3 }, None, 10).is_err()); + assert!(assigned_reader_row_range(0..11, None, 10).is_err()); + } + #[test] fn filter_binding_errors_are_deferred() -> VortexResult<()> { let dtype = DType::Primitive(PType::U8, Nullability::NonNullable);