fix: avoid compare fast path for misaligned chunks (#50609)

#49346

Signed-off-by: luzhang <luzhang@zilliz.com>
Co-authored-by: luzhang <luzhang@zilliz.com>

internal/core/src/exec/expression/CompareExpr.cpp

@@ -35,6 +35,47 @@ PhyCompareFilterExpr::IsStringExpr() {
            expr_->right_data_type_ == DataType::VARCHAR;
 }
 
+bool
+PhyCompareFilterExpr::CanUseBothDataFastPath() {
+    if (is_left_indexed_ || is_right_indexed_ || IsStringExpr()) {
+        return false;
+    }
+
+    // Offset-input path resolves each field's chunk from the row offset
+    // independently, so it does not require left/right chunk boundaries to
+    // align.
+    if (has_offset_input_) {
+        return true;
+    }
+
+    if (can_use_both_data_sequential_fast_path_.has_value()) {
+        return can_use_both_data_sequential_fast_path_.value();
+    }
+
+    auto segment = segment_chunk_reader_.segment_;
+    if (!segment->is_chunked() || segment->type() == SegmentType::Growing) {
+        can_use_both_data_sequential_fast_path_ = true;
+        return true;
+    }
+
+    auto left_chunks = segment->num_chunk_data(left_field_);
+    auto right_chunks = segment->num_chunk_data(right_field_);
+    if (left_chunks <= 0 || right_chunks <= 0 || left_chunks != right_chunks) {
+        can_use_both_data_sequential_fast_path_ = false;
+        return false;
+    }
+
+    for (int64_t i = 0; i <= left_chunks; ++i) {
+        if (segment->num_rows_until_chunk(left_field_, i) !=
+            segment->num_rows_until_chunk(right_field_, i)) {
+            can_use_both_data_sequential_fast_path_ = false;
+            return false;
+        }
+    }
+    can_use_both_data_sequential_fast_path_ = true;
+    return true;
+}
+
 int64_t
 PhyCompareFilterExpr::GetNextBatchSize() {
     auto current_rows = GetCurrentRows();
@@ -225,7 +266,7 @@ PhyCompareFilterExpr::Eval(EvalCtx& context, VectorPtr& result) {
     SetHasOffsetInput((input != nullptr));
     // For segment both fields has no index, can use SIMD to speed up.
     // Avoiding too much call stack that blocks SIMD.
-    if (!is_left_indexed_ && !is_right_indexed_ && !IsStringExpr()) {
+    if (CanUseBothDataFastPath()) {
         result = ExecCompareExprDispatcherForBothDataSegment(context);
         return;
     }

internal/core/src/exec/expression/CompareExpr.h

@@ -302,6 +302,9 @@ class PhyCompareFilterExpr : public Expr {
     bool
     IsStringExpr();
 
+    bool
+    CanUseBothDataFastPath();
+
     template <typename T, typename U, typename FUNC, typename... ValTypes>
     int64_t
     ProcessBothDataChunks(FUNC func,
@@ -606,6 +609,7 @@ class PhyCompareFilterExpr : public Expr {
     int64_t right_current_chunk_pos_{0};
     int64_t current_chunk_id_{0};
     int64_t current_chunk_pos_{0};
+    std::optional<bool> can_use_both_data_sequential_fast_path_;
 
     const segcore::SegmentChunkReader segment_chunk_reader_;
     int64_t batch_size_;

internal/core/src/segcore/ChunkedSegmentSealedStorageV2Test.cpp

@@ -66,6 +66,7 @@
 #include "segcore/SegmentChunkReader.h"
 #include "segcore/SegmentSealed.h"
 #include "segcore/Types.h"
+#include "segcore/storagev2translator/GroupCTMeta.h"
 #include "segcore/storagev1translator/ChunkTranslator.h"
 #include "storage/FileManager.h"
 #include "storage/Types.h"
@@ -166,6 +167,21 @@ class RawLookupOnlyIndex : public index::ScalarIndex<int64_t> {
 
     mutable size_t last_lookup_offset = 0;
 };
+
+class StorageV2CellTargetGuard {
+ public:
+    explicit StorageV2CellTargetGuard(int64_t bytes)
+        : old_bytes_(segcore::storagev2translator::GetCellTargetSizeBytes()) {
+        segcore::storagev2translator::SetCellTargetSizeBytes(bytes);
+    }
+
+    ~StorageV2CellTargetGuard() {
+        segcore::storagev2translator::SetCellTargetSizeBytes(old_bytes_);
+    }
+
+ private:
+    int64_t old_bytes_;
+};
 }  // namespace
 
 class TestChunkSegmentStorageV2 : public testing::TestWithParam<bool> {
@@ -560,6 +576,132 @@ TEST_P(TestChunkSegmentStorageV2, TestCompareExpr) {
     ASSERT_EQ(chunk_num * test_data_count, final.count());
 }
 
+TEST(TestChunkSegmentStorageV2Regression,
+     TestCompareExprFallsBackWhenColumnGroupChunksAreMisaligned) {
+    StorageV2CellTargetGuard cell_target_guard(1 * 1024 * 1024);
+
+    auto schema = std::make_shared<Schema>();
+    auto left_fid = schema->AddDebugField("left", DataType::INT64, false);
+    auto right_fid = schema->AddDebugField("right", DataType::INT64, false);
+    schema->AddDebugField("payload", DataType::VARCHAR, false);
+    schema->AddField(FieldName("ts"),
+                     TimestampFieldID,
+                     DataType::INT64,
+                     false,
+                     std::nullopt);
+    schema->set_primary_field_id(right_fid);
+
+    auto fs = milvus::segcore::GetDefaultArrowFileSystem();
+    const std::string root = "test_compare_expr_misaligned_storage_v2";
+    auto cleanup_status = fs->DeleteDir(root);
+    (void)cleanup_status;
+    ASSERT_TRUE(fs->CreateDir(root + "/0").ok());
+    ASSERT_TRUE(
+        fs->CreateDir(root + "/" + std::to_string(right_fid.get())).ok());
+
+    std::vector<std::string> paths = {
+        root + "/0/10000.parquet",
+        root + "/" + std::to_string(right_fid.get()) + "/10001.parquet"};
+    std::vector<std::vector<int>> column_groups = {{0, 2, 3}, {1}};
+    auto storage_config = milvus_storage::StorageConfig();
+    auto result = milvus_storage::PackedRecordBatchWriter::Make(
+        fs,
+        paths,
+        schema->ConvertToArrowSchema(),
+        storage_config,
+        column_groups,
+        16 * 1024 * 1024,
+        ::parquet::default_writer_properties());
+    ASSERT_TRUE(result.ok());
+    auto writer = result.ValueOrDie();
+
+    constexpr int64_t rows_per_batch = 10000;
+    constexpr int64_t batch_count = 2;
+    auto arrow_schema = schema->ConvertToArrowSchema();
+    for (int64_t batch = 0; batch < batch_count; ++batch) {
+        std::vector<std::shared_ptr<arrow::Array>> arrays;
+        auto start = batch * rows_per_batch;
+        for (int i = 0; i < arrow_schema->fields().size(); ++i) {
+            if (arrow_schema->fields()[i]->type()->id() == arrow::Type::INT64) {
+                std::vector<int64_t> values(rows_per_batch);
+                std::iota(values.begin(), values.end(), start);
+                arrow::Int64Builder builder;
+                ASSERT_TRUE(
+                    builder.AppendValues(values.data(), rows_per_batch).ok());
+                std::shared_ptr<arrow::Array> array;
+                ASSERT_TRUE(builder.Finish(&array).ok());
+                arrays.push_back(array);
+            } else {
+                arrow::StringBuilder builder;
+                std::vector<std::string> values;
+                values.reserve(rows_per_batch);
+                for (int64_t row = 0; row < rows_per_batch; ++row) {
+                    values.emplace_back(2048, 'x');
+                }
+                ASSERT_TRUE(builder.AppendValues(values).ok());
+                std::shared_ptr<arrow::Array> array;
+                ASSERT_TRUE(builder.Finish(&array).ok());
+                arrays.push_back(array);
+            }
+        }
+
+        auto record_batch =
+            arrow::RecordBatch::Make(arrow_schema, rows_per_batch, arrays);
+        ASSERT_TRUE(writer->Write(record_batch).ok());
+    }
+    ASSERT_TRUE(writer->Close().ok());
+
+    const int64_t row_count = rows_per_batch * batch_count;
+    LoadFieldDataInfo load_info;
+    load_info.storage_version = 2;
+    load_info.field_infos.emplace(
+        int64_t(0),
+        FieldBinlogInfo{int64_t(0),
+                        row_count,
+                        std::vector<int64_t>(row_count),
+                        std::vector<int64_t>(row_count * 4),
+                        false,
+                        "",
+                        std::vector<std::string>({paths[0]})});
+    load_info.field_infos.emplace(
+        right_fid.get(),
+        FieldBinlogInfo{right_fid.get(),
+                        row_count,
+                        std::vector<int64_t>(row_count),
+                        std::vector<int64_t>(row_count * 4),
+                        false,
+                        "",
+                        std::vector<std::string>({paths[1]})});
+
+    auto segment = segcore::CreateSealedSegment(
+        schema, nullptr, -1, segcore::SegcoreConfig::default_config(), true);
+    segment->AddFieldDataInfoForSealed(load_info);
+    for (auto& [id, info] : load_info.field_infos) {
+        LoadFieldDataInfo one_field;
+        one_field.storage_version = 2;
+        one_field.field_infos.emplace(id, info);
+        segment->LoadFieldData(one_field);
+    }
+
+    ASSERT_GT(segment->num_chunk_data(left_fid),
+              segment->num_chunk_data(right_fid));
+
+    auto expr =
+        std::make_shared<expr::CompareExpr>(left_fid,
+                                            right_fid,
+                                            milvus::DataType::INT64,
+                                            milvus::DataType::INT64,
+                                            proto::plan::OpType::GreaterEqual);
+    auto plan =
+        std::make_shared<plan::FilterBitsNode>(DEFAULT_PLANNODE_ID, expr);
+
+    auto final =
+        query::ExecuteQueryExpr(plan, segment.get(), row_count, MAX_TIMESTAMP);
+    ASSERT_EQ(row_count, final.count());
+
+    ASSERT_TRUE(fs->DeleteDir(root).ok());
+}
+
 TEST_P(TestChunkSegmentStorageV2, TestColumnExprWithScalarIndexRawData) {
     LoadInt64ScalarIndex(index::ASCENDING_SORT);
     ASSERT_TRUE(segment->HasRawData(fields.at("int64").get()));