feat(writer): Make LocationGenerator partition-aware (#1625)

## Which issue does this PR close?

- Closes #1624 

## What changes are included in this PR?
- Added `PartitionKey` to hold necessary partition-related info for path
generation
- [BREAKING] Changed `LocationGenerator::generate_location` to accept an
optional partition key
- [BREAKING] Have `ParquetWriter` take in an `Option<PartitionKey>`

## Are these changes tested?
Added ut

Author: CTTY

crates/iceberg/src/spec/partition.rs

@@ -176,6 +176,39 @@ impl PartitionSpec {
     }
 }
 
+/// A partition key represents a specific partition in a table, containing the partition spec,
+/// schema, and the actual partition values.
+#[derive(Clone, Debug)]
+pub struct PartitionKey {
+    /// The partition spec that contains the partition fields.
+    spec: PartitionSpec,
+    /// The schema to which the partition spec is bound.
+    schema: SchemaRef,
+    /// Partition fields' values in struct.
+    data: Struct,
+}
+
+impl PartitionKey {
+    /// Creates a new partition key with the given spec, schema, and data.
+    pub fn new(spec: PartitionSpec, schema: SchemaRef, data: Struct) -> Self {
+        Self { spec, schema, data }
+    }
+
+    /// Generates a partition path based on the partition values.
+    pub fn to_path(&self) -> String {
+        self.spec.partition_to_path(&self.data, self.schema.clone())
+    }
+
+    /// Returns `true` if the partition key is absent (`None`)
+    /// or represents an unpartitioned spec.
+    pub fn is_effectively_none(partition_key: Option<&PartitionKey>) -> bool {
+        match partition_key {
+            None => true,
+            Some(pk) => pk.spec.is_unpartitioned(),
+        }
+    }
+}
+
 /// Reference to [`UnboundPartitionSpec`].
 pub type UnboundPartitionSpecRef = Arc<UnboundPartitionSpec>;
 /// Unbound partition field can be built without a schema and later bound to a schema.
@@ -1772,7 +1805,7 @@ mod tests {
 
         assert_eq!(
             spec.partition_to_path(&data, schema.into()),
-            "id=42/name=\"alice\""
+            "id=42/name=alice"
         );
     }
 }

crates/iceberg/src/spec/snapshot_summary.rs

