Added progress for in_query_call operators

src/common/task_system/progress_bar.cpp

@@ -30,17 +30,23 @@
         return;
     }
     numPipelinesFinished++;
+    if (printing) {
+        std::cout << "\033[1A\033[2K\033[1B";
+    }
     // This ensures that the progress bar is updated back to 0% after a pipeline is finished.
     prevCurPipelineProgress = -0.01;
     updateProgress(0.0);
 }
 
 void ProgressBar::updateProgress(double curPipelineProgress) {
     // Only update the progress bar if the progress has changed by at least 1%.
-    if (!trackProgress || curPipelineProgress - prevCurPipelineProgress < 0.01) {
+    if (!trackProgress) {
         return;
     }
     std::lock_guard<std::mutex> lock(progressBarLock);
+    if (curPipelineProgress - prevCurPipelineProgress < 0.01) {
+        return;
+    }
     prevCurPipelineProgress = curPipelineProgress;
     if (printing) {
         std::cout << "\033[2A";

src/include/function/table/scan_functions.h

@@ -30,10 +30,11 @@ struct ScanSharedState : public BaseScanSharedState {
 
 struct ScanFileSharedState : public ScanSharedState {
     main::ClientContext* context;
+    uint64_t totalSize;
 
     ScanFileSharedState(
         common::ReaderConfig readerConfig, uint64_t numRows, main::ClientContext* context)
-        : ScanSharedState{std::move(readerConfig), numRows}, context{context} {}
+        : ScanSharedState{std::move(readerConfig), numRows}, context{context}, totalSize{0} {}
 };
 
 } // namespace function

src/include/function/table_functions.h

@@ -63,13 +63,15 @@ using table_func_init_shared_t =
 using table_func_init_local_t = std::function<std::unique_ptr<TableFuncLocalState>(
     TableFunctionInitInput&, TableFuncSharedState*, storage::MemoryManager*)>;
 using table_func_can_parallel_t = std::function<bool()>;
+using table_func_progress_t = std::function<double(TableFuncSharedState* sharedState)>;
 
 struct TableFunction final : public Function {
     table_func_t tableFunc;
     table_func_bind_t bindFunc;
     table_func_init_shared_t initSharedStateFunc;
     table_func_init_local_t initLocalStateFunc;
     table_func_can_parallel_t canParallelFunc = [] { return true; };
+    table_func_progress_t progressFunc = [](TableFuncSharedState* /*sharedState*/) { return 0.0; };
 
     TableFunction()
         : Function{}, tableFunc{nullptr}, bindFunc{nullptr}, initSharedStateFunc{nullptr},
@@ -80,6 +82,12 @@ struct TableFunction final : public Function {
         : Function{FunctionType::TABLE, std::move(name), std::move(inputTypes)},
           tableFunc{tableFunc}, bindFunc{bindFunc}, initSharedStateFunc{initSharedFunc},
           initLocalStateFunc{initLocalFunc} {}
+    TableFunction(std::string name, table_func_t tableFunc, table_func_bind_t bindFunc,
+        table_func_init_shared_t initSharedFunc, table_func_init_local_t initLocalFunc,
+        table_func_progress_t progressFunc, std::vector<common::LogicalTypeID> inputTypes)
+        : Function{FunctionType::TABLE, std::move(name), std::move(inputTypes)},
+          tableFunc{tableFunc}, bindFunc{bindFunc}, initSharedStateFunc{initSharedFunc},
+          initLocalStateFunc{initLocalFunc}, progressFunc{progressFunc} {}
 
     inline std::string signatureToString() const override {
         return common::LogicalTypeUtils::toString(parameterTypeIDs);

src/include/processor/operator/call/in_query_call.h

@@ -74,6 +74,8 @@ class InQueryCall : public PhysicalOperator {
 
     bool getNextTuplesInternal(ExecutionContext* context) override;
 
+    double getProgress(ExecutionContext* context) const override;
+
     std::unique_ptr<PhysicalOperator> clone() override {
         return std::make_unique<InQueryCall>(info.copy(), sharedState, id, paramsString);
     }

src/include/processor/operator/persistent/reader/csv/base_csv_reader.h

@@ -26,6 +26,9 @@ class BaseCSVReader {
 
     uint64_t countRows();
     bool isEOF() const;
+    uint64_t getFileSize();
+    // Get the file offset of the current buffer position.
+    uint64_t getFileOffset() const;
 
 protected:
     template<typename Driver>
@@ -56,8 +59,6 @@ class BaseCSVReader {
 
     inline bool isNewLine(char c) { return c == '\n' || c == '\r'; }
 
-    // Get the file offset of the current buffer position.
-    uint64_t getFileOffset() const;
     uint64_t getLineNumber();
 
 protected:

src/include/processor/operator/persistent/reader/parquet/parquet_reader.h

@@ -95,6 +95,7 @@ struct ParquetScanSharedState final : public function::ScanFileSharedState {
         const common::ReaderConfig readerConfig, uint64_t numRows, main::ClientContext* context);
 
     std::vector<std::unique_ptr<ParquetReader>> readers;
+    uint64_t totalRowsGroups;
 };
 
 struct ParquetScanLocalState final : public function::TableFuncLocalState {

src/processor/operator/call/in_query_call.cpp

@@ -65,5 +65,9 @@ bool InQueryCall::getNextTuplesInternal(ExecutionContext*) {
     return numTuplesScanned != 0;
 }
 
+double InQueryCall::getProgress(ExecutionContext* /*context*/) const {
+    return info.function.progressFunc(sharedState->funcState.get());
+}
+
 } // namespace processor
 } // namespace kuzu

src/processor/operator/persistent/reader/csv/base_csv_reader.cpp

@@ -114,6 +114,10 @@ bool BaseCSVReader::isEOF() const {
     return getFileOffset() >= fileInfo->getFileSize();
 }
 
+uint64_t BaseCSVReader::getFileSize() {
+    return fileInfo->getFileSize();
+}
+
 template<typename Driver>
 void BaseCSVReader::addValue(Driver& driver, uint64_t rowNum, column_id_t columnIdx,
     std::string_view strVal, std::vector<uint64_t>& escapePositions) {

src/processor/operator/persistent/reader/csv/parallel_csv_reader.cpp

@@ -168,8 +168,15 @@ static std::unique_ptr<TableFuncSharedState> initSharedState(TableFunctionInitIn
     auto bindData = reinterpret_cast<ScanBindData*>(input.bindData);
     auto csvConfig = CSVReaderConfig::construct(bindData->config.options);
     row_idx_t numRows = 0;
-    return std::make_unique<ParallelCSVScanSharedState>(bindData->config.copy(), numRows,
-        bindData->columnNames.size(), bindData->context, csvConfig.copy());
+    auto sharedState = std::make_unique<ParallelCSVScanSharedState>(bindData->config.copy(),
+        numRows, bindData->columnNames.size(), bindData->context, csvConfig.copy());
+    for (auto filePath : sharedState->readerConfig.filePaths) {
+        auto reader = std::make_unique<ParallelCSVReader>(filePath,
+            sharedState->csvReaderConfig.option.copy(), sharedState->numColumns,
+            sharedState->context);
+        sharedState->totalSize += reader->getFileSize();
+    }
+    return sharedState;
 }
 
 static std::unique_ptr<TableFuncLocalState> initLocalState(TableFunctionInitInput& /*input*/,
@@ -182,11 +189,24 @@ static std::unique_ptr<TableFuncLocalState> initLocalState(TableFunctionInitInpu
     return localState;
 }
 
+static double progressFunc(TableFuncSharedState* sharedState) {
+    auto state = reinterpret_cast<ParallelCSVScanSharedState*>(sharedState);
+    if (state->totalSize == 0) {
+        return 0.0;
+    } else if ((state->blockIdx * CopyConstants::PARALLEL_BLOCK_SIZE) * (state->fileIdx + 1) >=
+               state->totalSize) {
+        return 1.0;
+    }
+    return static_cast<double>(
+               (state->blockIdx * CopyConstants::PARALLEL_BLOCK_SIZE) * (state->fileIdx + 1)) /
+           state->totalSize;
+}
+
 function_set ParallelCSVScan::getFunctionSet() {
     function_set functionSet;
-    functionSet.push_back(
-        std::make_unique<TableFunction>(READ_CSV_PARALLEL_FUNC_NAME, tableFunc, bindFunc,
-            initSharedState, initLocalState, std::vector<LogicalTypeID>{LogicalTypeID::STRING}));
+    functionSet.push_back(std::make_unique<TableFunction>(READ_CSV_PARALLEL_FUNC_NAME, tableFunc,
+        bindFunc, initSharedState, initLocalState, progressFunc,
+        std::vector<LogicalTypeID>{LogicalTypeID::STRING}));
     return functionSet;
 }
 

src/processor/operator/persistent/reader/csv/serial_csv_reader.cpp

@@ -115,20 +115,42 @@
     auto bindData = reinterpret_cast<ScanBindData*>(input.bindData);
     auto csvConfig = CSVReaderConfig::construct(bindData->config.options);
     row_idx_t numRows = 0;
-    return std::make_unique<SerialCSVScanSharedState>(bindData->config.copy(), numRows,
+    auto sharedState = std::make_unique<SerialCSVScanSharedState>(bindData->config.copy(), numRows,
         bindData->columnNames.size(), csvConfig.copy(), bindData->context);
+    for (auto filePath : sharedState->readerConfig.filePaths) {
+        auto reader =
+            std::make_unique<SerialCSVReader>(filePath, sharedState->csvReaderConfig.option.copy(),
+                sharedState->numColumns, sharedState->context);
+        sharedState->totalSize += reader->getFileSize();
+    }
+    return sharedState;
 }
 
 static std::unique_ptr<TableFuncLocalState> initLocalState(TableFunctionInitInput& /*input*/,
     TableFuncSharedState* /*state*/, storage::MemoryManager* /*mm*/) {
     return std::make_unique<TableFuncLocalState>();
 }
 
+static double progressFunc(TableFuncSharedState* sharedState) {
+    auto state = reinterpret_cast<SerialCSVScanSharedState*>(sharedState);
+    if (state->totalSize == 0) {
+        return 0.0;
+    } else if (state->fileIdx >= state->readerConfig.getNumFiles()) {
+        return 1.0;
+    }
+    uint64_t totalReadSize = 0;
+    for (auto i = 0u; i < state->fileIdx; i++) {
+        totalReadSize += state->reader->getFileSize();
+    }
+    totalReadSize += state->reader->getFileOffset();
+    return static_cast<double>(totalReadSize) / state->totalSize;
+}
+
 function_set SerialCSVScan::getFunctionSet() {
     function_set functionSet;
-    functionSet.push_back(
-        std::make_unique<TableFunction>(READ_CSV_SERIAL_FUNC_NAME, tableFunc, bindFunc,
-            initSharedState, initLocalState, std::vector<LogicalTypeID>{LogicalTypeID::STRING}));
+    functionSet.push_back(std::make_unique<TableFunction>(READ_CSV_SERIAL_FUNC_NAME, tableFunc,
+        bindFunc, initSharedState, initLocalState, progressFunc,
+        std::vector<LogicalTypeID>{LogicalTypeID::STRING}));
     return functionSet;
 }
 

src/processor/operator/persistent/reader/parquet/parquet_reader.cpp

@@ -563,6 +563,11 @@
     : ScanFileSharedState{std::move(readerConfig), numRows, context} {
     readers.push_back(
         std::make_unique<ParquetReader>(this->readerConfig.filePaths[fileIdx], context));
+    totalRowsGroups = 0;
+    for (auto i = 0u; i < readerConfig.getNumFiles(); i++) {
+        auto reader = std::make_unique<ParquetReader>(readerConfig.filePaths[i], context);
+        totalRowsGroups += reader->getNumRowsGroups();
+    }
 }
 
 static bool parquetSharedStateNext(
@@ -674,11 +679,23 @@
     return localState;
 }
 
+static double progressFunc(TableFuncSharedState* state) {
+    auto parquetScanSharedState = reinterpret_cast<ParquetScanSharedState*>(state);
+    if (parquetScanSharedState->fileIdx >= parquetScanSharedState->readerConfig.getNumFiles()) {
+        return 1.0;
+    }
+    if (parquetScanSharedState->totalRowsGroups == 0) {
+        return 0.0;
+    }
+    return (double)(parquetScanSharedState->blockIdx * (parquetScanSharedState->fileIdx + 1)) /
+           parquetScanSharedState->totalRowsGroups;
+}
+
 function_set ParquetScanFunction::getFunctionSet() {
     function_set functionSet;
-    functionSet.push_back(
-        std::make_unique<TableFunction>(READ_PARQUET_FUNC_NAME, tableFunc, bindFunc,
-            initSharedState, initLocalState, std::vector<LogicalTypeID>{LogicalTypeID::STRING}));
+    functionSet.push_back(std::make_unique<TableFunction>(READ_PARQUET_FUNC_NAME, tableFunc,
+        bindFunc, initSharedState, initLocalState, progressFunc,
+        std::vector<LogicalTypeID>{LogicalTypeID::STRING}));
     return functionSet;
 }
 

tools/python_api/src_cpp/include/pandas/pandas_scan.h

@@ -17,10 +17,12 @@ struct PandasScanLocalState : public function::TableFuncLocalState {
 };
 
 struct PandasScanSharedState : public function::BaseScanSharedState {
-    explicit PandasScanSharedState(uint64_t numRows) : BaseScanSharedState{numRows}, position{0} {}
+    explicit PandasScanSharedState(uint64_t numRows)
+        : BaseScanSharedState{numRows}, position{0}, numReadRows{0} {}
 
     std::mutex lock;
     uint64_t position;
+    uint64_t numReadRows;
 };
 
 struct PandasScanFunction {

tools/python_api/src_cpp/pandas/pandas_scan.cpp

@@ -83,6 +83,7 @@ offset_t tableFunc(
     TableFuncInput& input, TableFuncOutput& output) {
     auto pandasScanData = reinterpret_cast<PandasScanFunctionData*>(input.bindData);
     auto pandasLocalState = reinterpret_cast<PandasScanLocalState*>(input.localState);
+    auto pandasSharedState = reinterpret_cast<PandasScanSharedState*>(input.sharedState);
 
     if (pandasLocalState->start >= pandasLocalState->end) {
         if (!sharedStateNext(input.bindData, pandasLocalState, input.sharedState)) {
@@ -97,6 +98,7 @@ offset_t tableFunc(
     }
     output.dataChunk.state->selVector->selectedSize = numValuesToOutput;
     pandasLocalState->start += numValuesToOutput;
+    pandasSharedState->numReadRows += numValuesToOutput;
     return numValuesToOutput;
 }
 
@@ -109,9 +111,17 @@ std::vector<std::unique_ptr<PandasColumnBindData>> PandasScanFunctionData::copyC
     return result;
 }
 
+static double progressFunc(TableFuncSharedState* sharedState) {
+	auto pandasSharedState = reinterpret_cast<PandasScanSharedState*>(sharedState);
+    if (pandasSharedState->numRows == 0) {
+		return 0.0;
+	}
+	return static_cast<double>(pandasSharedState->numReadRows) / pandasSharedState->numRows;
+}
+
 static TableFunction getFunction() {
     return TableFunction(READ_PANDAS_FUNC_NAME, tableFunc, bindFunc, initSharedState,
-        initLocalState, std::vector<LogicalTypeID>{LogicalTypeID::POINTER});
+        initLocalState, progressFunc, std::vector<LogicalTypeID>{LogicalTypeID::POINTER});
 }
 
 function_set PandasScanFunction::getFunctionSet() {

tools/python_api/src_cpp/pyarrow/pyarrow_scan.cpp

@@ -85,18 +85,27 @@
     return len;
 }
 
+double progressFunc(function::TableFuncSharedState* sharedState) {
+	PyArrowTableScanSharedState* state =
+		dynamic_cast<PyArrowTableScanSharedState*>(sharedState);
+    if (state->chunks.size() == 0) {
+		return 0.0;
+	}
+	return static_cast<double>(state->currentChunk) / state->chunks.size();
+}   
+
 function::function_set PyArrowTableScanFunction::getFunctionSet() {
 
     function_set functionSet;
     functionSet.push_back(
         std::make_unique<TableFunction>(READ_PYARROW_FUNC_NAME, tableFunc, bindFunc,
-            initSharedState, initLocalState, std::vector<LogicalTypeID>{LogicalTypeID::POINTER}));
+            initSharedState, initLocalState, progressFunc, std::vector<LogicalTypeID>{LogicalTypeID::POINTER}));
     return functionSet;
 }
 
 TableFunction PyArrowTableScanFunction::getFunction() {
     return TableFunction(READ_PYARROW_FUNC_NAME, tableFunc, bindFunc, initSharedState,
-        initLocalState, std::vector<LogicalTypeID>{LogicalTypeID::POINTER});
+        initLocalState, progressFunc, std::vector<LogicalTypeID>{LogicalTypeID::POINTER});
 }
 
 } // namespace kuzu