liurenjie1024 commented on code in PR #1781:
URL: https://github.com/apache/iceberg-rust/pull/1781#discussion_r2465115917
##########
crates/iceberg/src/arrow/record_batch_partition_splitter.rs:
##########
@@ -301,4 +342,117 @@ mod tests {
Struct::from_iter(vec![Some(Literal::int(3))]),
]);
}
+
+ #[test]
+ fn test_record_batch_partition_split_with_partition_column() {
+ use arrow_array::StructArray;
+ use arrow_schema::{Field, Schema as ArrowSchema};
+
+ let schema = Arc::new(
+ Schema::builder()
+ .with_fields(vec![
+ NestedField::required(
+ 1,
+ "id",
+ Type::Primitive(crate::spec::PrimitiveType::Int),
+ )
+ .into(),
+ NestedField::required(
+ 2,
+ "name",
+ Type::Primitive(crate::spec::PrimitiveType::String),
+ )
+ .into(),
+ ])
+ .build()
+ .unwrap(),
+ );
+ let partition_spec = Arc::new(
+ PartitionSpecBuilder::new(schema.clone())
+ .with_spec_id(1)
+ .add_unbound_field(UnboundPartitionField {
+ source_id: 1,
+ field_id: None,
+ name: "id_bucket".to_string(),
+ transform: Transform::Identity,
+ })
+ .unwrap()
+ .build()
+ .unwrap(),
+ );
+
+ // Create input schema with _partition column
+ // Note: partition field IDs start from 1000 by default
+ let partition_field = Field::new("id_bucket", DataType::Int32,
false).with_metadata(
+ HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(),
"1000".to_string())]),
+ );
+ let partition_struct_field = Field::new(
+ PROJECTED_PARTITION_VALUE_COLUMN,
+ DataType::Struct(vec![partition_field.clone()].into()),
+ false,
+ );
+
+ let input_schema = Arc::new(ArrowSchema::new(vec![
+ Field::new("id", DataType::Int32, false),
+ Field::new("name", DataType::Utf8, false),
+ partition_struct_field,
+ ]));
+
+ // Create splitter expecting pre-computed partition column
+ let partition_splitter =
RecordBatchPartitionSplitter::new_with_precomputed_values(
+ schema.clone(),
+ partition_spec,
+ )
+ .expect("Failed to create splitter");
+
+ // Create test data with pre-computed partition column
+ let id_array = Int32Array::from(vec![1, 2, 1, 3, 2, 3, 1]);
+ let data_array = StringArray::from(vec!["a", "b", "c", "d", "e", "f",
"g"]);
+
+ // Create partition column (same values as id for Identity transform)
+ let partition_values = Int32Array::from(vec![1, 2, 1, 3, 2, 3, 1]);
+ let partition_struct = StructArray::from(vec![(
+ Arc::new(partition_field),
+ Arc::new(partition_values) as ArrayRef,
+ )]);
+
+ let batch = RecordBatch::try_new(input_schema.clone(), vec![
+ Arc::new(id_array),
+ Arc::new(data_array),
+ Arc::new(partition_struct),
+ ])
+ .expect("Failed to create RecordBatch");
+
+ // Split using the pre-computed partition column
+ let mut partitioned_batches = partition_splitter
+ .split(&batch)
+ .expect("Failed to split RecordBatch");
+
+ partitioned_batches.sort_by_key(|(partition_key, _)| {
+ if let PrimitiveLiteral::Int(i) = partition_key.data().fields()[0]
+ .as_ref()
+ .unwrap()
+ .as_primitive_literal()
+ .unwrap()
+ {
+ i
+ } else {
+ panic!("The partition value is not a int");
+ }
+ });
+
+ assert_eq!(partitioned_batches.len(), 3);
+
+ // Verify partition values
+ let partition_values = partitioned_batches
+ .iter()
+ .map(|(partition_key, _)| partition_key.data().clone())
+ .collect::<Vec<_>>();
+
+ assert_eq!(partition_values, vec![
+ Struct::from_iter(vec![Some(Literal::int(1))]),
+ Struct::from_iter(vec![Some(Literal::int(2))]),
+ Struct::from_iter(vec![Some(Literal::int(3))]),
+ ]);
Review Comment:
This test didn't verify splitted record batch.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
