feat: Improve Column Remapping Performance in COPY FROM

extension/json/tests/json_test.cpp

@@ -20,6 +20,7 @@
 #include <memory>
 #include <vector>
 
+#include <arrow/array.h>
 #include <arrow/filesystem/localfs.h>
 #include <arrow/type.h>
 #include "neug/compiler/common/case_insensitive_map.h"
@@ -100,7 +101,7 @@ class JsonTest : public ::testing::Test {
       const std::string& csvFile, const std::vector<std::string>& columnNames,
       const std::vector<std::shared_ptr<::common::DataType>>& columnTypes,
       const common::case_insensitive_map_t<std::string>& options = {},
-      const std::vector<std::string>& skipColumns = {},
+      const std::vector<std::string>& projectColumns = {},
       std::shared_ptr<::common::Expression> skipRows = nullptr) {
     auto sharedState = std::make_shared<reader::ReadSharedState>();
 
@@ -120,7 +121,7 @@ class JsonTest : public ::testing::Test {
     externalSchema.file = fileSchema;
 
     sharedState->schema = std::move(externalSchema);
-    sharedState->skipColumns = skipColumns;
+    sharedState->projectColumns = projectColumns;
     sharedState->skipRows = skipRows;
 
     return sharedState;

extension/parquet/tests/parquet_test.cpp

@@ -534,7 +534,7 @@ TEST_F(ParquetTest, TestIntegration_ColumnPruning) {
   PARQUET_THROW_NOT_OK(
       parquet::arrow::WriteTable(*table, arrow::default_memory_pool(), outfile, 1));
 
-  // Set up shared state with skipColumns
+  // Set up shared state with projectColumns
   auto sharedState = std::make_shared<reader::ReadSharedState>();
   auto entrySchema = std::make_shared<reader::TableEntrySchema>();
   entrySchema->columnNames = {"id", "name", "score", "grade"};
@@ -551,20 +551,20 @@ TEST_F(ParquetTest, TestIntegration_ColumnPruning) {
   externalSchema.file = fileSchema;
   sharedState->schema = std::move(externalSchema);
 
-  // Neug's column pruning: skip "name" column
-  sharedState->skipColumns = {"name"};
+  // Neug's column projection: id, score, grade (exclude "name")
+  sharedState->projectColumns = {"id", "score", "grade"};
 
   auto reader = createParquetReader(sharedState);
   auto localState = std::make_shared<reader::ReadLocalState>();
   execution::Context ctx;
   reader->read(localState, ctx);
 
-  // Verify extension translates skipColumns to Arrow projection
-  // Should have 3 columns (id, score, grade - "name" is skipped)
+  // Verify extension translates projectColumns to Arrow projection
+  // Should have 3 columns (id, score, grade - "name" is excluded)
   EXPECT_EQ(ctx.col_num(), 3)
-      << "Extension should translate Neug's skipColumns to Arrow column projection";
+      << "Extension should translate Neug's projectColumns to Arrow column projection";
   EXPECT_EQ(sharedState->columnNum(), 3)
-      << "Extension should update columnNum after pruning";
+      << "Extension should update columnNum after projection";
 }
 
 TEST_F(ParquetTest, TestIntegration_FilterPushdown) {
@@ -761,7 +761,7 @@ TEST_F(ParquetTest, TestIntegration_CombinedFilterAndProjection) {
   externalSchema.entry = entrySchema;
   externalSchema.file = fileSchema;
   sharedState->schema = std::move(externalSchema);
-  sharedState->skipColumns = {"name"};  // Prune "name" column
+  sharedState->projectColumns = {"id", "score", "grade"};  // Exclude "name"
   sharedState->skipRows = filterExpr;   // Filter score > 90.0
 
   auto reader = createParquetReader(sharedState);

include/neug/compiler/function/table/bind_data.h

@@ -56,15 +56,14 @@ struct NEUG_API TableFuncBindData {
       : columns{other.columns},
         numRows{other.numRows},
         params(other.params),
-        columnSkips{other.columnSkips} {}
+        projectColumns{other.projectColumns} {}
   TableFuncBindData& operator=(const TableFuncBindData& other) = delete;
   virtual ~TableFuncBindData() = default;
 
   common::idx_t getNumColumns() const { return columns.size(); }
-  void setColumnSkips(std::vector<bool> skips) {
-    columnSkips = std::move(skips);
+  void setProjectColumns(std::vector<std::string> projectColumns) {
+    this->projectColumns = std::move(projectColumns);
   }
-  std::vector<bool> getColumnSkips() const;
 
   std::shared_ptr<binder::Expression> getRowSkips() const { return rowSkips; }
 
@@ -90,8 +89,12 @@ struct NEUG_API TableFuncBindData {
     return *common::neug_dynamic_cast<TARGET*>(this);
   }
 
+  const std::vector<std::string>& getProjectColumns() const {
+    return projectColumns;
+  }
+
  protected:
-  std::vector<bool> columnSkips;
+  std::vector<std::string> projectColumns;
   std::shared_ptr<binder::Expression> rowSkips;
 };
 

include/neug/compiler/optimizer/project_into_data_source_optimizer.h

@@ -0,0 +1,46 @@
+/**
+ * Copyright 2020 Alibaba Group Holding Limited.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include "logical_operator_visitor.h"
+#include "neug/compiler/planner/operator/logical_operator.h"
+#include "neug/compiler/planner/operator/logical_plan.h"
+
+namespace neug {
+namespace optimizer {
+
+// ProjectIntoDataSourceOptimizer pushes a Projection directly into
+// LogicalTableFunctionCall when the pattern is:
+//   LogicalTableFunctionCall -> LogicalProjection
+// After rewrite, the Projection is removed and projected expressions are
+// carried by LogicalTableFunctionCall::bindData->projectColumns.
+class ProjectIntoDataSourceOptimizer : public LogicalOperatorVisitor {
+ public:
+  void rewrite(planner::LogicalPlan* plan);
+
+ private:
+  std::shared_ptr<planner::LogicalOperator> visitOperator(
+      const std::shared_ptr<planner::LogicalOperator>& op);
+  std::shared_ptr<planner::LogicalOperator> visitProjectionReplace(
+      std::shared_ptr<planner::LogicalOperator> op) override;
+
+ private:
+  planner::LogicalPlan* plan;
+};
+
+}  // namespace optimizer
+}  // namespace neug

include/neug/compiler/planner/operator/logical_table_function_call.h

@@ -27,8 +27,8 @@ class NEUG_API LogicalTableFunctionCall final : public LogicalOperator {
   const function::TableFunction& getTableFunc() const { return tableFunc; }
   function::TableFuncBindData* getBindData() const { return bindData.get(); }
 
-  void setColumnSkips(std::vector<bool> columnSkips) {
-    bindData->setColumnSkips(std::move(columnSkips));
+  void setProjectColumns(std::vector<std::string> projectColumns) {
+    bindData->setProjectColumns(std::move(projectColumns));
   }
 
   void setNodeMaskRoots(std::vector<std::shared_ptr<LogicalOperator>> roots) {

include/neug/utils/reader/options.h

@@ -133,12 +133,12 @@ struct ReadOptions {
  * @brief Template base class for building format-specific scan options
  *
  * This template class provides a generic interface for building scan options
- * for different data formats. It handles column projection (skipColumns) and
- * row filtering (skipRows) operations. Derived classes implement
+ * for different data formats. It handles column projection (projectColumns)
+ * and row filtering (skipRows) operations. Derived classes implement
  * format-specific option building logic.
  *
  * @tparam T The type of options structure returned by build() and modified by
- *           skipColumns() and skipRows()
+ *           projectColumns() and skipRows()
  */
 template <class T>
 class OptionsBuilder {
@@ -158,11 +158,11 @@ class OptionsBuilder {
   virtual T build() const = 0;
 
   /**
-   * @brief Applies column projection to exclude skipped columns
+   * @brief Applies column projection to include only specified columns
    * @param options The options structure to modify
    * @return true if column projection was successfully applied, false otherwise
    */
-  virtual bool skipColumns(T& options) { return false; }
+  virtual bool projectColumns(T& options) { return false; }
 
   /**
    * @brief Applies row filtering based on filter expressions
@@ -206,8 +206,9 @@ struct ArrowOptions {
  *   and batch size
  * - File format: builds format-specific FileFormat via buildFileFormat()
  *
- * The skipColumns() method implements column pruning by setting the projection
- * expression in scanOptions, excluding columns listed in state.skipColumns.
+ * The projectColumns() method implements column pruning by setting the
+ * projection expression in scanOptions to include only columns listed in
+ * state.projectColumns.
  *
  * The skipRows() method implements filter pushdown by converting
  * common::Expression to arrow::compute::Expression and setting it as the
@@ -241,15 +242,15 @@ class ArrowOptionsBuilder : public OptionsBuilder<ArrowOptions> {
   virtual ArrowOptions build() const override = 0;
 
   /**
-   * @brief Applies column projection to exclude skipped columns
+   * @brief Applies column projection to include only specified columns
    *
-   * Modifies the projection expression in options.scanOptions to exclude
-   * columns listed in state.skipColumns, implementing column pruning.
+   * Modifies the projection expression in options.scanOptions to include
+   * only columns listed in state.projectColumns, implementing column pruning.
    *
    * @param options The ArrowOptions to modify
    * @return true if column projection was successfully applied, false otherwise
    */
-  virtual bool skipColumns(ArrowOptions& options) override;
+  virtual bool projectColumns(ArrowOptions& options) override;
 
   /**
    * @brief Applies row filtering based on filter expressions

include/neug/utils/reader/reader.h

@@ -86,7 +86,7 @@ struct ReadLocalState {
  * operations. It includes:
  * - Schema information: external table metadata (column names, types, file
  * info)
- * - Column pruning: list of columns to skip during projection
+ * - Column projection: list of columns to include in output
  * - Filter pushdown: predicate expression for row filtering
  *
  * The columnNum() method calculates the effective number of columns after
@@ -98,27 +98,19 @@ struct ReadLocalState {
  */
 struct ReadSharedState {
   ExternalSchema schema;
-  std::vector<std::string> skipColumns;
+  std::vector<std::string> projectColumns;
   std::shared_ptr<::common::Expression> skipRows;
 
   /**
-   * @brief Get the number of columns after column pruning
-   * @return The number of columns remaining (not skipped) according to
-   * skipColumns
+   * @brief Get the number of columns after column projection
+   * @return The number of projected columns, or all columns if no projection
    */
   int columnNum() {
     if (!schema.entry) {
       return 0;
     }
     const auto& allColumns = schema.entry->columnNames;
-    int total = static_cast<int>(allColumns.size());
-    for (const auto& column : allColumns) {
-      if (std::find(skipColumns.begin(), skipColumns.end(), column) !=
-          skipColumns.end()) {
-        --total;
-      }
-    }
-    return total;
+    return projectColumns.empty() ? allColumns.size() : projectColumns.size();
   }
 };
 

proto/cypher_dml.proto

@@ -27,7 +27,7 @@ message DataSource {
     string extension_name = 1;
     EntrySchema entry_schema = 4;
     FileSchema file_schema = 5;
-    repeated string skip_columns = 6;
+    repeated string project_columns = 6;
     common.Expression skip_rows = 7;
 }
 

src/compiler/binder/bind/read/bind_load_from.cpp

@@ -25,6 +25,7 @@
 #include "neug/compiler/binder/expression/expression_util.h"
 #include "neug/compiler/binder/query/reading_clause/bound_load_from.h"
 #include "neug/compiler/common/types/value/value.h"
+#include "neug/compiler/function/table/scan_file_function.h"
 #include "neug/compiler/parser/query/reading_clause/load_from.h"
 #include "neug/compiler/parser/scan_source.h"
 #include "neug/utils/exception/exception.h"
@@ -56,11 +57,6 @@ std::unique_ptr<BoundReadingClause> Binder::bindLoadFrom(
     auto boundScanSource = bindFileScanSource(
         *source, loadFrom.getParsingOptions(), columnNames, columnTypes);
     auto& scanInfo = boundScanSource->constCast<BoundTableScanSource>().info;
-    auto& bindData = scanInfo.bindData->cast<function::ScanFileBindData>();
-    auto& fileInfo = bindData.fileScanInfo;
-    // We support LOAD FROM by ArrowArrayContextColumn with BATCH_READ set to
-    // false.
-    fileInfo.options.insert({"BATCH_READ", common::Value::createValue(false)});
     boundLoadFrom = std::make_unique<BoundLoadFrom>(scanInfo.copy());
   } break;
   default:

src/compiler/function/table/bind_data.cpp

@@ -27,18 +27,6 @@
 namespace neug {
 namespace function {
 
-std::vector<bool> TableFuncBindData::getColumnSkips() const {
-  if (columnSkips
-          .empty()) {  // If not specified, all columns should be scanned.
-    std::vector<bool> skips;
-    for (auto i = 0u; i < getNumColumns(); ++i) {
-      skips.push_back(false);
-    }
-    return skips;
-  }
-  return columnSkips;
-}
-
 bool TableFuncBindData::getIgnoreErrorsOption() const {
   return common::CopyConstants::DEFAULT_IGNORE_ERRORS;
 }

src/compiler/gopt/g_query_converter.cpp

@@ -1173,15 +1173,10 @@ void GQueryConvertor::convertDataSource(
     sourcePB->set_allocated_skip_rows(
         exprConvertor->convert(*rowSkips, {}).release());
   }
-  auto columnSkips = scanBindData->getColumnSkips();
-  if (scanBindData->columns.size() != columnSkips.size()) {
-    THROW_EXCEPTION_WITH_FILE_LINE("Each column should have a skip flag");
-  }
-  for (auto idx = 0; idx < scanBindData->columns.size(); idx++) {
-    if (columnSkips[idx]) {
-      sourcePB->mutable_skip_columns()->Add(
-          scanBindData->columns[idx]->rawName());
-    }
+  // Proto field is named skip_columns for historical reasons; it now carries
+  // the list of columns to project (include), not skip.
+  for (const auto& column : scanBindData->getProjectColumns()) {
+    sourcePB->add_project_columns(column);
   }
 
   auto physicalPB = std::make_unique<::physical::PhysicalOpr>();

src/compiler/optimizer/CMakeLists.txt

@@ -21,7 +21,8 @@ add_library(neug_optimizer
         union_alias_map_optimizer.cpp
         remove_subquery_as_join.cpp
         common_pattern_reuse_optimizer.cpp
-        project_join_condition_optimizer.cpp)
+        project_join_condition_optimizer.cpp
+        project_into_data_source_optimizer.cpp)
 add_dependencies(neug_optimizer neug_plan_proto)
 
 set(ALL_OBJECT_FILES

src/compiler/optimizer/optimizer.cpp

@@ -37,6 +37,7 @@
 #include "neug/compiler/optimizer/filter_push_down_pattern.h"
 #include "neug/compiler/optimizer/flat_join_to_expand_optimizer.h"
 #include "neug/compiler/optimizer/limit_push_down_optimizer.h"
+#include "neug/compiler/optimizer/project_into_data_source_optimizer.h"
 #include "neug/compiler/optimizer/project_join_condition_optimizer.h"
 #include "neug/compiler/optimizer/projection_push_down_optimizer.h"
 #include "neug/compiler/optimizer/remove_factorization_rewriter.h"
@@ -74,6 +75,9 @@ void Optimizer::optimize(
         context->getClientConfig()->recursivePatternSemantic, context);
     projectionPushDownOptimizer.rewrite(plan);
 
+    auto projectIntoDataSourceOptimizer = ProjectIntoDataSourceOptimizer();
+    projectIntoDataSourceOptimizer.rewrite(plan);
+
     auto limitPushDownOptimizer = LimitPushDownOptimizer();
     limitPushDownOptimizer.rewrite(plan);