1use std::sync::Arc;
22
23use itertools::Itertools;
24use serde::{Deserialize, Serialize};
25use typed_builder::TypedBuilder;
26
27use super::transform::Transform;
28use super::{NestedField, Schema, SchemaRef, StructType};
29use crate::spec::Struct;
30use crate::{Error, ErrorKind, Result};
31
32pub(crate) const UNPARTITIONED_LAST_ASSIGNED_ID: i32 = 999;
33pub(crate) const DEFAULT_PARTITION_SPEC_ID: i32 = 0;
34
35#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone, TypedBuilder)]
37#[serde(rename_all = "kebab-case")]
38pub struct PartitionField {
39 pub source_id: i32,
41 pub field_id: i32,
44 pub name: String,
46 pub transform: Transform,
48}
49
50impl PartitionField {
51 pub fn into_unbound(self) -> UnboundPartitionField {
53 self.into()
54 }
55}
56
57pub type PartitionSpecRef = Arc<PartitionSpec>;
59#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
66#[serde(rename_all = "kebab-case")]
67pub struct PartitionSpec {
68 spec_id: i32,
70 fields: Vec<PartitionField>,
72}
73
74impl PartitionSpec {
75 pub fn builder(schema: impl Into<SchemaRef>) -> PartitionSpecBuilder {
77 PartitionSpecBuilder::new(schema)
78 }
79
80 pub fn fields(&self) -> &[PartitionField] {
82 &self.fields
83 }
84
85 pub fn spec_id(&self) -> i32 {
87 self.spec_id
88 }
89
90 pub fn unpartition_spec() -> Self {
92 Self {
93 spec_id: DEFAULT_PARTITION_SPEC_ID,
94 fields: vec![],
95 }
96 }
97
98 pub fn is_unpartitioned(&self) -> bool {
102 self.fields.is_empty() || self.fields.iter().all(|f| f.transform == Transform::Void)
103 }
104
105 pub fn partition_type(&self, schema: &Schema) -> Result<StructType> {
107 PartitionSpecBuilder::partition_type(&self.fields, schema)
108 }
109
110 pub fn into_unbound(self) -> UnboundPartitionSpec {
112 self.into()
113 }
114
115 pub fn with_spec_id(self, spec_id: i32) -> Self {
117 Self { spec_id, ..self }
118 }
119
120 pub fn has_sequential_ids(&self) -> bool {
124 has_sequential_ids(self.fields.iter().map(|f| f.field_id))
125 }
126
127 pub fn highest_field_id(&self) -> Option<i32> {
129 self.fields.iter().map(|f| f.field_id).max()
130 }
131
132 pub fn is_compatible_with(&self, other: &PartitionSpec) -> bool {
142 if self.fields.len() != other.fields.len() {
143 return false;
144 }
145
146 for (this_field, other_field) in self.fields.iter().zip(other.fields.iter()) {
147 if this_field.source_id != other_field.source_id
148 || this_field.name != other_field.name
149 || this_field.transform != other_field.transform
150 {
151 return false;
152 }
153 }
154
155 true
156 }
157
158 pub fn partition_to_path(&self, data: &Struct, schema: SchemaRef) -> String {
161 let partition_type = self.partition_type(&schema).unwrap();
162 let field_types = partition_type.fields();
163
164 self.fields
165 .iter()
166 .enumerate()
167 .map(|(i, field)| {
168 let value = data[i].as_ref();
169 form_urlencoded::Serializer::new(String::new())
170 .append_pair(
171 &field.name,
172 &field
173 .transform
174 .to_human_string(&field_types[i].field_type, value),
175 )
176 .finish()
177 })
178 .join("/")
179 }
180}
181
182#[derive(Clone, Debug)]
185pub struct PartitionKey {
186 spec: PartitionSpec,
188 schema: SchemaRef,
190 data: Struct,
192}
193
194impl PartitionKey {
195 pub fn new(spec: PartitionSpec, schema: SchemaRef, data: Struct) -> Self {
197 Self { spec, schema, data }
198 }
199
200 pub fn copy_with_data(&self, data: Struct) -> Self {
202 Self {
203 spec: self.spec.clone(),
204 schema: self.schema.clone(),
205 data,
206 }
207 }
208
209 pub fn to_path(&self) -> String {
211 self.spec.partition_to_path(&self.data, self.schema.clone())
212 }
213
214 pub fn is_effectively_none(partition_key: Option<&PartitionKey>) -> bool {
217 match partition_key {
218 None => true,
219 Some(pk) => pk.spec.is_unpartitioned(),
220 }
221 }
222
223 pub fn spec(&self) -> &PartitionSpec {
225 &self.spec
226 }
227
228 pub fn schema(&self) -> &SchemaRef {
230 &self.schema
231 }
232
233 pub fn data(&self) -> &Struct {
235 &self.data
236 }
237}
238
239pub type UnboundPartitionSpecRef = Arc<UnboundPartitionSpec>;
241#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone, TypedBuilder)]
243#[serde(rename_all = "kebab-case")]
244pub struct UnboundPartitionField {
245 pub source_id: i32,
247 #[builder(default, setter(strip_option(fallback = field_id_opt)))]
250 #[serde(skip_serializing_if = "Option::is_none")]
251 pub field_id: Option<i32>,
252 pub name: String,
254 pub transform: Transform,
256}
257
258#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone, Default)]
262#[serde(rename_all = "kebab-case")]
263pub struct UnboundPartitionSpec {
264 #[serde(skip_serializing_if = "Option::is_none")]
266 pub(crate) spec_id: Option<i32>,
267 pub(crate) fields: Vec<UnboundPartitionField>,
269}
270
271impl UnboundPartitionSpec {
272 pub fn builder() -> UnboundPartitionSpecBuilder {
274 UnboundPartitionSpecBuilder::default()
275 }
276
277 pub fn bind(self, schema: impl Into<SchemaRef>) -> Result<PartitionSpec> {
279 PartitionSpecBuilder::new_from_unbound(self, schema)?.build()
280 }
281
282 pub fn spec_id(&self) -> Option<i32> {
284 self.spec_id
285 }
286
287 pub fn fields(&self) -> &[UnboundPartitionField] {
289 &self.fields
290 }
291
292 pub fn with_spec_id(self, spec_id: i32) -> Self {
294 Self {
295 spec_id: Some(spec_id),
296 ..self
297 }
298 }
299}
300
301fn has_sequential_ids(field_ids: impl Iterator<Item = i32>) -> bool {
302 for (index, field_id) in field_ids.enumerate() {
303 let expected_id = (UNPARTITIONED_LAST_ASSIGNED_ID as i64)
304 .checked_add(1)
305 .and_then(|id| id.checked_add(index as i64))
306 .unwrap_or(i64::MAX);
307
308 if field_id as i64 != expected_id {
309 return false;
310 }
311 }
312
313 true
314}
315
316impl From<PartitionField> for UnboundPartitionField {
317 fn from(field: PartitionField) -> Self {
318 UnboundPartitionField {
319 source_id: field.source_id,
320 field_id: Some(field.field_id),
321 name: field.name,
322 transform: field.transform,
323 }
324 }
325}
326
327impl From<PartitionSpec> for UnboundPartitionSpec {
328 fn from(spec: PartitionSpec) -> Self {
329 UnboundPartitionSpec {
330 spec_id: Some(spec.spec_id),
331 fields: spec.fields.into_iter().map(Into::into).collect(),
332 }
333 }
334}
335
336#[derive(Debug, Default)]
338pub struct UnboundPartitionSpecBuilder {
339 spec_id: Option<i32>,
340 fields: Vec<UnboundPartitionField>,
341}
342
343impl UnboundPartitionSpecBuilder {
344 pub fn new() -> Self {
346 Self {
347 spec_id: None,
348 fields: vec![],
349 }
350 }
351
352 pub fn with_spec_id(mut self, spec_id: i32) -> Self {
354 self.spec_id = Some(spec_id);
355 self
356 }
357
358 pub fn add_partition_field(
360 self,
361 source_id: i32,
362 target_name: impl ToString,
363 transformation: Transform,
364 ) -> Result<Self> {
365 let field = UnboundPartitionField {
366 source_id,
367 field_id: None,
368 name: target_name.to_string(),
369 transform: transformation,
370 };
371 self.add_partition_field_internal(field)
372 }
373
374 pub fn add_partition_fields(
376 self,
377 fields: impl IntoIterator<Item = UnboundPartitionField>,
378 ) -> Result<Self> {
379 let mut builder = self;
380 for field in fields {
381 builder = builder.add_partition_field_internal(field)?;
382 }
383 Ok(builder)
384 }
385
386 fn add_partition_field_internal(mut self, field: UnboundPartitionField) -> Result<Self> {
387 self.check_name_set_and_unique(&field.name)?;
388 self.check_for_redundant_partitions(field.source_id, &field.transform)?;
389 if let Some(partition_field_id) = field.field_id {
390 self.check_partition_id_unique(partition_field_id)?;
391 }
392 self.fields.push(field);
393 Ok(self)
394 }
395
396 pub fn build(self) -> UnboundPartitionSpec {
398 UnboundPartitionSpec {
399 spec_id: self.spec_id,
400 fields: self.fields,
401 }
402 }
403}
404
405#[derive(Debug)]
407pub struct PartitionSpecBuilder {
408 spec_id: Option<i32>,
409 last_assigned_field_id: i32,
410 fields: Vec<UnboundPartitionField>,
411 schema: SchemaRef,
412}
413
414impl PartitionSpecBuilder {
415 pub fn new(schema: impl Into<SchemaRef>) -> Self {
417 Self {
418 spec_id: None,
419 fields: vec![],
420 last_assigned_field_id: UNPARTITIONED_LAST_ASSIGNED_ID,
421 schema: schema.into(),
422 }
423 }
424
425 pub fn new_from_unbound(
427 unbound: UnboundPartitionSpec,
428 schema: impl Into<SchemaRef>,
429 ) -> Result<Self> {
430 let mut builder =
431 Self::new(schema).with_spec_id(unbound.spec_id.unwrap_or(DEFAULT_PARTITION_SPEC_ID));
432
433 for field in unbound.fields {
434 builder = builder.add_unbound_field(field)?;
435 }
436 Ok(builder)
437 }
438
439 pub fn with_last_assigned_field_id(mut self, last_assigned_field_id: i32) -> Self {
445 self.last_assigned_field_id = last_assigned_field_id;
446 self
447 }
448
449 pub fn with_spec_id(mut self, spec_id: i32) -> Self {
451 self.spec_id = Some(spec_id);
452 self
453 }
454
455 pub fn add_partition_field(
457 self,
458 source_name: impl AsRef<str>,
459 target_name: impl Into<String>,
460 transform: Transform,
461 ) -> Result<Self> {
462 let source_id = self
463 .schema
464 .field_by_name(source_name.as_ref())
465 .ok_or_else(|| {
466 Error::new(
467 ErrorKind::DataInvalid,
468 format!(
469 "Cannot find source column with name: {} in schema",
470 source_name.as_ref()
471 ),
472 )
473 })?
474 .id;
475 let field = UnboundPartitionField {
476 source_id,
477 field_id: None,
478 name: target_name.into(),
479 transform,
480 };
481
482 self.add_unbound_field(field)
483 }
484
485 pub fn add_unbound_field(mut self, field: UnboundPartitionField) -> Result<Self> {
490 self.check_name_set_and_unique(&field.name)?;
491 self.check_for_redundant_partitions(field.source_id, &field.transform)?;
492 Self::check_name_does_not_collide_with_schema(&field, &self.schema)?;
493 Self::check_transform_compatibility(&field, &self.schema)?;
494 if let Some(partition_field_id) = field.field_id {
495 self.check_partition_id_unique(partition_field_id)?;
496 }
497
498 self.fields.push(field);
500 Ok(self)
501 }
502
503 pub fn add_unbound_fields(
505 self,
506 fields: impl IntoIterator<Item = UnboundPartitionField>,
507 ) -> Result<Self> {
508 let mut builder = self;
509 for field in fields {
510 builder = builder.add_unbound_field(field)?;
511 }
512 Ok(builder)
513 }
514
515 pub fn build(self) -> Result<PartitionSpec> {
517 let fields = Self::set_field_ids(self.fields, self.last_assigned_field_id)?;
518 Ok(PartitionSpec {
519 spec_id: self.spec_id.unwrap_or(DEFAULT_PARTITION_SPEC_ID),
520 fields,
521 })
522 }
523
524 fn set_field_ids(
525 fields: Vec<UnboundPartitionField>,
526 last_assigned_field_id: i32,
527 ) -> Result<Vec<PartitionField>> {
528 let mut last_assigned_field_id = last_assigned_field_id;
529 let assigned_ids = fields
532 .iter()
533 .filter_map(|f| f.field_id)
534 .collect::<std::collections::HashSet<_>>();
535
536 fn _check_add_1(prev: i32) -> Result<i32> {
537 prev.checked_add(1).ok_or_else(|| {
538 Error::new(
539 ErrorKind::DataInvalid,
540 "Cannot assign more partition ids. Overflow.",
541 )
542 })
543 }
544
545 let mut bound_fields = Vec::with_capacity(fields.len());
546 for field in fields.into_iter() {
547 let partition_field_id = if let Some(partition_field_id) = field.field_id {
548 last_assigned_field_id = std::cmp::max(last_assigned_field_id, partition_field_id);
549 partition_field_id
550 } else {
551 last_assigned_field_id = _check_add_1(last_assigned_field_id)?;
552 while assigned_ids.contains(&last_assigned_field_id) {
553 last_assigned_field_id = _check_add_1(last_assigned_field_id)?;
554 }
555 last_assigned_field_id
556 };
557
558 bound_fields.push(PartitionField {
559 source_id: field.source_id,
560 field_id: partition_field_id,
561 name: field.name,
562 transform: field.transform,
563 })
564 }
565
566 Ok(bound_fields)
567 }
568
569 fn partition_type(fields: &Vec<PartitionField>, schema: &Schema) -> Result<StructType> {
571 let mut struct_fields = Vec::with_capacity(fields.len());
572 for partition_field in fields {
573 let field = schema
574 .field_by_id(partition_field.source_id)
575 .ok_or_else(|| {
576 Error::new(
577 ErrorKind::Unexpected,
580 format!(
581 "No column with source column id {} in schema {:?}",
582 partition_field.source_id, schema
583 ),
584 )
585 })?;
586 let res_type = partition_field.transform.result_type(&field.field_type)?;
587 let field =
588 NestedField::optional(partition_field.field_id, &partition_field.name, res_type)
589 .into();
590 struct_fields.push(field);
591 }
592 Ok(StructType::new(struct_fields))
593 }
594
595 fn check_name_does_not_collide_with_schema(
600 field: &UnboundPartitionField,
601 schema: &Schema,
602 ) -> Result<()> {
603 match schema.field_by_name(field.name.as_str()) {
604 Some(schema_collision) => {
605 if field.transform == Transform::Identity {
606 if schema_collision.id == field.source_id {
607 Ok(())
608 } else {
609 Err(Error::new(
610 ErrorKind::DataInvalid,
611 format!(
612 "Cannot create identity partition sourced from different field in schema. Field name '{}' has id `{}` in schema but partition source id is `{}`",
613 field.name, schema_collision.id, field.source_id
614 ),
615 ))
616 }
617 } else {
618 Err(Error::new(
619 ErrorKind::DataInvalid,
620 format!(
621 "Cannot create partition with name: '{}' that conflicts with schema field and is not an identity transform.",
622 field.name
623 ),
624 ))
625 }
626 }
627 None => Ok(()),
628 }
629 }
630
631 fn check_transform_compatibility(field: &UnboundPartitionField, schema: &Schema) -> Result<()> {
634 let schema_field = schema.field_by_id(field.source_id).ok_or_else(|| {
635 Error::new(
636 ErrorKind::DataInvalid,
637 format!(
638 "Cannot find partition source field with id `{}` in schema",
639 field.source_id
640 ),
641 )
642 })?;
643
644 if field.transform != Transform::Void {
645 if !schema_field.field_type.is_primitive() {
646 return Err(Error::new(
647 ErrorKind::DataInvalid,
648 format!(
649 "Cannot partition by non-primitive source field: '{}'.",
650 schema_field.field_type
651 ),
652 ));
653 }
654
655 if field
656 .transform
657 .result_type(&schema_field.field_type)
658 .is_err()
659 {
660 return Err(Error::new(
661 ErrorKind::DataInvalid,
662 format!(
663 "Invalid source type: '{}' for transform: '{}'.",
664 schema_field.field_type,
665 field.transform.dedup_name()
666 ),
667 ));
668 }
669 }
670
671 Ok(())
672 }
673}
674
675trait CorePartitionSpecValidator {
677 fn check_name_set_and_unique(&self, name: &str) -> Result<()> {
679 if name.is_empty() {
680 return Err(Error::new(
681 ErrorKind::DataInvalid,
682 "Cannot use empty partition name",
683 ));
684 }
685
686 if self.fields().iter().any(|f| f.name == name) {
687 return Err(Error::new(
688 ErrorKind::DataInvalid,
689 format!("Cannot use partition name more than once: {name}"),
690 ));
691 }
692 Ok(())
693 }
694
695 fn check_for_redundant_partitions(&self, source_id: i32, transform: &Transform) -> Result<()> {
697 let collision = self.fields().iter().find(|f| {
698 f.source_id == source_id && f.transform.dedup_name() == transform.dedup_name()
699 });
700
701 if let Some(collision) = collision {
702 Err(Error::new(
703 ErrorKind::DataInvalid,
704 format!(
705 "Cannot add redundant partition with source id `{}` and transform `{}`. A partition with the same source id and transform already exists with name `{}`",
706 source_id,
707 transform.dedup_name(),
708 collision.name
709 ),
710 ))
711 } else {
712 Ok(())
713 }
714 }
715
716 fn check_partition_id_unique(&self, field_id: i32) -> Result<()> {
718 if self.fields().iter().any(|f| f.field_id == Some(field_id)) {
719 return Err(Error::new(
720 ErrorKind::DataInvalid,
721 format!("Cannot use field id more than once in one PartitionSpec: {field_id}"),
722 ));
723 }
724
725 Ok(())
726 }
727
728 fn fields(&self) -> &Vec<UnboundPartitionField>;
729}
730
731impl CorePartitionSpecValidator for PartitionSpecBuilder {
732 fn fields(&self) -> &Vec<UnboundPartitionField> {
733 &self.fields
734 }
735}
736
737impl CorePartitionSpecValidator for UnboundPartitionSpecBuilder {
738 fn fields(&self) -> &Vec<UnboundPartitionField> {
739 &self.fields
740 }
741}
742
743#[cfg(test)]
744mod tests {
745 use super::*;
746 use crate::spec::{Literal, PrimitiveType, Type};
747
748 #[test]
749 fn test_partition_spec() {
750 let spec = r#"
751 {
752 "spec-id": 1,
753 "fields": [ {
754 "source-id": 4,
755 "field-id": 1000,
756 "name": "ts_day",
757 "transform": "day"
758 }, {
759 "source-id": 1,
760 "field-id": 1001,
761 "name": "id_bucket",
762 "transform": "bucket[16]"
763 }, {
764 "source-id": 2,
765 "field-id": 1002,
766 "name": "id_truncate",
767 "transform": "truncate[4]"
768 } ]
769 }
770 "#;
771
772 let partition_spec: PartitionSpec = serde_json::from_str(spec).unwrap();
773 assert_eq!(4, partition_spec.fields[0].source_id);
774 assert_eq!(1000, partition_spec.fields[0].field_id);
775 assert_eq!("ts_day", partition_spec.fields[0].name);
776 assert_eq!(Transform::Day, partition_spec.fields[0].transform);
777
778 assert_eq!(1, partition_spec.fields[1].source_id);
779 assert_eq!(1001, partition_spec.fields[1].field_id);
780 assert_eq!("id_bucket", partition_spec.fields[1].name);
781 assert_eq!(Transform::Bucket(16), partition_spec.fields[1].transform);
782
783 assert_eq!(2, partition_spec.fields[2].source_id);
784 assert_eq!(1002, partition_spec.fields[2].field_id);
785 assert_eq!("id_truncate", partition_spec.fields[2].name);
786 assert_eq!(Transform::Truncate(4), partition_spec.fields[2].transform);
787 }
788
789 #[test]
790 fn test_is_unpartitioned() {
791 let schema = Schema::builder()
792 .with_fields(vec![
793 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
794 NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
795 ])
796 .build()
797 .unwrap();
798 let partition_spec = PartitionSpec::builder(schema.clone())
799 .with_spec_id(1)
800 .build()
801 .unwrap();
802 assert!(
803 partition_spec.is_unpartitioned(),
804 "Empty partition spec should be unpartitioned"
805 );
806
807 let partition_spec = PartitionSpec::builder(schema.clone())
808 .add_unbound_fields(vec![
809 UnboundPartitionField::builder()
810 .source_id(1)
811 .name("id".to_string())
812 .transform(Transform::Identity)
813 .build(),
814 UnboundPartitionField::builder()
815 .source_id(2)
816 .name("name_string".to_string())
817 .transform(Transform::Void)
818 .build(),
819 ])
820 .unwrap()
821 .with_spec_id(1)
822 .build()
823 .unwrap();
824 assert!(
825 !partition_spec.is_unpartitioned(),
826 "Partition spec with one non void transform should not be unpartitioned"
827 );
828
829 let partition_spec = PartitionSpec::builder(schema.clone())
830 .with_spec_id(1)
831 .add_unbound_fields(vec![
832 UnboundPartitionField::builder()
833 .source_id(1)
834 .name("id_void".to_string())
835 .transform(Transform::Void)
836 .build(),
837 UnboundPartitionField::builder()
838 .source_id(2)
839 .name("name_void".to_string())
840 .transform(Transform::Void)
841 .build(),
842 ])
843 .unwrap()
844 .build()
845 .unwrap();
846 assert!(
847 partition_spec.is_unpartitioned(),
848 "Partition spec with all void field should be unpartitioned"
849 );
850 }
851
852 #[test]
853 fn test_unbound_partition_spec() {
854 let spec = r#"
855 {
856 "spec-id": 1,
857 "fields": [ {
858 "source-id": 4,
859 "field-id": 1000,
860 "name": "ts_day",
861 "transform": "day"
862 }, {
863 "source-id": 1,
864 "field-id": 1001,
865 "name": "id_bucket",
866 "transform": "bucket[16]"
867 }, {
868 "source-id": 2,
869 "field-id": 1002,
870 "name": "id_truncate",
871 "transform": "truncate[4]"
872 } ]
873 }
874 "#;
875
876 let partition_spec: UnboundPartitionSpec = serde_json::from_str(spec).unwrap();
877 assert_eq!(Some(1), partition_spec.spec_id);
878
879 assert_eq!(4, partition_spec.fields[0].source_id);
880 assert_eq!(Some(1000), partition_spec.fields[0].field_id);
881 assert_eq!("ts_day", partition_spec.fields[0].name);
882 assert_eq!(Transform::Day, partition_spec.fields[0].transform);
883
884 assert_eq!(1, partition_spec.fields[1].source_id);
885 assert_eq!(Some(1001), partition_spec.fields[1].field_id);
886 assert_eq!("id_bucket", partition_spec.fields[1].name);
887 assert_eq!(Transform::Bucket(16), partition_spec.fields[1].transform);
888
889 assert_eq!(2, partition_spec.fields[2].source_id);
890 assert_eq!(Some(1002), partition_spec.fields[2].field_id);
891 assert_eq!("id_truncate", partition_spec.fields[2].name);
892 assert_eq!(Transform::Truncate(4), partition_spec.fields[2].transform);
893
894 let spec = r#"
895 {
896 "fields": [ {
897 "source-id": 4,
898 "name": "ts_day",
899 "transform": "day"
900 } ]
901 }
902 "#;
903 let partition_spec: UnboundPartitionSpec = serde_json::from_str(spec).unwrap();
904 assert_eq!(None, partition_spec.spec_id);
905
906 assert_eq!(4, partition_spec.fields[0].source_id);
907 assert_eq!(None, partition_spec.fields[0].field_id);
908 assert_eq!("ts_day", partition_spec.fields[0].name);
909 assert_eq!(Transform::Day, partition_spec.fields[0].transform);
910 }
911
912 #[test]
913 fn test_unbound_partition_spec_serialization_skips_none_fields() {
914 let spec = UnboundPartitionSpec::builder()
915 .add_partition_field(4, "ts_day".to_string(), Transform::Day)
916 .unwrap()
917 .build();
918
919 let value = serde_json::to_value(&spec).unwrap();
920 let object = value.as_object().unwrap();
921 assert!(!object.contains_key("spec-id"));
922 let field = object["fields"][0].as_object().unwrap();
923 assert!(!field.contains_key("field-id"));
924
925 let value = serde_json::to_value(spec.with_spec_id(1)).unwrap();
926 let object = value.as_object().unwrap();
927 assert_eq!(Some(&serde_json::json!(1)), object.get("spec-id"));
928
929 let spec: UnboundPartitionSpec = serde_json::from_str(
932 r#"{
933 "spec-id": null,
934 "fields": [
935 {"source-id": 4, "name": "ts_day", "transform": "day", "field-id": null}
936 ]
937 }"#,
938 )
939 .unwrap();
940 assert_eq!(None, spec.spec_id);
941 assert_eq!(None, spec.fields[0].field_id);
942 }
943
944 #[test]
945 fn test_new_unpartition() {
946 let schema = Schema::builder()
947 .with_fields(vec![
948 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
949 NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
950 ])
951 .build()
952 .unwrap();
953 let partition_spec = PartitionSpec::builder(schema.clone())
954 .with_spec_id(0)
955 .build()
956 .unwrap();
957 let partition_type = partition_spec.partition_type(&schema).unwrap();
958 assert_eq!(0, partition_type.fields().len());
959
960 let unpartition_spec = PartitionSpec::unpartition_spec();
961 assert_eq!(partition_spec, unpartition_spec);
962 }
963
964 #[test]
965 fn test_partition_type() {
966 let spec = r#"
967 {
968 "spec-id": 1,
969 "fields": [ {
970 "source-id": 4,
971 "field-id": 1000,
972 "name": "ts_day",
973 "transform": "day"
974 }, {
975 "source-id": 1,
976 "field-id": 1001,
977 "name": "id_bucket",
978 "transform": "bucket[16]"
979 }, {
980 "source-id": 2,
981 "field-id": 1002,
982 "name": "id_truncate",
983 "transform": "truncate[4]"
984 } ]
985 }
986 "#;
987
988 let partition_spec: PartitionSpec = serde_json::from_str(spec).unwrap();
989 let schema = Schema::builder()
990 .with_fields(vec![
991 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
992 NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
993 NestedField::required(3, "ts", Type::Primitive(PrimitiveType::Timestamp)).into(),
994 NestedField::required(4, "ts_day", Type::Primitive(PrimitiveType::Timestamp))
995 .into(),
996 NestedField::required(5, "id_bucket", Type::Primitive(PrimitiveType::Int)).into(),
997 NestedField::required(6, "id_truncate", Type::Primitive(PrimitiveType::Int)).into(),
998 ])
999 .build()
1000 .unwrap();
1001
1002 let partition_type = partition_spec.partition_type(&schema).unwrap();
1003 assert_eq!(3, partition_type.fields().len());
1004 assert_eq!(
1005 *partition_type.fields()[0],
1006 NestedField::optional(
1007 partition_spec.fields[0].field_id,
1008 &partition_spec.fields[0].name,
1009 Type::Primitive(PrimitiveType::Date)
1010 )
1011 );
1012 assert_eq!(
1013 *partition_type.fields()[1],
1014 NestedField::optional(
1015 partition_spec.fields[1].field_id,
1016 &partition_spec.fields[1].name,
1017 Type::Primitive(PrimitiveType::Int)
1018 )
1019 );
1020 assert_eq!(
1021 *partition_type.fields()[2],
1022 NestedField::optional(
1023 partition_spec.fields[2].field_id,
1024 &partition_spec.fields[2].name,
1025 Type::Primitive(PrimitiveType::String)
1026 )
1027 );
1028 }
1029
1030 #[test]
1031 fn test_partition_empty() {
1032 let spec = r#"
1033 {
1034 "spec-id": 1,
1035 "fields": []
1036 }
1037 "#;
1038
1039 let partition_spec: PartitionSpec = serde_json::from_str(spec).unwrap();
1040 let schema = Schema::builder()
1041 .with_fields(vec![
1042 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1043 NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
1044 NestedField::required(3, "ts", Type::Primitive(PrimitiveType::Timestamp)).into(),
1045 NestedField::required(4, "ts_day", Type::Primitive(PrimitiveType::Timestamp))
1046 .into(),
1047 NestedField::required(5, "id_bucket", Type::Primitive(PrimitiveType::Int)).into(),
1048 NestedField::required(6, "id_truncate", Type::Primitive(PrimitiveType::Int)).into(),
1049 ])
1050 .build()
1051 .unwrap();
1052
1053 let partition_type = partition_spec.partition_type(&schema).unwrap();
1054 assert_eq!(0, partition_type.fields().len());
1055 }
1056
1057 #[test]
1058 fn test_partition_error() {
1059 let spec = r#"
1060 {
1061 "spec-id": 1,
1062 "fields": [ {
1063 "source-id": 4,
1064 "field-id": 1000,
1065 "name": "ts_day",
1066 "transform": "day"
1067 }, {
1068 "source-id": 1,
1069 "field-id": 1001,
1070 "name": "id_bucket",
1071 "transform": "bucket[16]"
1072 }, {
1073 "source-id": 2,
1074 "field-id": 1002,
1075 "name": "id_truncate",
1076 "transform": "truncate[4]"
1077 } ]
1078 }
1079 "#;
1080
1081 let partition_spec: PartitionSpec = serde_json::from_str(spec).unwrap();
1082 let schema = Schema::builder()
1083 .with_fields(vec![
1084 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1085 NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
1086 ])
1087 .build()
1088 .unwrap();
1089
1090 assert!(partition_spec.partition_type(&schema).is_err());
1091 }
1092
1093 #[test]
1094 fn test_builder_disallow_duplicate_names() {
1095 UnboundPartitionSpec::builder()
1096 .add_partition_field(1, "ts_day".to_string(), Transform::Day)
1097 .unwrap()
1098 .add_partition_field(2, "ts_day".to_string(), Transform::Day)
1099 .unwrap_err();
1100 }
1101
1102 #[test]
1103 fn test_builder_disallow_duplicate_field_ids() {
1104 let schema = Schema::builder()
1105 .with_fields(vec![
1106 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1107 NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
1108 ])
1109 .build()
1110 .unwrap();
1111 PartitionSpec::builder(schema.clone())
1112 .add_unbound_field(UnboundPartitionField {
1113 source_id: 1,
1114 field_id: Some(1000),
1115 name: "id".to_string(),
1116 transform: Transform::Identity,
1117 })
1118 .unwrap()
1119 .add_unbound_field(UnboundPartitionField {
1120 source_id: 2,
1121 field_id: Some(1000),
1122 name: "id_bucket".to_string(),
1123 transform: Transform::Bucket(16),
1124 })
1125 .unwrap_err();
1126 }
1127
1128 #[test]
1129 fn test_builder_auto_assign_field_ids() {
1130 let schema = Schema::builder()
1131 .with_fields(vec![
1132 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1133 NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
1134 NestedField::required(3, "ts", Type::Primitive(PrimitiveType::Timestamp)).into(),
1135 ])
1136 .build()
1137 .unwrap();
1138 let spec = PartitionSpec::builder(schema.clone())
1139 .with_spec_id(1)
1140 .add_unbound_field(UnboundPartitionField {
1141 source_id: 1,
1142 name: "id".to_string(),
1143 transform: Transform::Identity,
1144 field_id: Some(1012),
1145 })
1146 .unwrap()
1147 .add_unbound_field(UnboundPartitionField {
1148 source_id: 2,
1149 name: "name_void".to_string(),
1150 transform: Transform::Void,
1151 field_id: None,
1152 })
1153 .unwrap()
1154 .add_unbound_field(UnboundPartitionField {
1156 source_id: 3,
1157 name: "year".to_string(),
1158 transform: Transform::Year,
1159 field_id: Some(1),
1160 })
1161 .unwrap()
1162 .build()
1163 .unwrap();
1164
1165 assert_eq!(1012, spec.fields[0].field_id);
1166 assert_eq!(1013, spec.fields[1].field_id);
1167 assert_eq!(1, spec.fields[2].field_id);
1168 }
1169
1170 #[test]
1171 fn test_builder_valid_schema() {
1172 let schema = Schema::builder()
1173 .with_fields(vec![
1174 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1175 NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
1176 ])
1177 .build()
1178 .unwrap();
1179
1180 PartitionSpec::builder(schema.clone())
1181 .with_spec_id(1)
1182 .build()
1183 .unwrap();
1184
1185 let spec = PartitionSpec::builder(schema.clone())
1186 .with_spec_id(1)
1187 .add_partition_field("id", "id_bucket[16]", Transform::Bucket(16))
1188 .unwrap()
1189 .build()
1190 .unwrap();
1191
1192 assert_eq!(spec, PartitionSpec {
1193 spec_id: 1,
1194 fields: vec![PartitionField {
1195 source_id: 1,
1196 field_id: 1000,
1197 name: "id_bucket[16]".to_string(),
1198 transform: Transform::Bucket(16),
1199 }],
1200 });
1201 assert_eq!(
1202 spec.partition_type(&schema).unwrap(),
1203 StructType::new(vec![
1204 NestedField::optional(1000, "id_bucket[16]", Type::Primitive(PrimitiveType::Int))
1205 .into()
1206 ])
1207 )
1208 }
1209
1210 #[test]
1211 fn test_collision_with_schema_name() {
1212 let schema = Schema::builder()
1213 .with_fields(vec![
1214 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1215 ])
1216 .build()
1217 .unwrap();
1218
1219 PartitionSpec::builder(schema.clone())
1220 .with_spec_id(1)
1221 .build()
1222 .unwrap();
1223
1224 let err = PartitionSpec::builder(schema)
1225 .with_spec_id(1)
1226 .add_unbound_field(UnboundPartitionField {
1227 source_id: 1,
1228 field_id: None,
1229 name: "id".to_string(),
1230 transform: Transform::Bucket(16),
1231 })
1232 .unwrap_err();
1233 assert!(err.message().contains("conflicts with schema"))
1234 }
1235
1236 #[test]
1237 fn test_builder_collision_is_ok_for_identity_transforms() {
1238 let schema = Schema::builder()
1239 .with_fields(vec![
1240 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1241 NestedField::required(2, "number", Type::Primitive(PrimitiveType::Int)).into(),
1242 ])
1243 .build()
1244 .unwrap();
1245
1246 PartitionSpec::builder(schema.clone())
1247 .with_spec_id(1)
1248 .build()
1249 .unwrap();
1250
1251 PartitionSpec::builder(schema.clone())
1252 .with_spec_id(1)
1253 .add_unbound_field(UnboundPartitionField {
1254 source_id: 1,
1255 field_id: None,
1256 name: "id".to_string(),
1257 transform: Transform::Identity,
1258 })
1259 .unwrap()
1260 .build()
1261 .unwrap();
1262
1263 PartitionSpec::builder(schema)
1265 .with_spec_id(1)
1266 .add_unbound_field(UnboundPartitionField {
1267 source_id: 2,
1268 field_id: None,
1269 name: "id".to_string(),
1270 transform: Transform::Identity,
1271 })
1272 .unwrap_err();
1273 }
1274
1275 #[test]
1276 fn test_builder_all_source_ids_must_exist() {
1277 let schema = Schema::builder()
1278 .with_fields(vec![
1279 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1280 NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
1281 NestedField::required(3, "ts", Type::Primitive(PrimitiveType::Timestamp)).into(),
1282 ])
1283 .build()
1284 .unwrap();
1285
1286 PartitionSpec::builder(schema.clone())
1288 .with_spec_id(1)
1289 .add_unbound_fields(vec![
1290 UnboundPartitionField {
1291 source_id: 1,
1292 field_id: None,
1293 name: "id_bucket".to_string(),
1294 transform: Transform::Bucket(16),
1295 },
1296 UnboundPartitionField {
1297 source_id: 2,
1298 field_id: None,
1299 name: "name".to_string(),
1300 transform: Transform::Identity,
1301 },
1302 ])
1303 .unwrap()
1304 .build()
1305 .unwrap();
1306
1307 PartitionSpec::builder(schema)
1309 .with_spec_id(1)
1310 .add_unbound_fields(vec![
1311 UnboundPartitionField {
1312 source_id: 1,
1313 field_id: None,
1314 name: "id_bucket".to_string(),
1315 transform: Transform::Bucket(16),
1316 },
1317 UnboundPartitionField {
1318 source_id: 4,
1319 field_id: None,
1320 name: "name".to_string(),
1321 transform: Transform::Identity,
1322 },
1323 ])
1324 .unwrap_err();
1325 }
1326
1327 #[test]
1328 fn test_builder_disallows_variant_source() {
1329 let schema = Schema::builder()
1330 .with_fields(vec![
1331 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1332 NestedField::optional(2, "v", Type::Variant(crate::spec::VariantType)).into(),
1333 ])
1334 .build()
1335 .unwrap();
1336
1337 let err = PartitionSpec::builder(schema)
1338 .with_spec_id(1)
1339 .add_unbound_fields(vec![UnboundPartitionField {
1340 source_id: 2,
1341 field_id: None,
1342 name: "v_part".to_string(),
1343 transform: Transform::Identity,
1344 }])
1345 .expect_err("variant must not be allowed as a partition source");
1346
1347 assert_eq!(
1348 err.message(),
1349 "Cannot partition by non-primitive source field: 'variant'."
1350 );
1351 }
1352
1353 #[test]
1354 fn test_builder_disallows_redundant() {
1355 let err = UnboundPartitionSpec::builder()
1356 .with_spec_id(1)
1357 .add_partition_field(1, "id_bucket[16]".to_string(), Transform::Bucket(16))
1358 .unwrap()
1359 .add_partition_field(
1360 1,
1361 "id_bucket_with_other_name".to_string(),
1362 Transform::Bucket(16),
1363 )
1364 .unwrap_err();
1365 assert!(err.message().contains("redundant partition"));
1366 }
1367
1368 #[test]
1369 fn test_builder_incompatible_transforms_disallowed() {
1370 let schema = Schema::builder()
1371 .with_fields(vec![
1372 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1373 ])
1374 .build()
1375 .unwrap();
1376
1377 PartitionSpec::builder(schema)
1378 .with_spec_id(1)
1379 .add_unbound_field(UnboundPartitionField {
1380 source_id: 1,
1381 field_id: None,
1382 name: "id_year".to_string(),
1383 transform: Transform::Year,
1384 })
1385 .unwrap_err();
1386 }
1387
1388 #[test]
1389 fn test_build_unbound_specs_without_partition_id() {
1390 let spec = UnboundPartitionSpec::builder()
1391 .with_spec_id(1)
1392 .add_partition_fields(vec![UnboundPartitionField {
1393 source_id: 1,
1394 field_id: None,
1395 name: "id_bucket[16]".to_string(),
1396 transform: Transform::Bucket(16),
1397 }])
1398 .unwrap()
1399 .build();
1400
1401 assert_eq!(spec, UnboundPartitionSpec {
1402 spec_id: Some(1),
1403 fields: vec![UnboundPartitionField {
1404 source_id: 1,
1405 field_id: None,
1406 name: "id_bucket[16]".to_string(),
1407 transform: Transform::Bucket(16),
1408 }]
1409 });
1410 }
1411
1412 #[test]
1413 fn test_is_compatible_with() {
1414 let schema = Schema::builder()
1415 .with_fields(vec![
1416 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1417 NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
1418 ])
1419 .build()
1420 .unwrap();
1421
1422 let partition_spec_1 = PartitionSpec::builder(schema.clone())
1423 .with_spec_id(1)
1424 .add_unbound_field(UnboundPartitionField {
1425 source_id: 1,
1426 field_id: None,
1427 name: "id_bucket".to_string(),
1428 transform: Transform::Bucket(16),
1429 })
1430 .unwrap()
1431 .build()
1432 .unwrap();
1433
1434 let partition_spec_2 = PartitionSpec::builder(schema)
1435 .with_spec_id(1)
1436 .add_unbound_field(UnboundPartitionField {
1437 source_id: 1,
1438 field_id: None,
1439 name: "id_bucket".to_string(),
1440 transform: Transform::Bucket(16),
1441 })
1442 .unwrap()
1443 .build()
1444 .unwrap();
1445
1446 assert!(partition_spec_1.is_compatible_with(&partition_spec_2));
1447 }
1448
1449 #[test]
1450 fn test_not_compatible_with_transform_different() {
1451 let schema = Schema::builder()
1452 .with_fields(vec![
1453 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1454 ])
1455 .build()
1456 .unwrap();
1457
1458 let partition_spec_1 = PartitionSpec::builder(schema.clone())
1459 .with_spec_id(1)
1460 .add_unbound_field(UnboundPartitionField {
1461 source_id: 1,
1462 field_id: None,
1463 name: "id_bucket".to_string(),
1464 transform: Transform::Bucket(16),
1465 })
1466 .unwrap()
1467 .build()
1468 .unwrap();
1469
1470 let partition_spec_2 = PartitionSpec::builder(schema)
1471 .with_spec_id(1)
1472 .add_unbound_field(UnboundPartitionField {
1473 source_id: 1,
1474 field_id: None,
1475 name: "id_bucket".to_string(),
1476 transform: Transform::Bucket(32),
1477 })
1478 .unwrap()
1479 .build()
1480 .unwrap();
1481
1482 assert!(!partition_spec_1.is_compatible_with(&partition_spec_2));
1483 }
1484
1485 #[test]
1486 fn test_not_compatible_with_source_id_different() {
1487 let schema = Schema::builder()
1488 .with_fields(vec![
1489 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1490 NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
1491 ])
1492 .build()
1493 .unwrap();
1494
1495 let partition_spec_1 = PartitionSpec::builder(schema.clone())
1496 .with_spec_id(1)
1497 .add_unbound_field(UnboundPartitionField {
1498 source_id: 1,
1499 field_id: None,
1500 name: "id_bucket".to_string(),
1501 transform: Transform::Bucket(16),
1502 })
1503 .unwrap()
1504 .build()
1505 .unwrap();
1506
1507 let partition_spec_2 = PartitionSpec::builder(schema)
1508 .with_spec_id(1)
1509 .add_unbound_field(UnboundPartitionField {
1510 source_id: 2,
1511 field_id: None,
1512 name: "id_bucket".to_string(),
1513 transform: Transform::Bucket(16),
1514 })
1515 .unwrap()
1516 .build()
1517 .unwrap();
1518
1519 assert!(!partition_spec_1.is_compatible_with(&partition_spec_2));
1520 }
1521
1522 #[test]
1523 fn test_not_compatible_with_order_different() {
1524 let schema = Schema::builder()
1525 .with_fields(vec![
1526 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1527 NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
1528 ])
1529 .build()
1530 .unwrap();
1531
1532 let partition_spec_1 = PartitionSpec::builder(schema.clone())
1533 .with_spec_id(1)
1534 .add_unbound_field(UnboundPartitionField {
1535 source_id: 1,
1536 field_id: None,
1537 name: "id_bucket".to_string(),
1538 transform: Transform::Bucket(16),
1539 })
1540 .unwrap()
1541 .add_unbound_field(UnboundPartitionField {
1542 source_id: 2,
1543 field_id: None,
1544 name: "name".to_string(),
1545 transform: Transform::Identity,
1546 })
1547 .unwrap()
1548 .build()
1549 .unwrap();
1550
1551 let partition_spec_2 = PartitionSpec::builder(schema)
1552 .with_spec_id(1)
1553 .add_unbound_field(UnboundPartitionField {
1554 source_id: 2,
1555 field_id: None,
1556 name: "name".to_string(),
1557 transform: Transform::Identity,
1558 })
1559 .unwrap()
1560 .add_unbound_field(UnboundPartitionField {
1561 source_id: 1,
1562 field_id: None,
1563 name: "id_bucket".to_string(),
1564 transform: Transform::Bucket(16),
1565 })
1566 .unwrap()
1567 .build()
1568 .unwrap();
1569
1570 assert!(!partition_spec_1.is_compatible_with(&partition_spec_2));
1571 }
1572
1573 #[test]
1574 fn test_highest_field_id_unpartitioned() {
1575 let spec = PartitionSpec::builder(Schema::builder().with_fields(vec![]).build().unwrap())
1576 .with_spec_id(1)
1577 .build()
1578 .unwrap();
1579
1580 assert!(spec.highest_field_id().is_none());
1581 }
1582
1583 #[test]
1584 fn test_highest_field_id() {
1585 let schema = Schema::builder()
1586 .with_fields(vec![
1587 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1588 NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
1589 ])
1590 .build()
1591 .unwrap();
1592
1593 let spec = PartitionSpec::builder(schema)
1594 .with_spec_id(1)
1595 .add_unbound_field(UnboundPartitionField {
1596 source_id: 1,
1597 field_id: Some(1001),
1598 name: "id".to_string(),
1599 transform: Transform::Identity,
1600 })
1601 .unwrap()
1602 .add_unbound_field(UnboundPartitionField {
1603 source_id: 2,
1604 field_id: Some(1000),
1605 name: "name".to_string(),
1606 transform: Transform::Identity,
1607 })
1608 .unwrap()
1609 .build()
1610 .unwrap();
1611
1612 assert_eq!(Some(1001), spec.highest_field_id());
1613 }
1614
1615 #[test]
1616 fn test_has_sequential_ids() {
1617 let schema = Schema::builder()
1618 .with_fields(vec![
1619 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1620 NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
1621 ])
1622 .build()
1623 .unwrap();
1624
1625 let spec = PartitionSpec::builder(schema)
1626 .with_spec_id(1)
1627 .add_unbound_field(UnboundPartitionField {
1628 source_id: 1,
1629 field_id: Some(1000),
1630 name: "id".to_string(),
1631 transform: Transform::Identity,
1632 })
1633 .unwrap()
1634 .add_unbound_field(UnboundPartitionField {
1635 source_id: 2,
1636 field_id: Some(1001),
1637 name: "name".to_string(),
1638 transform: Transform::Identity,
1639 })
1640 .unwrap()
1641 .build()
1642 .unwrap();
1643
1644 assert_eq!(1000, spec.fields[0].field_id);
1645 assert_eq!(1001, spec.fields[1].field_id);
1646 assert!(spec.has_sequential_ids());
1647 }
1648
1649 #[test]
1650 fn test_sequential_ids_must_start_at_1000() {
1651 let schema = Schema::builder()
1652 .with_fields(vec![
1653 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1654 NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
1655 ])
1656 .build()
1657 .unwrap();
1658
1659 let spec = PartitionSpec::builder(schema)
1660 .with_spec_id(1)
1661 .add_unbound_field(UnboundPartitionField {
1662 source_id: 1,
1663 field_id: Some(999),
1664 name: "id".to_string(),
1665 transform: Transform::Identity,
1666 })
1667 .unwrap()
1668 .add_unbound_field(UnboundPartitionField {
1669 source_id: 2,
1670 field_id: Some(1000),
1671 name: "name".to_string(),
1672 transform: Transform::Identity,
1673 })
1674 .unwrap()
1675 .build()
1676 .unwrap();
1677
1678 assert_eq!(999, spec.fields[0].field_id);
1679 assert_eq!(1000, spec.fields[1].field_id);
1680 assert!(!spec.has_sequential_ids());
1681 }
1682
1683 #[test]
1684 fn test_sequential_ids_must_have_no_gaps() {
1685 let schema = Schema::builder()
1686 .with_fields(vec![
1687 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1688 NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
1689 ])
1690 .build()
1691 .unwrap();
1692
1693 let spec = PartitionSpec::builder(schema)
1694 .with_spec_id(1)
1695 .add_unbound_field(UnboundPartitionField {
1696 source_id: 1,
1697 field_id: Some(1000),
1698 name: "id".to_string(),
1699 transform: Transform::Identity,
1700 })
1701 .unwrap()
1702 .add_unbound_field(UnboundPartitionField {
1703 source_id: 2,
1704 field_id: Some(1002),
1705 name: "name".to_string(),
1706 transform: Transform::Identity,
1707 })
1708 .unwrap()
1709 .build()
1710 .unwrap();
1711
1712 assert_eq!(1000, spec.fields[0].field_id);
1713 assert_eq!(1002, spec.fields[1].field_id);
1714 assert!(!spec.has_sequential_ids());
1715 }
1716
1717 #[test]
1718 fn test_partition_to_path() {
1719 let schema = Schema::builder()
1720 .with_fields(vec![
1721 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
1722 NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
1723 NestedField::required(3, "timestamp", Type::Primitive(PrimitiveType::Timestamp))
1724 .into(),
1725 NestedField::required(4, "empty", Type::Primitive(PrimitiveType::String)).into(),
1726 ])
1727 .build()
1728 .unwrap();
1729
1730 let spec = PartitionSpec::builder(schema.clone())
1731 .add_partition_field("id", "id", Transform::Identity)
1732 .unwrap()
1733 .add_partition_field("name", "name", Transform::Identity)
1734 .unwrap()
1735 .add_partition_field("timestamp", "ts_hour", Transform::Hour)
1736 .unwrap()
1737 .add_partition_field("empty", "empty_void", Transform::Void)
1738 .unwrap()
1739 .build()
1740 .unwrap();
1741
1742 let data = Struct::from_iter([
1743 Some(Literal::int(42)),
1744 Some(Literal::string("alice")),
1745 Some(Literal::int(1000)),
1746 Some(Literal::string("empty")),
1747 ]);
1748
1749 assert_eq!(
1750 spec.partition_to_path(&data, schema.into()),
1751 "id=42/name=alice/ts_hour=1000/empty_void=null"
1752 );
1753 }
1754
1755 #[test]
1756 fn test_partition_to_path_escaped_strings() {
1757 let schema = Schema::builder()
1758 .with_fields(vec![
1759 NestedField::required(1, "\"esc\"#1", Type::Primitive(PrimitiveType::String))
1760 .into(),
1761 NestedField::required(2, "data", Type::Primitive(PrimitiveType::String)).into(),
1762 ])
1763 .build()
1764 .unwrap();
1765
1766 let spec = PartitionSpec::builder(schema.clone())
1767 .add_partition_field("\"esc\"#1", "\"esc\"#1", Transform::Identity)
1768 .unwrap()
1769 .build()
1770 .unwrap();
1771
1772 let data = Struct::from_iter([
1773 Some(Literal::string("a/b/c/d")),
1774 Some(Literal::string("val#1")),
1775 ]);
1776
1777 assert_eq!(
1778 spec.partition_to_path(&data, schema.into()),
1779 "%22esc%22%231=a%2Fb%2Fc%2Fd"
1780 );
1781 }
1782
1783 #[test]
1784 fn test_partition_to_path_escaped_field_name() {
1785 let schema = Schema::builder()
1786 .with_fields(vec![
1787 NestedField::required(1, "\"esc\"#1", Type::Primitive(PrimitiveType::String))
1788 .into(),
1789 NestedField::required(2, "data", Type::Primitive(PrimitiveType::String)).into(),
1790 ])
1791 .build()
1792 .unwrap();
1793
1794 let spec = PartitionSpec::builder(schema.clone())
1795 .add_partition_field("data", "data", Transform::Identity)
1796 .unwrap()
1797 .add_partition_field("data", "data_truc_10", Transform::Truncate(10))
1798 .unwrap()
1799 .build()
1800 .unwrap();
1801
1802 let data = Struct::from_iter([
1803 Some(Literal::string("a/b/c/d")),
1804 Some(Literal::string("a/b/c/d")),
1805 ]);
1806
1807 assert_eq!(
1808 spec.partition_to_path(&data, schema.into()),
1809 "data=a%2Fb%2Fc%2Fd/data_truc_10=a%2Fb%2Fc%2Fd"
1810 );
1811 }
1812}