Always include partition column when get-table-layout (#3045)

ZhitongYan · Leo Yan · burhan94 · web-flow · commit 2d31a1b6a43d · 2025-10-29T12:19:14.000-04:00
Co-authored-by: Leo Yan &lt;leozy@amazon.com&gt;
Co-authored-by: burhan94 &lt;hamburha@amazon.com&gt;
diff --git a/athena-snowflake/src/main/java/com/amazonaws/athena/connectors/snowflake/SnowflakeMetadataHandler.java b/athena-snowflake/src/main/java/com/amazonaws/athena/connectors/snowflake/SnowflakeMetadataHandler.java
@@ -231,6 +231,10 @@ public GetDataSourceCapabilitiesResponse doGetDataSourceCapabilities(BlockAlloca
     public void enhancePartitionSchema(SchemaBuilder partitionSchemaBuilder, GetTableLayoutRequest request)
     {
         LOGGER.info("{}: Catalog {}, table {}", request.getQueryId(), request.getTableName().getSchemaName(), request.getTableName());
+        // Always ensure the partition column exists in the schema
+        if (partitionSchemaBuilder.getField(BLOCK_PARTITION_COLUMN_NAME) == null) {
+            partitionSchemaBuilder.addField(BLOCK_PARTITION_COLUMN_NAME, Types.MinorType.VARCHAR.getType());
+        }
         partitionSchemaBuilder.addField(QUERY_ID, new ArrowType.Utf8());
         partitionSchemaBuilder.addField(PREPARED_STMT, new ArrowType.Utf8());
     }

Original file line number	Diff line number	Diff line change
`@@ -231,6 +231,10 @@ public GetDataSourceCapabilitiesResponse doGetDataSourceCapabilities(BlockAlloca`
`231`	`231`	`public void enhancePartitionSchema(SchemaBuilder partitionSchemaBuilder, GetTableLayoutRequest request)`
`232`	`232`	`{`
`233`	`233`	`LOGGER.info("{}: Catalog {}, table {}", request.getQueryId(), request.getTableName().getSchemaName(), request.getTableName());`
	`234`	`+ // Always ensure the partition column exists in the schema`
	`235`	`+ if (partitionSchemaBuilder.getField(BLOCK_PARTITION_COLUMN_NAME) == null) {`
	`236`	`+ partitionSchemaBuilder.addField(BLOCK_PARTITION_COLUMN_NAME, Types.MinorType.VARCHAR.getType());`
	`237`	`+ }`
`234`	`238`	`partitionSchemaBuilder.addField(QUERY_ID, new ArrowType.Utf8());`
`235`	`239`	`partitionSchemaBuilder.addField(PREPARED_STMT, new ArrowType.Utf8());`
`236`	`240`	`}`