Partition subfolder is not showing

[Tanvo149]

I've been learning Rust and iceberg recently and ran into an issue while writing to Parquet files. It doesn't seem to create a partition subdirectory as expected. I'm using the REST Catalog locally as my Iceberg catalog.

It writes to "/tmp/iceberg_warehouse255034623058200441/iceberg_data/cat/t1/data instead of "/tmp/iceberg_warehouse255034623058200441/iceberg_data/cat/t1/data/[partition]...

Am i doing something wrong? Also, how can I just write to parquet by the columns instead of partition_value?

I tried looking at github folder to see if there was any example, and I tried to mirror it.

Thank you for taking the time in reading this. Thanks.

        let props = WriterProperties::builder()
        .set_compression(parquet::basic::Compression::SNAPPY)
        .set_created_by("Data".to_string())
        .build();

     let parquet_writer_builder = ParquetWriterBuilder::new(
            props,
            _created_table.metadata().current_schema().clone(),
            _created_table.file_io().clone(),
            location_generator.clone(),
            file_name_generator.clone()
    );


    let partition_value = Struct::from_iter([Some(Literal::int(1))]);

    let mut data_file_writer =
            DataFileWriterBuilder::new(parquet_writer_builder, Some(partition_value.clone()))
                .build()
                .await?;

    let arrow_schema = arrow_schema::Schema::new(vec![
            Field::new("name", DataType::Utf8, false).with_metadata(HashMap::from([(
                PARQUET_FIELD_ID_META_KEY.to_string(),
                1.to_string(),
            )])),
            Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
                PARQUET_FIELD_ID_META_KEY.to_string(),
                2.to_string(),
            )])),
            Field::new("partition_date", DataType::Date32, false).with_metadata(HashMap::from([(
                PARQUET_FIELD_ID_META_KEY.to_string(),
                3.to_string(),
            )])),
        ]);


        let date_strs = vec!["2024-01-01", "2024-01-01", "2024-01-01"]; // Added a duplicate for testing grouping
        let date32_values: Vec<i32> = date_strs.iter().map(|s| to_date32(s)).collect();

        let batch = RecordBatch::try_new(Arc::new(arrow_schema.clone()), vec![
            Arc::new(StringArray::from(vec!["Alice", "Bob", "Charlie"])),
            Arc::new(Int32Array::from(vec![1, 1, 1])),
            Arc::new(Date32Array::from(date32_values.clone())),
        ])?;
        data_file_writer.write(batch).await?;

[macpie]

I think this is a known "issue" see #891

[liurenjie1024]

In summary, it's not part of iceberg spec to generate partition folder. User may provide its own location generator for do that.

[macpie]

This is just a quick PartitionedLocationGenerator I wrote. Probably not great but it's a start if someone want to generate partition folders.

use iceberg::{
    Error, ErrorKind, spec::TableMetadata,
    writer::file_writer::location_generator::LocationGenerator,
};

const WRITE_DATA_LOCATION: &str = "write.data.path";
const WRITE_FOLDER_STORAGE_LOCATION: &str = "write.folder-storage.path";
const DEFAULT_DATA_DIR: &str = "/data";

#[derive(Clone, Debug)]
pub struct PartitionedLocationGenerator {
    dir_path: String,
    partitions: Vec<(String, String)>,
}

impl PartitionedLocationGenerator {
    pub fn new(
        table_metadata: TableMetadata,
        partition_spec_id: i32,
        part_values: Vec<String>,
    ) -> anyhow::Result<Self> {
        let part_spec = table_metadata
            .partition_spec_by_id(partition_spec_id)
            .ok_or_else(|| anyhow::anyhow!("Partition spec {} not found", partition_spec_id))?;

        let part_fields = part_spec.fields();

        if part_fields.len() != part_values.len() {
            anyhow::bail!(
                "Partition spec {} mismatch: {} fields vs {} values",
                partition_spec_id,
                part_fields.len(),
                part_values.len()
            );
        }

        let mut partitions = vec![];

        for (i, part_field) in part_fields.into_iter().enumerate() {
            partitions.push((part_field.name.clone(), part_values.get(i).unwrap().clone()));
        }

        let table_location = table_metadata.location();
        let rel_dir_path = {
            let prop = table_metadata.properties();
            let data_location = prop
                .get(WRITE_DATA_LOCATION)
                .or(prop.get(WRITE_FOLDER_STORAGE_LOCATION));
            if let Some(data_location) = data_location {
                data_location.strip_prefix(table_location).ok_or_else(|| {
                    Error::new(
                        ErrorKind::DataInvalid,
                        format!(
                            "data location {} is not a subpath of table location {}",
                            data_location, table_location
                        ),
                    )
                })?
            } else {
                DEFAULT_DATA_DIR
            }
        };

        Ok(Self {
            dir_path: format!("{}{}", table_location, rel_dir_path),
            partitions,
        })
    }
}

impl LocationGenerator for PartitionedLocationGenerator {
    fn generate_location(&self, file_name: &str) -> String {
        let part_dir = self
            .partitions
            .iter()
            .map(|(k, v)| format!("{}={}", k, v))
            .collect::<Vec<_>>()
            .join("/");

        format!("{}/{}/{}", self.dir_path, part_dir, file_name)
    }
}