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