fix: invalidate table scan schema cache (#758)

## Summary
- invalidate the cached scan snapshot schema whenever
snapshot/ref/time-travel state changes
- preserve requested column-stat names and re-resolve them during Build
so stats IDs match the final snapshot schema
- add regression coverage for IncludeColumnStats before UseSnapshot and
non-main UseRef with an evolved schema

## Tests
- cmake --build build --target scan_test
- ctest --test-dir build -R scan_test --output-on-failure

Author: fallintoplace

src/iceberg/table_scan.cc

@@ -239,24 +239,16 @@ TableScanBuilder<ScanType>& TableScanBuilder<ScanType>::CaseSensitive(
 template <typename ScanType>
 TableScanBuilder<ScanType>& TableScanBuilder<ScanType>::IncludeColumnStats() {
   context_.return_column_stats = true;
+  context_.columns_to_keep_stats.clear();
+  requested_column_stats_.reset();
   return *this;
 }
 
 template <typename ScanType>
 TableScanBuilder<ScanType>& TableScanBuilder<ScanType>::IncludeColumnStats(
     const std::vector<std::string>& requested_columns) {
   context_.return_column_stats = true;
-  context_.columns_to_keep_stats.clear();
-  context_.columns_to_keep_stats.reserve(requested_columns.size());
-
-  ICEBERG_BUILDER_ASSIGN_OR_RETURN(auto schema_ref, ResolveSnapshotSchema());
-  const auto& schema = schema_ref.get();
-  for (const auto& column_name : requested_columns) {
-    ICEBERG_BUILDER_ASSIGN_OR_RETURN(auto field, schema->FindFieldByName(column_name));
-    if (field.has_value()) {
-      context_.columns_to_keep_stats.insert(field.value().get().field_id());
-    }
-  }
+  requested_column_stats_ = requested_columns;
 
   return *this;
 }
@@ -307,7 +299,6 @@ TableScanBuilder<ScanType>& TableScanBuilder<ScanType>::UseSnapshot(int64_t snap
 template <typename ScanType>
 TableScanBuilder<ScanType>& TableScanBuilder<ScanType>::UseRef(const std::string& ref) {
   if (ref == SnapshotRef::kMainBranch) {
-    snapshot_schema_ = nullptr;
     context_.snapshot_id.reset();
     return *this;
   }
@@ -396,6 +387,27 @@ TableScanBuilder<ScanType>& TableScanBuilder<ScanType>::UseBranch(
   return *this;
 }
 
+template <typename ScanType>
+Status TableScanBuilder<ScanType>::ResolveColumnStatsSelection() {
+  if (!requested_column_stats_.has_value()) {
+    return {};
+  }
+
+  context_.columns_to_keep_stats.clear();
+  context_.columns_to_keep_stats.reserve(requested_column_stats_->size());
+
+  ICEBERG_ASSIGN_OR_RAISE(auto schema_ref, ResolveSnapshotSchema());
+  const auto& schema = schema_ref.get();
+  for (const auto& column_name : *requested_column_stats_) {
+    ICEBERG_ASSIGN_OR_RAISE(auto field, schema->FindFieldByName(column_name));
+    if (field.has_value()) {
+      context_.columns_to_keep_stats.insert(field.value().get().field_id());
+    }
+  }
+
+  return {};
+}
+
 template <typename ScanType>
 Result<std::reference_wrapper<const std::shared_ptr<Schema>>>
 TableScanBuilder<ScanType>::ResolveSnapshotSchema() {
@@ -416,6 +428,7 @@ TableScanBuilder<ScanType>::ResolveSnapshotSchema() {
 template <typename ScanType>
 Result<std::unique_ptr<ScanType>> TableScanBuilder<ScanType>::Build() {
   ICEBERG_RETURN_UNEXPECTED(CheckErrors());
+  ICEBERG_RETURN_UNEXPECTED(ResolveColumnStatsSelection());
   ICEBERG_RETURN_UNEXPECTED(context_.Validate());
 
   ICEBERG_ASSIGN_OR_RAISE(auto schema, ResolveSnapshotSchema());

src/iceberg/table_scan.h

@@ -392,11 +392,13 @@ class ICEBERG_TEMPLATE_CLASS_EXPORT TableScanBuilder : public ErrorCollector {
 
   // Return the schema bound to the specified snapshot.
   Result<std::reference_wrapper<const std::shared_ptr<Schema>>> ResolveSnapshotSchema();
+  Status ResolveColumnStatsSelection();
 
   std::shared_ptr<TableMetadata> metadata_;
   std::shared_ptr<FileIO> io_;
   internal::TableScanContext context_;
   std::shared_ptr<Schema> snapshot_schema_;
+  std::optional<std::vector<std::string>> requested_column_stats_;
 };
 
 /// \brief Represents a configured scan operation on a table.

src/iceberg/test/table_scan_test.cc

@@ -231,6 +231,60 @@ TEST_P(TableScanTest, UseRefPreservesInt64SnapshotIds) {
   EXPECT_EQ(snapshot->snapshot_id, kLargeSnapshotId);
 }
 
+TEST_P(TableScanTest, IncludeColumnStatsUsesFinalSnapshotSchema) {
+  constexpr int64_t kBaseSnapshotId = 1000L;
+  constexpr int64_t kEvolvedSnapshotId = 2000L;
+  constexpr int32_t kBaseIdFieldId = 1;
+  constexpr int32_t kEvolvedIdFieldId = 10;
+  constexpr int32_t kEvolvedDataFieldId = 11;
+  constexpr int32_t kEvolvedSchemaId = 1;
+
+  auto evolved_schema = std::make_shared<Schema>(
+      std::vector<SchemaField>{
+          SchemaField::MakeRequired(kEvolvedIdFieldId, "id", int32()),
+          SchemaField::MakeRequired(kEvolvedDataFieldId, "data", string())},
+      kEvolvedSchemaId);
+  table_metadata_->schemas.push_back(evolved_schema);
+  table_metadata_->last_column_id = kEvolvedDataFieldId;
+  table_metadata_->snapshots.push_back(std::make_shared<Snapshot>(
+      Snapshot{.snapshot_id = kEvolvedSnapshotId,
+               .parent_snapshot_id = kBaseSnapshotId,
+               .sequence_number = 2L,
+               .timestamp_ms = TimePointMsFromUnixMs(1609459201000L),
+               .manifest_list = "/tmp/metadata/snap-2000-2-manifest-list.avro",
+               .schema_id = evolved_schema->schema_id()}));
+  table_metadata_->refs["evolved-branch"] = std::make_shared<SnapshotRef>(
+      SnapshotRef{.snapshot_id = kEvolvedSnapshotId, .retention = SnapshotRef::Branch{}});
+
+  {
+    ICEBERG_UNWRAP_OR_FAIL(auto builder,
+                           DataTableScanBuilder::Make(table_metadata_, file_io_));
+    builder->IncludeColumnStats({"id"}).UseSnapshot(kEvolvedSnapshotId);
+    ICEBERG_UNWRAP_OR_FAIL(auto scan, builder->Build());
+    ICEBERG_UNWRAP_OR_FAIL(auto scan_schema, scan->schema());
+
+    EXPECT_EQ(scan_schema->schema_id(), evolved_schema->schema_id());
+    const auto& stats_fields = scan->context().columns_to_keep_stats;
+    EXPECT_EQ(stats_fields.size(), 1);
+    EXPECT_TRUE(stats_fields.contains(kEvolvedIdFieldId));
+    EXPECT_FALSE(stats_fields.contains(kBaseIdFieldId));
+  }
+
+  {
+    ICEBERG_UNWRAP_OR_FAIL(auto builder,
+                           DataTableScanBuilder::Make(table_metadata_, file_io_));
+    builder->IncludeColumnStats({"id"}).UseRef("evolved-branch");
+    ICEBERG_UNWRAP_OR_FAIL(auto scan, builder->Build());
+    ICEBERG_UNWRAP_OR_FAIL(auto scan_schema, scan->schema());
+
+    EXPECT_EQ(scan_schema->schema_id(), evolved_schema->schema_id());
+    const auto& stats_fields = scan->context().columns_to_keep_stats;
+    EXPECT_EQ(stats_fields.size(), 1);
+    EXPECT_TRUE(stats_fields.contains(kEvolvedIdFieldId));
+    EXPECT_FALSE(stats_fields.contains(kBaseIdFieldId));
+  }
+}
+
 TEST_P(TableScanTest, TableScanBuilderValidationErrors) {
   // Test negative min rows
   ICEBERG_UNWRAP_OR_FAIL(auto builder,