Supporting export one embedding column to FVECS file (#1321)

### What problem does this PR solve?

SQL: COPY test_knn_ip(c2) TO '/var/infinity/embedding.FVECS' WITH
(DELIMITER ',', FORMAT fvecs);

### Type of change

- [x] New Feature (non-breaking change which adds functionality)

Signed-off-by: Jin Hai <haijin.chn@gmail.com>

Author: JinHai-CN

src/executor/operator/physical_export.cpp

@@ -33,6 +33,8 @@ import file_system_type;
 import defer_op;
 import stl;
 import logical_type;
+import embedding_info;
+import status;
 
 namespace infinity {
 
@@ -50,6 +52,10 @@ bool PhysicalExport::Execute(QueryContext *query_context, OperatorState *operato
             exported_row_count = ExportToJSONL(query_context, export_op_state);
             break;
         }
+        case CopyFileType::kFVECS: {
+            exported_row_count = ExportToFVECS(query_context, export_op_state);
+            break;
+        }
         default: {
             String error_message = "Not supported file type";
             LOG_CRITICAL(error_message);
@@ -192,4 +198,70 @@ SizeT PhysicalExport::ExportToJSONL(QueryContext *query_context, ExportOperatorS
     return row_count;
 }
 
+SizeT PhysicalExport::ExportToFVECS(QueryContext *query_context, ExportOperatorState *export_op_state) {
+
+    if(column_idx_array_.size() != 1) {
+        String error_message = "Only one column with embedding data type can be exported as FVECS file";
+        LOG_CRITICAL(error_message);
+        UnrecoverableError(error_message);
+    }
+
+    u64 exported_column_idx = column_idx_array_[0];
+    const Vector<SharedPtr<ColumnDef>>& column_defs = table_entry_->column_defs();
+    DataType* data_type = column_defs[exported_column_idx]->type().get();
+    if(data_type->type() != LogicalType::kEmbedding) {
+        String error_message = fmt::format("Only embedding column can be exported as FVECS file, but it is {}", data_type->ToString());
+        LOG_CRITICAL(error_message);
+        UnrecoverableError(error_message);
+    }
+
+    EmbeddingInfo* embedding_type_info = static_cast<EmbeddingInfo*>(data_type->type_info().get());
+    if(embedding_type_info->Type() != EmbeddingDataType::kElemFloat) {
+        Status status = Status::NotSupport("Only float element type embedding is supported now.");
+        LOG_ERROR(status.message());
+        RecoverableError(status);
+    }
+
+    i32 dimension = embedding_type_info->Dimension();
+
+    LocalFileSystem fs;
+    auto [file_handler, status] = fs.OpenFile(file_path_, FileFlags::WRITE_FLAG | FileFlags::CREATE_FLAG, FileLockType::kWriteLock);
+    if(!status.ok()) {
+        RecoverableError(status);
+    }
+    DeferFn file_defer([&]() { fs.Close(*file_handler); });
+
+    SizeT row_count{0};
+    Map<SegmentID, SegmentSnapshot>& segment_block_index_ref = block_index_->segment_block_index_;
+
+    // Write header
+    LOG_DEBUG(fmt::format("Going to export segment count: {}", segment_block_index_ref.size()));
+    for(auto& [segment_id, segment_snapshot]: segment_block_index_ref) {
+        SizeT block_count = segment_snapshot.block_map_.size();
+        LOG_DEBUG(fmt::format("Export segment_id: {}, with block count: {}", segment_id, block_count));
+        for(SizeT block_idx = 0; block_idx < block_count; ++ block_idx) {
+            LOG_DEBUG(fmt::format("Export block_idx: {}", block_idx));
+            BlockEntry *block_entry = segment_snapshot.block_map_[block_idx];
+            SizeT block_row_count = block_entry->row_count();
+
+            ColumnVector exported_column_vector = block_entry->GetColumnBlockEntry(exported_column_idx)->GetColumnVector(query_context->storage()->buffer_manager());
+            if(exported_column_vector.Size() != block_row_count) {
+                String error_message = "Unmatched row_count between block and block_column";
+                LOG_CRITICAL(error_message);
+                UnrecoverableError(error_message);
+            }
+
+            for(SizeT row_idx = 0; row_idx < block_row_count;  ++ row_idx) {
+                Value v = exported_column_vector.GetValue(row_idx);
+                Span<char> embedding = v.GetEmbedding();
+                fs.Write(*file_handler, &dimension, sizeof(dimension));
+                fs.Write(*file_handler, embedding.data(), embedding.size_bytes());
+                ++ row_count;
+            }
+        }
+    }
+    LOG_DEBUG(fmt::format("Export to FVECS, db {}, table {}, file: {}, row: {}", schema_name_, table_name_, file_path_, row_count));
+    return row_count;
+}
+
 } // namespace infinity

src/executor/operator/physical_export.cppm

@@ -43,10 +43,11 @@ public:
                             bool header,
                             char delimiter,
                             CopyFileType type,
+                            Vector<u64> column_idx_array,
                             SharedPtr<BlockIndex> block_index,
                             SharedPtr<Vector<LoadMeta>> load_metas)
         : PhysicalOperator(PhysicalOperatorType::kExport, nullptr, nullptr, id, load_metas), table_entry_(table_entry), file_type_(type), file_path_(std::move(file_path)),
-          table_name_(std::move(table_name)), schema_name_(std::move(schema_name)), header_(header), delimiter_(delimiter), block_index_(std::move(block_index)) {}
+          table_name_(std::move(table_name)), schema_name_(std::move(schema_name)), header_(header), delimiter_(delimiter), column_idx_array_(std::move(column_idx_array)), block_index_(std::move(block_index)) {}
 
     ~PhysicalExport() override = default;
 
@@ -69,6 +70,8 @@ public:
 
     SizeT ExportToJSONL(QueryContext *query_context, ExportOperatorState *export_op_state);
 
+    SizeT ExportToFVECS(QueryContext *query_context, ExportOperatorState *export_op_state);
+
     inline CopyFileType FileType() const { return file_type_; }
 
     inline const String &file_path() const { return file_path_; }
@@ -92,7 +95,7 @@ private:
     String schema_name_{"default_db"};
     bool header_{false};
     char delimiter_{','};
-
+    Vector<u64> column_idx_array_;
     SharedPtr<BlockIndex> block_index_{};
 };
 

