Improve copy-to-parquet perf

src/include/processor/operator/persistent/copy_to_parquet.h

@@ -13,16 +13,18 @@ struct CopyToParquetInfo final : public CopyToInfo {
         kuzu_parquet::format::CompressionCodec::SNAPPY;
     std::unique_ptr<FactorizedTableSchema> tableSchema;
     std::vector<std::unique_ptr<common::LogicalType>> types;
+    DataPos countingVecPos;
 
     CopyToParquetInfo(std::unique_ptr<FactorizedTableSchema> tableSchema,
         std::vector<std::unique_ptr<common::LogicalType>> types, std::vector<std::string> names,
-        std::vector<DataPos> dataPoses, std::string fileName)
+        std::vector<DataPos> dataPoses, std::string fileName, DataPos countingVecPos)
         : CopyToInfo{std::move(names), std::move(dataPoses), std::move(fileName)},
-          tableSchema{std::move(tableSchema)}, types{std::move(types)} {}
+          tableSchema{std::move(tableSchema)}, types{std::move(types)}, countingVecPos{std::move(
+                                                                            countingVecPos)} {}
 
-    inline std::unique_ptr<CopyToInfo> copy() override {
-        return std::make_unique<CopyToParquetInfo>(
-            tableSchema->copy(), common::LogicalType::copy(types), names, dataPoses, fileName);
+    std::unique_ptr<CopyToInfo> copy() override {
+        return std::make_unique<CopyToParquetInfo>(tableSchema->copy(),
+            common::LogicalType::copy(types), names, dataPoses, fileName, countingVecPos);
     }
 };
 
@@ -35,8 +37,10 @@ class CopyToParquetLocalState final : public CopyToLocalState {
 
 private:
     std::unique_ptr<FactorizedTable> ft;
+    uint64_t numTuplesInFT;
     std::vector<common::ValueVector*> vectorsToAppend;
     storage::MemoryManager* mm;
+    common::ValueVector* countingVec;
 };
 
 class CopyToParquetSharedState final : public CopyToSharedState {

src/processor/map/map_copy_to.cpp

@@ -18,21 +18,24 @@ std::unique_ptr<CopyToInfo> getCopyToInfo(Schema* childSchema, std::string fileP
     case FileType::PARQUET: {
         auto copyToSchema = std::make_unique<FactorizedTableSchema>();
         auto copyToExpressions = childSchema->getExpressionsInScope();
+        auto countingVecPos = DataPos(childSchema->getExpressionPos(*copyToExpressions[0]));
         for (auto& copyToExpression : copyToExpressions) {
             auto [dataChunkPos, vectorPos] = childSchema->getExpressionPos(*copyToExpression);
             std::unique_ptr<ColumnSchema> columnSchema;
             if (!childSchema->getGroup(dataChunkPos)->isFlat()) {
                 // payload is unFlat and not in the same group as keys
                 columnSchema = std::make_unique<ColumnSchema>(
                     true /* isUnFlat */, dataChunkPos, sizeof(overflow_value_t));
+                countingVecPos = DataPos(childSchema->getExpressionPos(*copyToExpression));
             } else {
                 columnSchema = std::make_unique<ColumnSchema>(false /* isUnFlat */, dataChunkPos,
                     LogicalTypeUtils::getRowLayoutSize(copyToExpression->getDataType()));
             }
             copyToSchema->appendColumn(std::move(columnSchema));
         }
         return std::make_unique<CopyToParquetInfo>(std::move(copyToSchema), std::move(columnsTypes),
-            std::move(columnNames), std::move(vectorsToCopyPos), std::move(filePath));
+            std::move(columnNames), std::move(vectorsToCopyPos), std::move(filePath),
+            std::move(countingVecPos));
     }
     case FileType::CSV: {
         return std::make_unique<CopyToCSVInfo>(std::move(columnNames), std::move(vectorsToCopyPos),

src/processor/operator/persistent/copy_to_parquet.cpp

@@ -10,7 +10,10 @@
     CopyToInfo* info, storage::MemoryManager* mm, ResultSet* resultSet) {
     auto copyToInfo = reinterpret_cast<CopyToParquetInfo*>(info);
     ft = std::make_unique<FactorizedTable>(mm, std::move(copyToInfo->tableSchema));
+    numTuplesInFT = 0;
+    countingVec = nullptr;
     vectorsToAppend.reserve(info->dataPoses.size());
+    countingVec = resultSet->getValueVector(copyToInfo->countingVecPos).get();
     for (auto& pos : info->dataPoses) {
         vectorsToAppend.push_back(resultSet->getValueVector(pos).get());
     }
@@ -19,8 +22,10 @@
 
 void CopyToParquetLocalState::sink(CopyToSharedState* sharedState, CopyToInfo* /*info*/) {
     ft->append(vectorsToAppend);
-    if (ft->getTotalNumFlatTuples() > StorageConstants::NODE_GROUP_SIZE) {
+    numTuplesInFT += countingVec->state->selVector->selectedSize;
+    if (numTuplesInFT > StorageConstants::NODE_GROUP_SIZE) {
         reinterpret_cast<CopyToParquetSharedState*>(sharedState)->flush(*ft);
+        numTuplesInFT = 0;
     }
 }