[Cpp]: fix reader bug (#118)

Signed-off-by: sunby <[email protected]>
milvus-io · Feb 6, 2024 · 9d1ad9c · 9d1ad9c
1 parent 88bc6a7
commit 9d1ad9c
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 4 deletions.
diff --git a/cpp/src/common/utils.cpp b/cpp/src/common/utils.cpp
@@ -15,6 +15,7 @@
 #include "common/utils.h"
 #include <arrow/type_fwd.h>
 #include <arrow/util/key_value_metadata.h>
+#include <cstdint>
 #include <iterator>
 #include <memory>
 #include <boost/uuid/random_generator.hpp>
@@ -40,7 +41,6 @@ ReadOptions CreateInternalReadOptions(std::shared_ptr<arrow::Schema> schema,
     internal_option.columns =
         std::set<std::string>(std::make_move_iterator(field_names.begin()), std::make_move_iterator(field_names.end()));
   } else {
-    internal_option.columns.insert(schema_options.primary_column);
     if (schema_options.has_version_column()) {
       internal_option.columns.insert(schema_options.version_column);
     }

diff --git a/cpp/src/reader/common/combine_reader.cpp b/cpp/src/reader/common/combine_reader.cpp
@@ -15,6 +15,7 @@
 #include "reader/common/combine_reader.h"
 #include <memory>
 #include "arrow/type.h"
+#include "common/constants.h"
 namespace milvus_storage {
 std::unique_ptr<CombineReader> CombineReader::Make(std::unique_ptr<arrow::RecordBatchReader> scalar_reader,
                                                    std::unique_ptr<arrow::RecordBatchReader> vector_reader,
@@ -35,6 +36,13 @@ arrow::Status CombineReader::ReadNext(std::shared_ptr<arrow::RecordBatch>* batch
     return arrow::Status::OK();
   }
 
+  for (int i = 0; i < scalar_batch->num_columns(); ++i) {
+    if (scalar_batch->column_name(i) == kOffsetFieldName) {
+      scalar_batch->RemoveColumn(i);
+      break;
+    }
+  }
+
   assert(scalar_batch->num_rows() == vector_batch->num_rows());
 
   auto vec_column = vector_batch->GetColumnByName(schema_->options().vector_column);

diff --git a/cpp/src/storage/schema.cpp b/cpp/src/storage/schema.cpp
@@ -14,6 +14,7 @@
 
 #include "storage/schema.h"
 #include <memory>
+#include "common/constants.h"
 #include "common/macro.h"
 #include "common/utils.h"
 #include "common/log.h"
@@ -68,7 +69,7 @@ Status Schema::BuildScalarSchema() {
     }
     RETURN_ARROW_NOT_OK(scalar_schema_builder.AddField(field));
   }
-  auto offset_field = std::make_shared<arrow::Field>("off_set", arrow::int64());
+  auto offset_field = std::make_shared<arrow::Field>(kOffsetFieldName, arrow::int64());
   RETURN_ARROW_NOT_OK(scalar_schema_builder.AddField(offset_field));
   ASSIGN_OR_RETURN_ARROW_NOT_OK(scalar_schema_, scalar_schema_builder.Finish());
   return Status::OK();

diff --git a/cpp/test/schema_test.cpp b/cpp/test/schema_test.cpp
@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include "storage/schema.h"
+#include "common/constants.h"
 #include "gtest/gtest.h"
 #include "storage/options.h"
 #include "test_util.h"
@@ -55,7 +56,7 @@ TEST(SchemaValidateTest, SchemaValidateNoVersionColTest) {
   /// scalar schema has no version column but has offset column
   ASSERT_EQ(scalar_schema->num_fields(), 2);
   ASSERT_EQ(scalar_schema->field(0)->name(), schema_options.primary_column);
-  ASSERT_EQ(scalar_schema->field(1)->name(), "off_set");
+  ASSERT_EQ(scalar_schema->field(1)->name(), kOffsetFieldName);
 
   auto vector_schema = space_schema1->vector_schema();
   ASSERT_EQ(vector_schema->num_fields(), 2);
@@ -106,7 +107,7 @@ TEST(SchemaValidateTest, SchemaValidateVersionColTest) {
   ASSERT_EQ(scalar_schema->num_fields(), 3);
   ASSERT_EQ(scalar_schema->field(0)->name(), schema_options.primary_column);
   ASSERT_EQ(scalar_schema->field(1)->name(), schema_options.version_column);
-  ASSERT_EQ(scalar_schema->field(2)->name(), "off_set");
+  ASSERT_EQ(scalar_schema->field(2)->name(), kOffsetFieldName);
 
   auto vector_schema = space_schema1->vector_schema();
   ASSERT_EQ(vector_schema->num_fields(), 3);