src/executor/operator/physical_import.cpp

@@ -136,7 +136,7 @@ void PhysicalImport::ImportFVECS(QueryContext *query_context, ImportOperatorStat
     }
     DeferFn defer_fn([&]() { fs.Close(*file_handler); });
 
-    int dimension = 0;
+    i32 dimension = 0;
     i64 nbytes = fs.Read(*file_handler, &dimension, sizeof(dimension));
     fs.Seek(*file_handler, 0);
     if (nbytes == 0) {

src/executor/physical_planner.cpp

@@ -567,6 +567,7 @@ UniquePtr<PhysicalOperator> PhysicalPlanner::BuildExport(const SharedPtr<Logical
                                       logical_export->header(),
                                       logical_export->delimiter(),
                                       logical_export->FileType(),
+                                      logical_export->column_idx_array(),
                                       logical_export->block_index(),
                                       logical_operator->load_metas());
 }

src/parser/expression_lexer.cpp

@@ -1,6 +1,6 @@
-#line 1 "expression_lexer.cpp"
+#line 2 "expression_lexer.cpp"
 
-#line 3 "expression_lexer.cpp"
+#line 4 "expression_lexer.cpp"
 
 #define  YY_INT_ALIGNED short int
 
@@ -808,10 +808,10 @@ static const flex_int16_t yy_rule_linenum[28] =
 
 static thread_local std::stringstream string_buffer;
 
-#line 811 "expression_lexer.cpp"
+#line 812 "expression_lexer.cpp"
 #define YY_NO_INPUT 1
 
-#line 814 "expression_lexer.cpp"
+#line 815 "expression_lexer.cpp"
 
 #define INITIAL 0
 #define SINGLE_QUOTED_STRING 1
@@ -1165,7 +1165,7 @@ YY_DECL
 #line 27 "expression_lexer.l"
 
 
-#line 1168 "expression_lexer.cpp"
+#line 1169 "expression_lexer.cpp"
 
 	while ( /*CONSTCOND*/1 )		/* loops until end-of-file is reached */
 		{
@@ -1405,7 +1405,7 @@ YY_RULE_SETUP
 #line 88 "expression_lexer.l"
 ECHO;
 	YY_BREAK
-#line 1408 "expression_lexer.cpp"
+#line 1409 "expression_lexer.cpp"
 case YY_STATE_EOF(INITIAL):
 	yyterminate();
 

src/parser/expression_lexer.h

@@ -2,9 +2,9 @@
 #define expressionHEADER_H 1
 #define expressionIN_HEADER 1
 
-#line 5 "expression_lexer.h"
+#line 6 "expression_lexer.h"
 
-#line 7 "expression_lexer.h"
+#line 8 "expression_lexer.h"
 
 #define  YY_INT_ALIGNED short int
 
@@ -849,6 +849,6 @@ extern int yylex \
 #line 88 "expression_lexer.l"
 
 
-#line 852 "expression_lexer.h"
+#line 853 "expression_lexer.h"
 #undef expressionIN_HEADER
 #endif /* expressionHEADER_H */

src/parser/lexer.cpp

@@ -1,6 +1,6 @@
-#line 1 "lexer.cpp"
+#line 2 "lexer.cpp"
 
-#line 3 "lexer.cpp"
+#line 4 "lexer.cpp"
 
 #define  YY_INT_ALIGNED short int
 
@@ -1258,10 +1258,10 @@ static const flex_int16_t yy_rule_linenum[180] =
 
 static thread_local std::stringstream string_buffer;
 
-#line 1261 "lexer.cpp"
+#line 1262 "lexer.cpp"
 #define YY_NO_INPUT 1
 
-#line 1264 "lexer.cpp"
+#line 1265 "lexer.cpp"
 
 #define INITIAL 0
 #define SINGLE_QUOTED_STRING 1
@@ -1615,7 +1615,7 @@ YY_DECL
 #line 27 "lexer.l"
 
 
-#line 1618 "lexer.cpp"
+#line 1619 "lexer.cpp"
 
 	while ( /*CONSTCOND*/1 )		/* loops until end-of-file is reached */
 		{
@@ -2615,7 +2615,7 @@ YY_RULE_SETUP
 #line 240 "lexer.l"
 ECHO;
 	YY_BREAK
-#line 2618 "lexer.cpp"
+#line 2619 "lexer.cpp"
 case YY_STATE_EOF(INITIAL):
 	yyterminate();
 

src/parser/lexer.h

@@ -2,9 +2,9 @@
 #define sqlHEADER_H 1
 #define sqlIN_HEADER 1
 
-#line 5 "lexer.h"
+#line 6 "lexer.h"
 
-#line 7 "lexer.h"
+#line 8 "lexer.h"
 
 #define  YY_INT_ALIGNED short int
 
@@ -849,6 +849,6 @@ extern int yylex \
 #line 240 "lexer.l"
 
 
-#line 852 "lexer.h"
+#line 853 "lexer.h"
 #undef sqlIN_HEADER
 #endif /* sqlHEADER_H */