|
20 | 20 | */
|
21 | 21 | use std::sync::Arc;
|
22 | 22 |
|
| 23 | +use itertools::Itertools; |
23 | 24 | use serde::{Deserialize, Serialize};
|
24 | 25 | use typed_builder::TypedBuilder;
|
25 | 26 |
|
26 | 27 | use super::transform::Transform;
|
27 | 28 | use super::{NestedField, Schema, SchemaRef, StructType};
|
| 29 | +use crate::spec::Struct; |
28 | 30 | use crate::{Error, ErrorKind, Result};
|
29 | 31 |
|
30 | 32 | pub(crate) const UNPARTITIONED_LAST_ASSIGNED_ID: i32 = 999;
|
@@ -152,6 +154,30 @@ impl PartitionSpec {
|
152 | 154 |
|
153 | 155 | true
|
154 | 156 | }
|
| 157 | + |
| 158 | + pub(crate) fn partition_to_path(&self, data: &Struct, schema: SchemaRef) -> String { |
| 159 | + let partition_type = self.partition_type(&schema).unwrap(); |
| 160 | + let field_types = partition_type.fields(); |
| 161 | + |
| 162 | + self.fields |
| 163 | + .iter() |
| 164 | + .enumerate() |
| 165 | + .map(|(i, field)| { |
| 166 | + let value = if data.is_null_at_index(i) { |
| 167 | + None |
| 168 | + } else { |
| 169 | + Some(&data.fields()[i]) |
| 170 | + }; |
| 171 | + format!( |
| 172 | + "{}={}", |
| 173 | + field.name, |
| 174 | + field |
| 175 | + .transform |
| 176 | + .to_human_string(&field_types[i].field_type, value) |
| 177 | + ) |
| 178 | + }) |
| 179 | + .join("/") |
| 180 | + } |
155 | 181 | }
|
156 | 182 |
|
157 | 183 | /// Reference to [`UnboundPartitionSpec`].
|
@@ -660,7 +686,7 @@ impl CorePartitionSpecValidator for UnboundPartitionSpecBuilder {
|
660 | 686 | #[cfg(test)]
|
661 | 687 | mod tests {
|
662 | 688 | use super::*;
|
663 |
| - use crate::spec::{PrimitiveType, Type}; |
| 689 | + use crate::spec::{Literal, PrimitiveType, Type}; |
664 | 690 |
|
665 | 691 | #[test]
|
666 | 692 | fn test_partition_spec() {
|
@@ -1733,4 +1759,30 @@ mod tests {
|
1733 | 1759 | assert_eq!(1002, spec.fields[1].field_id);
|
1734 | 1760 | assert!(!spec.has_sequential_ids());
|
1735 | 1761 | }
|
| 1762 | + |
| 1763 | + #[test] |
| 1764 | + fn test_partition_to_path() { |
| 1765 | + let schema = Schema::builder() |
| 1766 | + .with_fields(vec![ |
| 1767 | + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), |
| 1768 | + NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(), |
| 1769 | + ]) |
| 1770 | + .build() |
| 1771 | + .unwrap(); |
| 1772 | + |
| 1773 | + let spec = PartitionSpec::builder(schema.clone()) |
| 1774 | + .add_partition_field("id", "id", Transform::Identity) |
| 1775 | + .unwrap() |
| 1776 | + .add_partition_field("name", "name", Transform::Identity) |
| 1777 | + .unwrap() |
| 1778 | + .build() |
| 1779 | + .unwrap(); |
| 1780 | + |
| 1781 | + let data = Struct::from_iter([Some(Literal::int(42)), Some(Literal::string("alice"))]); |
| 1782 | + |
| 1783 | + assert_eq!( |
| 1784 | + spec.partition_to_path(&data, schema.into()), |
| 1785 | + "id=42/name=\"alice\"" |
| 1786 | + ); |
| 1787 | + } |
1736 | 1788 | }
|
0 commit comments