@@ -819,7 +819,7 @@ mod tests {
         assert_eq!(props.get(ADDED_FILE_SIZE).unwrap(), "300");
         assert_eq!(props.get(REMOVED_FILE_SIZE).unwrap(), "100");
 
-        let partition_key = format!("{}{}", CHANGED_PARTITION_PREFIX, "year=\"2025\"");
+        let partition_key = format!("{}{}", CHANGED_PARTITION_PREFIX, "year=2025");
 
         assert!(props.contains_key(&partition_key));
 

crates/iceberg/src/spec/transform.rs

@@ -142,7 +142,7 @@ impl Transform {
                 let field_type = field_type.as_primitive_type().unwrap();
                 let datum = Datum::new(field_type.clone(), value);
                 match self {
-                    Self::Identity => datum.to_string(),
+                    Self::Identity => datum.to_human_string(),
                     Self::Void => "null".to_string(),
                     _ => {
                         todo!()

crates/iceberg/src/spec/values.rs

@@ -1266,6 +1266,17 @@ impl Datum {
             _ => false,
         }
     }
+
+    /// Returns a human-readable string representation of this literal.
+    ///
+    /// For string literals, this returns the raw string value without quotes.
+    /// For all other literals, it falls back to [`to_string()`].
+    pub(crate) fn to_human_string(&self) -> String {
+        match self.literal() {
+            PrimitiveLiteral::String(s) => s.to_string(),
+            _ => self.to_string(),
+        }
+    }
 }
 
 /// Map is a collection of key-value pairs with a key type and a value type.

crates/iceberg/src/writer/base_writer/data_file_writer.rs

@@ -144,6 +144,7 @@ mod test {
         let pw = ParquetWriterBuilder::new(
             WriterProperties::builder().build(),
             Arc::new(schema),
+            None,
             file_io.clone(),
             location_gen,
             file_name_gen,
@@ -221,6 +222,7 @@ mod test {
         let parquet_writer_builder = ParquetWriterBuilder::new(
             WriterProperties::builder().build(),
             Arc::new(schema.clone()),
+            None,
             file_io.clone(),
             location_gen,
             file_name_gen,

crates/iceberg/src/writer/base_writer/equality_delete_writer.rs

@@ -404,6 +404,7 @@ mod test {
         let pb = ParquetWriterBuilder::new(
             WriterProperties::builder().build(),
             Arc::new(delete_schema),
+            None,
             file_io.clone(),
             location_gen,
             file_name_gen,
@@ -569,6 +570,7 @@ mod test {
         let pb = ParquetWriterBuilder::new(
             WriterProperties::builder().build(),
             Arc::new(delete_schema),
+            None,
             file_io.clone(),
             location_gen,
             file_name_gen,

crates/iceberg/src/writer/file_writer/location_generator.rs

@@ -21,14 +21,24 @@ use std::sync::Arc;
 use std::sync::atomic::AtomicU64;
 
 use crate::Result;
-use crate::spec::{DataFileFormat, TableMetadata};
+use crate::spec::{DataFileFormat, PartitionKey, TableMetadata};
 
 /// `LocationGenerator` used to generate the location of data file.
 pub trait LocationGenerator: Clone + Send + 'static {
-    /// Generate an absolute path for the given file name.
-    /// e.g.
-    /// For file name "part-00000.parquet", the generated location maybe "/table/data/part-00000.parquet"
-    fn generate_location(&self, file_name: &str) -> String;
+    /// Generate an absolute path for the given file name that includes the partition path.
+    ///
+    /// # Arguments
+    ///
+    /// * `partition_key` - The partition key of the file. If None, generate a non-partitioned path.
+    /// * `file_name` - The name of the file
+    ///
+    /// # Returns
+    ///
+    /// An absolute path that includes the partition path, e.g.,
+    /// "/table/data/id=1/name=alice/part-00000.parquet"
+    /// or non-partitioned path:
+    /// "/table/data/part-00000.parquet"
+    fn generate_location(&self, partition_key: Option<&PartitionKey>, file_name: &str) -> String;
 }
 
 const WRITE_DATA_LOCATION: &str = "write.data.path";
@@ -39,29 +49,38 @@ const DEFAULT_DATA_DIR: &str = "/data";
 /// `DefaultLocationGenerator` used to generate the data dir location of data file.
 /// The location is generated based on the table location and the data location in table properties.
 pub struct DefaultLocationGenerator {
-    dir_path: String,
+    data_location: String,
 }
 
 impl DefaultLocationGenerator {
     /// Create a new `DefaultLocationGenerator`.
     pub fn new(table_metadata: TableMetadata) -> Result<Self> {
         let table_location = table_metadata.location();
         let prop = table_metadata.properties();
-        let data_location = prop
+        let configured_data_location = prop
             .get(WRITE_DATA_LOCATION)
             .or(prop.get(WRITE_FOLDER_STORAGE_LOCATION));
-        let dir_path = if let Some(data_location) = data_location {
+        let data_location = if let Some(data_location) = configured_data_location {
             data_location.clone()
         } else {
             format!("{}{}", table_location, DEFAULT_DATA_DIR)
         };
-        Ok(Self { dir_path })
+        Ok(Self { data_location })
     }
 }
 
 impl LocationGenerator for DefaultLocationGenerator {
-    fn generate_location(&self, file_name: &str) -> String {
-        format!("{}/{}", self.dir_path, file_name)
+    fn generate_location(&self, partition_key: Option<&PartitionKey>, file_name: &str) -> String {
+        if PartitionKey::is_effectively_none(partition_key) {
+            format!("{}/{}", self.data_location, file_name)
+        } else {
+            format!(
+                "{}/{}/{}",
+                self.data_location,
+                partition_key.unwrap().to_path(),
+                file_name
+            )
+        }
     }
 }
 
@@ -115,11 +134,15 @@ impl FileNameGenerator for DefaultFileNameGenerator {
 #[cfg(test)]
 pub(crate) mod test {
     use std::collections::HashMap;
+    use std::sync::Arc;
 
     use uuid::Uuid;
 
     use super::LocationGenerator;
-    use crate::spec::{FormatVersion, PartitionSpec, StructType, TableMetadata};
+    use crate::spec::{
+        FormatVersion, Literal, NestedField, PartitionKey, PartitionSpec, PrimitiveType, Schema,
+        Struct, StructType, TableMetadata, Transform, Type,
+    };
     use crate::writer::file_writer::location_generator::{
         FileNameGenerator, WRITE_DATA_LOCATION, WRITE_FOLDER_STORAGE_LOCATION,
     };
@@ -136,8 +159,17 @@ pub(crate) mod test {
     }
 
     impl LocationGenerator for MockLocationGenerator {
-        fn generate_location(&self, file_name: &str) -> String {
-            format!("{}/{}", self.root, file_name)
+        fn generate_location(&self, partition: Option<&PartitionKey>, file_name: &str) -> String {
+            if PartitionKey::is_effectively_none(partition) {
+                format!("{}/{}", self.root, file_name)
+            } else {
+                format!(
+                    "{}/{}/{}",
+                    self.root,
+                    partition.unwrap().to_path(),
+                    file_name
+                )
+            }
         }
     }
 
@@ -169,7 +201,7 @@ pub(crate) mod test {
             encryption_keys: HashMap::new(),
         };
 
-        let file_name_genertaor = super::DefaultFileNameGenerator::new(
+        let file_name_generator = super::DefaultFileNameGenerator::new(
             "part".to_string(),
             Some("test".to_string()),
             crate::spec::DataFileFormat::Parquet,
@@ -179,7 +211,7 @@ pub(crate) mod test {
         let location_generator =
             super::DefaultLocationGenerator::new(table_metadata.clone()).unwrap();
         let location =
-            location_generator.generate_location(&file_name_genertaor.generate_file_name());
+            location_generator.generate_location(None, &file_name_generator.generate_file_name());
         assert_eq!(location, "s3://data.db/table/data/part-00000-test.parquet");
 
         // test custom data location
@@ -190,7 +222,7 @@ pub(crate) mod test {
         let location_generator =
             super::DefaultLocationGenerator::new(table_metadata.clone()).unwrap();
         let location =
-            location_generator.generate_location(&file_name_genertaor.generate_file_name());
+            location_generator.generate_location(None, &file_name_generator.generate_file_name());
         assert_eq!(
             location,
             "s3://data.db/table/data_1/part-00001-test.parquet"
@@ -203,7 +235,7 @@ pub(crate) mod test {
         let location_generator =
             super::DefaultLocationGenerator::new(table_metadata.clone()).unwrap();
         let location =
-            location_generator.generate_location(&file_name_genertaor.generate_file_name());
+            location_generator.generate_location(None, &file_name_generator.generate_file_name());
         assert_eq!(
             location,
             "s3://data.db/table/data_2/part-00002-test.parquet"
@@ -217,7 +249,79 @@ pub(crate) mod test {
         let location_generator =
             super::DefaultLocationGenerator::new(table_metadata.clone()).unwrap();
         let location =
-            location_generator.generate_location(&file_name_genertaor.generate_file_name());
+            location_generator.generate_location(None, &file_name_generator.generate_file_name());
         assert_eq!(location, "s3://data.db/data_3/part-00003-test.parquet");
     }
+
+    #[test]
+    fn test_location_generate_with_partition() {
+        // Create a schema with two fields: id (int) and name (string)
+        let schema = Arc::new(
+            Schema::builder()
+                .with_schema_id(1)
+                .with_fields(vec![
+                    NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
+                    NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
+                ])
+                .build()
+                .unwrap(),
+        );
+
+        // Create a partition spec with both fields
+        let partition_spec = PartitionSpec::builder(schema.clone())
+            .add_partition_field("id", "id", Transform::Identity)
+            .unwrap()
+            .add_partition_field("name", "name", Transform::Identity)
+            .unwrap()
+            .build()
+            .unwrap();
+
+        // Create partition data with values
+        let partition_data =
+            Struct::from_iter([Some(Literal::int(42)), Some(Literal::string("alice"))]);
+
+        // Create a partition key
+        let partition_key = PartitionKey::new(partition_spec, schema, partition_data);
+
+        // Test with MockLocationGenerator
+        let mock_location_gen = MockLocationGenerator::new("/base/path".to_string());
+        let file_name = "data-00000.parquet";
+        let location = mock_location_gen.generate_location(Some(&partition_key), file_name);
+        assert_eq!(location, "/base/path/id=42/name=alice/data-00000.parquet");
+
+        // Create a table metadata for DefaultLocationGenerator
+        let table_metadata = TableMetadata {
+            format_version: FormatVersion::V2,
+            table_uuid: Uuid::parse_str("fb072c92-a02b-11e9-ae9c-1bb7bc9eca94").unwrap(),
+            location: "s3://data.db/table".to_string(),
+            last_updated_ms: 1515100955770,
+            last_column_id: 2,
+            schemas: HashMap::new(),
+            current_schema_id: 1,
+            partition_specs: HashMap::new(),
+            default_spec: PartitionSpec::unpartition_spec().into(),
+            default_partition_type: StructType::new(vec![]),
+            last_partition_id: 1000,
+            default_sort_order_id: 0,
+            sort_orders: HashMap::from_iter(vec![]),
+            snapshots: HashMap::default(),
+            current_snapshot_id: None,
+            last_sequence_number: 1,
+            properties: HashMap::new(),
+            snapshot_log: Vec::new(),
+            metadata_log: vec![],
+            refs: HashMap::new(),
+            statistics: HashMap::new(),
+            partition_statistics: HashMap::new(),
+            encryption_keys: HashMap::new(),
+        };
+
+        // Test with DefaultLocationGenerator
+        let default_location_gen = super::DefaultLocationGenerator::new(table_metadata).unwrap();
+        let location = default_location_gen.generate_location(Some(&partition_key), file_name);
+        assert_eq!(
+            location,
+            "s3://data.db/table/data/id=42/name=alice/data-00000.parquet"
+        );
+    }
 }