Make schema select fail on missing columns

Author: fallintoplace

src/iceberg/manifest/manifest_reader.cc

@@ -681,7 +681,19 @@ Result<std::shared_ptr<Schema>> ProjectSchema(std::shared_ptr<Schema> schema,
                                               const std::vector<std::string>& columns,
                                               bool case_sensitive) {
   if (!columns.empty()) {
-    return schema->Select(columns, case_sensitive);
+    if (std::ranges::contains(columns, Schema::kAllColumns)) {
+      return schema->Select(columns, case_sensitive);
+    }
+
+    std::vector<std::string> existing_columns;
+    existing_columns.reserve(columns.size());
+    for (const auto& column : columns) {
+      ICEBERG_ASSIGN_OR_RAISE(auto field, schema->FindFieldByName(column, case_sensitive));
+      if (field.has_value()) {
+        existing_columns.push_back(column);
+      }
+    }
+    return schema->Select(existing_columns, case_sensitive);
   }
   return schema;
 }

src/iceberg/schema.cc

@@ -366,9 +366,10 @@ Result<std::unique_ptr<Schema>> Schema::Select(std::span<const std::string> name
   std::unordered_set<int32_t> selected_ids;
   for (const auto& name : names) {
     ICEBERG_ASSIGN_OR_RAISE(auto result, FindFieldByName(name, case_sensitive));
-    if (result.has_value()) {
-      selected_ids.insert(result.value().get().field_id());
+    if (!result.has_value()) {
+      return InvalidArgument("Cannot find selected column: {}", name);
     }
+    selected_ids.insert(result.value().get().field_id());
   }
 
   PruneColumnVisitor visitor(selected_ids, /*select_full_types=*/true);

src/iceberg/schema.h

@@ -160,7 +160,8 @@ class ICEBERG_EXPORT Schema : public StructType {
   ///
   /// \param names Selected field names and nested names are dot-concatenated.
   /// \param case_sensitive Whether name matching is case-sensitive (default: true).
-  /// \return Projected schema containing only selected fields.
+  /// \return Projected schema containing only selected fields, or an error if any
+  /// requested field cannot be found.
   /// \note If the field name of a nested type has been selected, all of its
   /// sub-fields will be selected.
   Result<std::unique_ptr<Schema>> Select(std::span<const std::string> names,

src/iceberg/table_scan.cc

@@ -413,6 +413,11 @@ Result<std::unique_ptr<ScanType>> TableScanBuilder<ScanType>::Build() {
   ICEBERG_RETURN_UNEXPECTED(context_.Validate());
 
   ICEBERG_ASSIGN_OR_RAISE(auto schema, ResolveSnapshotSchema());
+  if (!context_.selected_columns.empty()) {
+    ICEBERG_ASSIGN_OR_RAISE(
+        std::ignore,
+        schema.get()->Select(context_.selected_columns, context_.case_sensitive));
+  }
   return ScanType::Make(metadata_, schema.get(), io_, std::move(context_));
 }
 

src/iceberg/test/schema_test.cc

@@ -816,13 +816,15 @@ INSTANTIATE_TEST_SUITE_P(
                         .create_schema = []() { return BasicSchema(); },
                         .select_fields = {"nonexistent"},
                         .expected_schema = []() { return MakeSchema(); },
-                        .should_succeed = true},
+                        .should_succeed = false,
+                        .expected_error_message = "Cannot find selected column: nonexistent"},
 
         SelectTestParam{.test_name = "SelectCaseSensitive",
                         .create_schema = []() { return BasicSchema(); },
                         .select_fields = {"Name"},  // case-sensitive
                         .expected_schema = []() { return MakeSchema(); },
-                        .should_succeed = true},
+                        .should_succeed = false,
+                        .expected_error_message = "Cannot find selected column: Name"},
 
         SelectTestParam{.test_name = "SelectCaseInsensitive",
                         .create_schema = []() { return BasicSchema(); },
@@ -840,6 +842,14 @@ INSTANTIATE_TEST_SUITE_P(
                           .expected_schema = []() { return MakeSchema(Id(), Name()); },
                           .should_succeed = true},
 
+                      SelectTestParam{
+                          .test_name = "SelectPartialMissingField",
+                          .create_schema = []() { return AddressSchema(); },
+                          .select_fields = {"id", "typo"},
+                          .expected_schema = []() { return MakeSchema(); },
+                          .should_succeed = false,
+                          .expected_error_message = "Cannot find selected column: typo"},
+
                       SelectTestParam{.test_name = "SelectNestedField",
                                       .create_schema = []() { return AddressSchema(); },
                                       .select_fields = {"address.street"},

src/iceberg/test/table_scan_test.cc

@@ -677,6 +677,13 @@ TEST_P(TableScanTest, SchemaWithSelectedColumnsAndFilter) {
   }
 }
 
+TEST_P(TableScanTest, SelectMissingColumnFails) {
+  ICEBERG_UNWRAP_OR_FAIL(auto builder,
+                         DataTableScanBuilder::Make(table_metadata_, file_io_));
+  builder->Select({"id", "typo"});
+  EXPECT_THAT(builder->Build(), IsError(ErrorKind::kInvalidArgument));
+}
+
 INSTANTIATE_TEST_SUITE_P(TableScanVersions, TableScanTest, testing::Values(1, 2, 3));
 
 }  // namespace iceberg