fix: trim vortex reader surface

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: aoiasd <zhicheng.yue@zilliz.com>
This commit is contained in:
aoiasd
2026-07-15 14:27:58 +08:00
co-authored by Claude Opus 4.6
parent 9eb312fe80
commit 4060829e64
5 changed files with 25 additions and 610 deletions
@@ -15,10 +15,8 @@
// limitations under the License.
#include <algorithm>
#include <cstring>
#include <limits>
#include <memory>
#include <optional>
#include <type_traits>
#include <utility>
#include <vector>
@@ -28,478 +26,6 @@ namespace milvus {
namespace detail {
inline bool
IsSupportedUnaryScanOp(proto::plan::OpType op_type) {
switch (op_type) {
case proto::plan::OpType::GreaterThan:
case proto::plan::OpType::GreaterEqual:
case proto::plan::OpType::LessThan:
case proto::plan::OpType::LessEqual:
case proto::plan::OpType::Equal:
case proto::plan::OpType::NotEqual:
return true;
default:
return false;
}
}
template <typename T>
inline bool
TryGetScanValue(const proto::plan::GenericValue& value, T* out) {
if constexpr (std::is_same_v<T, bool>) {
if (value.val_case() != proto::plan::GenericValue::kBoolVal) {
return false;
}
*out = value.bool_val();
return true;
} else if constexpr (std::is_integral_v<T>) {
if (value.val_case() != proto::plan::GenericValue::kInt64Val) {
return false;
}
if (value.int64_val() <
static_cast<int64_t>(std::numeric_limits<T>::min()) ||
value.int64_val() >
static_cast<int64_t>(std::numeric_limits<T>::max())) {
return false;
}
*out = static_cast<T>(value.int64_val());
return true;
} else if constexpr (std::is_floating_point_v<T>) {
if (value.val_case() == proto::plan::GenericValue::kFloatVal) {
*out = static_cast<T>(value.float_val());
return true;
}
if (value.val_case() == proto::plan::GenericValue::kInt64Val) {
*out = static_cast<T>(value.int64_val());
return true;
}
return false;
}
return false;
}
template <typename T>
inline bool
CanUseScanValue(const proto::plan::GenericValue& value) {
T typed_value{};
return TryGetScanValue<T>(value, &typed_value);
}
template <typename T>
inline bool
UnaryScanCompare(T lhs, T rhs, proto::plan::OpType op_type) {
switch (op_type) {
case proto::plan::OpType::GreaterThan:
return lhs > rhs;
case proto::plan::OpType::GreaterEqual:
return lhs >= rhs;
case proto::plan::OpType::LessThan:
return lhs < rhs;
case proto::plan::OpType::LessEqual:
return lhs <= rhs;
case proto::plan::OpType::Equal:
return lhs == rhs;
case proto::plan::OpType::NotEqual:
return lhs != rhs;
default:
return false;
}
}
inline void
ResetRowIdPayloadOutput(ChunkedColumnInterface::ScanBatch* out) {
out->values = ChunkedColumnInterface::ValueView{};
out->validity = ChunkedColumnInterface::ValidityView{};
out->row_ids.clear();
out->owner.reset();
out->row_id_start = 0;
out->size = 0;
}
inline void
AppendRowIdPayloadEntry(ChunkedColumnInterface::ScanBatch* out,
FixedVector<bool>* validity,
bool* has_invalid,
int64_t row_id,
bool valid) {
AssertInfo(validity != nullptr, "row id payload validity owner is null");
AssertInfo(has_invalid != nullptr, "row id payload invalid flag is null");
if (!out->row_ids.empty()) {
AssertInfo(out->row_ids.back() <= row_id,
"row id payload is not ordered: {} before {}",
out->row_ids.back(),
row_id);
}
out->row_ids.emplace_back(row_id);
validity->push_back(valid);
*has_invalid = *has_invalid || !valid;
}
inline void
FinalizeRowIdPayloadOutput(ChunkedColumnInterface::ScanBatch* out,
std::shared_ptr<FixedVector<bool>> validity,
bool has_invalid,
bool nullable) {
AssertInfo(validity != nullptr, "row id payload validity owner is null");
AssertInfo(validity->size() == out->row_ids.size(),
"row id payload validity size {} does not match row ids size {}",
validity->size(),
out->row_ids.size());
out->row_id_start = out->row_ids.empty() ? 0 : out->row_ids.front();
out->size = static_cast<int64_t>(out->row_ids.size());
out->validity.size = out->size;
out->validity.nullable = nullable;
if (out->row_ids.empty() || !has_invalid) {
out->validity.encoding =
ChunkedColumnInterface::ValidityEncoding::AllValid;
out->validity.all_valid = true;
return;
}
out->validity.encoding =
ChunkedColumnInterface::ValidityEncoding::BoolArray;
out->validity.data = validity->data();
out->validity.offset = 0;
out->validity.all_valid = false;
out->owner = std::move(validity);
}
class FixedWidthRowIdScanCursor final
: public ChunkedColumnInterface::ScanCursor {
public:
FixedWidthRowIdScanCursor(const ChunkedColumnInterface* column,
milvus::OpContext* op_ctx,
int64_t start_offset,
int64_t length,
DataType data_type,
proto::plan::OpType op_type,
const proto::plan::GenericValue& value)
: column_(column),
op_ctx_(op_ctx),
data_type_(data_type),
op_type_(op_type),
value_(value),
scan_pos_(start_offset),
scan_end_(start_offset + length) {
if (start_offset < scan_end_) {
auto [chunk_id, offset] = column_->GetChunkIDByOffset(start_offset);
current_chunk_id_ = static_cast<int64_t>(chunk_id);
current_chunk_offset_ = static_cast<int64_t>(offset);
}
}
bool
Next(ChunkedColumnInterface::ScanBatch* out) override {
AssertInfo(out != nullptr, "row id scan output batch is null");
ResetRowIdPayloadOutput(out);
if (scan_pos_ >= scan_end_) {
return false;
}
while (scan_pos_ < scan_end_) {
auto& span = GetCurrentSpan();
const auto rows = span.get().row_count();
if (current_chunk_offset_ >= rows) {
++current_chunk_id_;
current_chunk_offset_ = 0;
cached_span_.reset();
continue;
}
const auto rows_left_in_chunk = rows - current_chunk_offset_;
const auto rows_left_in_scan = scan_end_ - scan_pos_;
const auto rows_to_scan =
std::min<int64_t>(rows_left_in_chunk, rows_left_in_scan);
auto validity = std::make_shared<FixedVector<bool>>();
bool has_invalid = false;
ScanSpan(
span.get(), rows_to_scan, out, validity.get(), &has_invalid);
scan_pos_ += rows_to_scan;
current_chunk_offset_ += rows_to_scan;
if (!out->row_ids.empty()) {
FinalizeRowIdPayloadOutput(out,
std::move(validity),
has_invalid,
column_->IsNullable());
return true;
}
}
return false;
}
private:
PinWrapper<SpanBase>&
GetCurrentSpan() {
if (!cached_span_.has_value() ||
cached_chunk_id_ != current_chunk_id_) {
cached_span_ = column_->Span(op_ctx_, current_chunk_id_);
cached_chunk_id_ = current_chunk_id_;
}
return cached_span_.value();
}
template <typename T>
void
ScanTypedSpan(const SpanBase& span,
int64_t rows_to_scan,
ChunkedColumnInterface::ScanBatch* out,
FixedVector<bool>* validity,
bool* has_invalid) {
T value{};
if (!TryGetScanValue<T>(value_, &value)) {
return;
}
const auto* data = static_cast<const T*>(span.data());
const auto* valid = span.valid_data();
for (int64_t i = 0; i < rows_to_scan; ++i) {
const auto local_offset = current_chunk_offset_ + i;
if (valid != nullptr && !valid[local_offset]) {
AppendRowIdPayloadEntry(
out, validity, has_invalid, scan_pos_ + i, false);
continue;
}
if (UnaryScanCompare<T>(data[local_offset], value, op_type_)) {
AppendRowIdPayloadEntry(
out, validity, has_invalid, scan_pos_ + i, true);
}
}
}
void
ScanSpan(const SpanBase& span,
int64_t rows_to_scan,
ChunkedColumnInterface::ScanBatch* out,
FixedVector<bool>* validity,
bool* has_invalid) {
switch (data_type_) {
case DataType::BOOL:
ScanTypedSpan<bool>(
span, rows_to_scan, out, validity, has_invalid);
break;
case DataType::INT8:
ScanTypedSpan<int8_t>(
span, rows_to_scan, out, validity, has_invalid);
break;
case DataType::INT16:
ScanTypedSpan<int16_t>(
span, rows_to_scan, out, validity, has_invalid);
break;
case DataType::INT32:
ScanTypedSpan<int32_t>(
span, rows_to_scan, out, validity, has_invalid);
break;
case DataType::INT64:
case DataType::TIMESTAMPTZ:
ScanTypedSpan<int64_t>(
span, rows_to_scan, out, validity, has_invalid);
break;
case DataType::FLOAT:
ScanTypedSpan<float>(
span, rows_to_scan, out, validity, has_invalid);
break;
case DataType::DOUBLE:
ScanTypedSpan<double>(
span, rows_to_scan, out, validity, has_invalid);
break;
default:
break;
}
}
const ChunkedColumnInterface* column_;
milvus::OpContext* op_ctx_;
DataType data_type_;
proto::plan::OpType op_type_;
proto::plan::GenericValue value_;
int64_t scan_pos_;
int64_t scan_end_;
int64_t current_chunk_id_{0};
int64_t current_chunk_offset_{0};
int64_t cached_chunk_id_{-1};
std::optional<PinWrapper<SpanBase>> cached_span_{std::nullopt};
};
template <typename T>
inline bool
BinaryRangeScanCompare(
T value, T lower, bool lower_inclusive, T upper, bool upper_inclusive) {
const bool lower_ok = lower_inclusive ? value >= lower : value > lower;
const bool upper_ok = upper_inclusive ? value <= upper : value < upper;
return lower_ok && upper_ok;
}
class FixedWidthBinaryRangeRowIdScanCursor final
: public ChunkedColumnInterface::ScanCursor {
public:
FixedWidthBinaryRangeRowIdScanCursor(
const ChunkedColumnInterface* column,
milvus::OpContext* op_ctx,
int64_t start_offset,
int64_t length,
DataType data_type,
const proto::plan::GenericValue& lower_value,
bool lower_inclusive,
const proto::plan::GenericValue& upper_value,
bool upper_inclusive)
: column_(column),
op_ctx_(op_ctx),
data_type_(data_type),
lower_value_(lower_value),
upper_value_(upper_value),
lower_inclusive_(lower_inclusive),
upper_inclusive_(upper_inclusive),
scan_pos_(start_offset),
scan_end_(start_offset + length) {
if (start_offset < scan_end_) {
auto [chunk_id, offset] = column_->GetChunkIDByOffset(start_offset);
current_chunk_id_ = static_cast<int64_t>(chunk_id);
current_chunk_offset_ = static_cast<int64_t>(offset);
}
}
bool
Next(ChunkedColumnInterface::ScanBatch* out) override {
AssertInfo(out != nullptr,
"binary range row id scan output batch is null");
ResetRowIdPayloadOutput(out);
if (scan_pos_ >= scan_end_) {
return false;
}
while (scan_pos_ < scan_end_) {
auto& span = GetCurrentSpan();
const auto rows = span.get().row_count();
if (current_chunk_offset_ >= rows) {
++current_chunk_id_;
current_chunk_offset_ = 0;
cached_span_.reset();
continue;
}
const auto rows_left_in_chunk = rows - current_chunk_offset_;
const auto rows_left_in_scan = scan_end_ - scan_pos_;
const auto rows_to_scan =
std::min<int64_t>(rows_left_in_chunk, rows_left_in_scan);
auto validity = std::make_shared<FixedVector<bool>>();
bool has_invalid = false;
ScanSpan(
span.get(), rows_to_scan, out, validity.get(), &has_invalid);
scan_pos_ += rows_to_scan;
current_chunk_offset_ += rows_to_scan;
if (!out->row_ids.empty()) {
FinalizeRowIdPayloadOutput(out,
std::move(validity),
has_invalid,
column_->IsNullable());
return true;
}
}
return false;
}
private:
PinWrapper<SpanBase>&
GetCurrentSpan() {
if (!cached_span_.has_value() ||
cached_chunk_id_ != current_chunk_id_) {
cached_span_ = column_->Span(op_ctx_, current_chunk_id_);
cached_chunk_id_ = current_chunk_id_;
}
return cached_span_.value();
}
template <typename T>
void
ScanTypedSpan(const SpanBase& span,
int64_t rows_to_scan,
ChunkedColumnInterface::ScanBatch* out,
FixedVector<bool>* validity,
bool* has_invalid) {
T lower{};
T upper{};
if (!TryGetScanValue<T>(lower_value_, &lower) ||
!TryGetScanValue<T>(upper_value_, &upper)) {
return;
}
const auto* data = static_cast<const T*>(span.data());
const auto* valid = span.valid_data();
for (int64_t i = 0; i < rows_to_scan; ++i) {
const auto local_offset = current_chunk_offset_ + i;
if (valid != nullptr && !valid[local_offset]) {
AppendRowIdPayloadEntry(
out, validity, has_invalid, scan_pos_ + i, false);
continue;
}
if (BinaryRangeScanCompare<T>(data[local_offset],
lower,
lower_inclusive_,
upper,
upper_inclusive_)) {
AppendRowIdPayloadEntry(
out, validity, has_invalid, scan_pos_ + i, true);
}
}
}
void
ScanSpan(const SpanBase& span,
int64_t rows_to_scan,
ChunkedColumnInterface::ScanBatch* out,
FixedVector<bool>* validity,
bool* has_invalid) {
switch (data_type_) {
case DataType::BOOL:
ScanTypedSpan<bool>(
span, rows_to_scan, out, validity, has_invalid);
break;
case DataType::INT8:
ScanTypedSpan<int8_t>(
span, rows_to_scan, out, validity, has_invalid);
break;
case DataType::INT16:
ScanTypedSpan<int16_t>(
span, rows_to_scan, out, validity, has_invalid);
break;
case DataType::INT32:
ScanTypedSpan<int32_t>(
span, rows_to_scan, out, validity, has_invalid);
break;
case DataType::INT64:
case DataType::TIMESTAMPTZ:
ScanTypedSpan<int64_t>(
span, rows_to_scan, out, validity, has_invalid);
break;
case DataType::FLOAT:
ScanTypedSpan<float>(
span, rows_to_scan, out, validity, has_invalid);
break;
case DataType::DOUBLE:
ScanTypedSpan<double>(
span, rows_to_scan, out, validity, has_invalid);
break;
default:
break;
}
}
const ChunkedColumnInterface* column_;
milvus::OpContext* op_ctx_;
DataType data_type_;
proto::plan::GenericValue lower_value_;
proto::plan::GenericValue upper_value_;
bool lower_inclusive_;
bool upper_inclusive_;
int64_t scan_pos_;
int64_t scan_end_;
int64_t current_chunk_id_{0};
int64_t current_chunk_offset_{0};
int64_t cached_chunk_id_{-1};
std::optional<PinWrapper<SpanBase>> cached_span_{std::nullopt};
};
class FixedWidthDataScanCursor final
: public ChunkedColumnInterface::ScanCursor {
public:
@@ -879,124 +405,6 @@ class ViewDataScanCursor final : public ChunkedColumnInterface::ScanCursor {
int64_t scan_end_;
};
inline bool
CanUseFixedWidthRowIdScan(DataType data_type,
proto::plan::OpType op_type,
const proto::plan::GenericValue& value) {
if (!IsSupportedUnaryScanOp(op_type)) {
return false;
}
switch (data_type) {
case DataType::BOOL:
return CanUseScanValue<bool>(value);
case DataType::INT8:
return CanUseScanValue<int8_t>(value);
case DataType::INT16:
return CanUseScanValue<int16_t>(value);
case DataType::INT32:
return CanUseScanValue<int32_t>(value);
case DataType::INT64:
case DataType::TIMESTAMPTZ:
return CanUseScanValue<int64_t>(value);
case DataType::FLOAT:
return CanUseScanValue<float>(value);
case DataType::DOUBLE:
return CanUseScanValue<double>(value);
default:
return false;
}
}
inline bool
CanUseFixedWidthBinaryRangeRowIdScan(
DataType data_type,
const proto::plan::GenericValue& lower_value,
const proto::plan::GenericValue& upper_value) {
switch (data_type) {
case DataType::BOOL:
return CanUseScanValue<bool>(lower_value) &&
CanUseScanValue<bool>(upper_value);
case DataType::INT8:
return CanUseScanValue<int8_t>(lower_value) &&
CanUseScanValue<int8_t>(upper_value);
case DataType::INT16:
return CanUseScanValue<int16_t>(lower_value) &&
CanUseScanValue<int16_t>(upper_value);
case DataType::INT32:
return CanUseScanValue<int32_t>(lower_value) &&
CanUseScanValue<int32_t>(upper_value);
case DataType::INT64:
case DataType::TIMESTAMPTZ:
return CanUseScanValue<int64_t>(lower_value) &&
CanUseScanValue<int64_t>(upper_value);
case DataType::FLOAT:
return CanUseScanValue<float>(lower_value) &&
CanUseScanValue<float>(upper_value);
case DataType::DOUBLE:
return CanUseScanValue<double>(lower_value) &&
CanUseScanValue<double>(upper_value);
default:
return false;
}
}
inline std::unique_ptr<ChunkedColumnInterface::ScanCursor>
MakeFixedWidthRowIdScanCursor(const ChunkedColumnInterface* column,
milvus::OpContext* op_ctx,
int64_t start_offset,
int64_t length,
DataType data_type,
proto::plan::OpType op_type,
const proto::plan::GenericValue& value) {
AssertInfo(
start_offset >= 0 && length >= 0 &&
start_offset + length <= static_cast<int64_t>(column->NumRows()),
"row id scan range [{}, {}) out of rows {}",
start_offset,
start_offset + length,
column->NumRows());
if (!CanUseFixedWidthRowIdScan(data_type, op_type, value)) {
return nullptr;
}
return std::make_unique<FixedWidthRowIdScanCursor>(
column, op_ctx, start_offset, length, data_type, op_type, value);
}
inline std::unique_ptr<ChunkedColumnInterface::ScanCursor>
MakeFixedWidthBinaryRangeRowIdScanCursor(
const ChunkedColumnInterface* column,
milvus::OpContext* op_ctx,
int64_t start_offset,
int64_t length,
DataType data_type,
const proto::plan::GenericValue& lower_value,
bool lower_inclusive,
const proto::plan::GenericValue& upper_value,
bool upper_inclusive) {
AssertInfo(
start_offset >= 0 && length >= 0 &&
start_offset + length <= static_cast<int64_t>(column->NumRows()),
"binary range row id scan range [{}, {}) out of rows {}",
start_offset,
start_offset + length,
column->NumRows());
if (!CanUseFixedWidthBinaryRangeRowIdScan(
data_type, lower_value, upper_value)) {
return nullptr;
}
return std::make_unique<FixedWidthBinaryRangeRowIdScanCursor>(
column,
op_ctx,
start_offset,
length,
data_type,
lower_value,
lower_inclusive,
upper_value,
upper_inclusive);
}
inline std::unique_ptr<ChunkedColumnInterface::ScanCursor>
MakeFixedWidthDataScanCursor(const ChunkedColumnInterface* column,
milvus::OpContext* op_ctx,
+9 -7
View File
@@ -48,6 +48,8 @@ ThrowVortexStatus(const arrow::Status& status,
ErrorCode fallback_code,
std::string_view action) {
auto code = fallback_code;
// The Vortex bridge also uses IOError for decode failures, so the caller
// owns the fallback classification instead of mapping IOError globally.
if (status.IsOutOfMemory()) {
code = ErrorCode::MemAllocateFailed;
} else if (status.IsCancelled()) {
@@ -2261,6 +2263,13 @@ VortexColumn::BulkRawBsonAt(
const uint32_t* row_offsets,
const uint32_t* value_offsets,
int64_t count) const {
if (data_type_ != DataType::STRING) {
ThrowInfo(ErrorCode::Unsupported,
"VortexColumn::BulkRawBsonAt only supports BSON fields");
}
if (count == 0) {
return;
}
AssertInfo(row_offsets != nullptr && value_offsets != nullptr,
"row_offsets and value_offsets must be provided");
std::vector<int64_t> offsets(count);
@@ -2396,13 +2405,6 @@ VortexColumn::TakeOwn(milvus::OpContext* op_ctx,
return result;
}
std::shared_ptr<VortexColumn::TakeResult>
VortexColumn::Take(milvus::OpContext* op_ctx,
const int64_t* offsets,
int64_t count) const {
return std::make_shared<TakeResult>(TakeOwn(op_ctx, offsets, count));
}
ChunkedColumnInterface::ScanResult
VortexColumn::Scan(milvus::OpContext* op_ctx,
const ScanOptions& options) const {
+6 -11
View File
@@ -226,22 +226,12 @@ class VortexColumn final : public ChunkedColumnInterface {
ScanResult
Scan(milvus::OpContext* op_ctx, const ScanOptions& options) const override;
private:
struct TakeResult {
std::vector<std::shared_ptr<Chunk>> chunks;
std::vector<int64_t> offsets;
};
TakeResult
TakeOwn(milvus::OpContext* op_ctx,
const int64_t* offsets,
int64_t count) const;
std::shared_ptr<TakeResult>
Take(milvus::OpContext* op_ctx,
const int64_t* offsets,
int64_t count) const;
private:
std::optional<DataType>
GetDefaultScanDataType() const override;
@@ -354,6 +344,11 @@ class VortexColumn final : public ChunkedColumnInterface {
int64_t chunk_id,
const std::vector<int64_t>& offsets) const;
TakeResult
TakeOwn(milvus::OpContext* op_ctx,
const int64_t* offsets,
int64_t count) const;
ArrowTakeResult
TakeArrowFromFile(milvus::OpContext* op_ctx,
int64_t chunk_id,
@@ -42,6 +42,8 @@ ThrowVortexStatus(const arrow::Status& status,
ErrorCode fallback_code,
std::string_view action) {
auto code = fallback_code;
// The Vortex bridge also uses IOError for decode failures, so the caller
// owns the fallback classification instead of mapping IOError globally.
if (status.IsOutOfMemory()) {
code = ErrorCode::MemAllocateFailed;
} else if (status.IsCancelled()) {
@@ -1013,6 +1013,14 @@ TEST(VortexColumnTest, NullableAllScalarTypesScanCorrectness) {
CheckNoDataScan(column);
CheckApplyValidDataInChunk(column);
CheckDataScan(column, type);
if (type == DataType::STRING) {
EXPECT_NO_THROW(column.BulkRawBsonAt(
nullptr,
[](BsonView, uint32_t, uint32_t) {},
nullptr,
nullptr,
0));
}
if (IsVortexStringPushdownType(type)) {
CheckNullableFilteredScanReturnsValidity(column, type);
} else if (type == DataType::TEXT) {