diff --git a/crates/paimon/src/spec/binary_row.rs b/crates/paimon/src/spec/binary_row.rs index 2c2e2f1ae..17c4820ab 100644 --- a/crates/paimon/src/spec/binary_row.rs +++ b/crates/paimon/src/spec/binary_row.rs @@ -1332,6 +1332,7 @@ enum TypedColumn<'a> { Utf8View(&'a arrow_array::StringViewArray), LargeUtf8(&'a arrow_array::LargeStringArray), Date32(&'a arrow_array::Date32Array), + Time32Ms(&'a arrow_array::Time32MillisecondArray), Decimal128(&'a arrow_array::Decimal128Array, u32, u32), // (array, precision, scale) Binary(&'a arrow_array::BinaryArray), Variant(&'a arrow_array::StructArray), @@ -1408,6 +1409,18 @@ fn downcast_columns<'a>( .downcast_ref() .ok_or_else(|| type_mismatch_err("Date", col_idx))?, ), + // Java writes TIME with the same `writeInt` as INTEGER and DATE + // (`BinaryWriter`), so in the binary format a TIME *is* an int + // millis-of-day and no precision can add detail. Hence one arm + // rather than the `match precision` the `Timestamp` arm below + // needs; `paimon_type_to_arrow` maps every precision to + // `Time32(Millisecond)` and the batch is validated against the + // schema built from it, so that is what arrives here. + DataType::Time(_) => TypedColumn::Time32Ms( + col.as_any() + .downcast_ref() + .ok_or_else(|| type_mismatch_err("Time", col_idx))?, + ), DataType::Decimal(d) => TypedColumn::Decimal128( col.as_any() .downcast_ref() @@ -1587,6 +1600,16 @@ fn write_typed_value( builder.write_int(pos, arr.value(row_idx)); } } + // Java keeps TIME in the fixed-length part and writes it with `writeInt` + // (`BinaryWriter#write` -> `BinaryRowWriter#writeInt`), so the four bytes and + // the zeroed upper half of the slot have to match for `hash_code` parity. + TypedColumn::Time32Ms(arr) => { + if arr.is_null(row_idx) { + builder.set_null_at(pos); + } else { + builder.write_int(pos, arr.value(row_idx)); + } + } TypedColumn::Decimal128(arr, precision, _scale) => { if arr.is_null(row_idx) { builder.set_null_at(pos); @@ -2299,19 +2322,25 @@ mod tests { #[test] fn test_batch_vs_per_row_equivalence() { - use arrow_array::{Int32Array, StringArray}; - use arrow_schema::{DataType as ArrowDT, Field, Schema}; + use arrow_array::{Int32Array, StringArray, Time32MillisecondArray}; + use arrow_schema::{DataType as ArrowDT, Field, Schema, TimeUnit}; use std::sync::Arc; let schema = Arc::new(Schema::new(vec![ Field::new("id", ArrowDT::Int32, true), Field::new("name", ArrowDT::Utf8, true), + Field::new("tm", ArrowDT::Time32(TimeUnit::Millisecond), true), ])); let batch = RecordBatch::try_new( schema, vec![ Arc::new(Int32Array::from(vec![Some(1), None, Some(3)])), Arc::new(StringArray::from(vec![Some("hello"), Some("world"), None])), + Arc::new(Time32MillisecondArray::from(vec![ + Some(45_296_123), + None, + Some(0), + ])), ], ) .unwrap(); @@ -2323,8 +2352,13 @@ mod tests { "name".into(), DataType::VarChar(crate::spec::VarCharType::string_type()), ), + crate::spec::DataField::new( + 2, + "tm".into(), + DataType::Time(crate::spec::TimeType::new(3).unwrap()), + ), ]; - let indices = vec![0, 1]; + let indices = vec![0, 1, 2]; // Batch results let batch_bytes = batch_to_serialized_bytes(&batch, &indices, &fields).unwrap(); diff --git a/crates/paimon/src/table/table_write.rs b/crates/paimon/src/table/table_write.rs index 1a0c39938..15862f07e 100644 --- a/crates/paimon/src/table/table_write.rs +++ b/crates/paimon/src/table/table_write.rs @@ -2600,6 +2600,135 @@ pub(in crate::table) mod tests { assert_eq!(total_rows, 4); } + /// TIME as an append table's `bucket-key`. Java allows it: `validateBucket` + /// rejects only ARRAY, MULTISET, MAP and ROW there. Routing has to agree with + /// `BinaryRow::hash_code`, so the bucket is pinned against the row the per-row + /// encoder produces rather than against a hard-coded number. + #[tokio::test] + async fn test_time_bucket_key_routes_by_binary_row_hash() { + let file_io = test_file_io(); + let table_path = "memory:/test_time_bucket_key"; + setup_dirs(&file_io, table_path).await; + + let time_type = DataType::Time(TimeType::new(3).unwrap()); + let schema = Schema::builder() + .column("tm", time_type.clone()) + .column("value", DataType::Int(IntType::new())) + .option("bucket", "4") + .option("bucket-key", "tm") + .build() + .unwrap(); + let table = Table::new( + file_io.clone(), + Identifier::new("default", "test_time_bucket_key_table"), + table_path.to_string(), + TableSchema::new(0, &schema), + None, + ); + + // 12:34:56.123 and midnight. + let times = [45_296_123_i32, 0]; + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("tm", ArrowDataType::Time32(TimeUnit::Millisecond), true), + ArrowField::new("value", ArrowDataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + arrow_schema, + vec![ + Arc::new(Time32MillisecondArray::from(times.map(Some).to_vec())), + Arc::new(Int32Array::from(vec![Some(10), Some(20)])), + ], + ) + .unwrap(); + + let fields = table.schema().fields().to_vec(); + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + let output = table_write + .bucket_assigner + .assign_batch(&batch, &fields) + .await + .unwrap(); + + let expected: Vec = times + .iter() + .map(|&millis| { + let row = BinaryRow::from_datums(&[(Some(&Datum::Time(millis)), &time_type)]); + // Mirrors `default_bucket`: `(hash % n).abs()`, which is Java's + // `Math.abs(hashcode % numBuckets)` and is *not* a euclidean + // remainder for negative hashes. + (row.hash_code() % 4).wrapping_abs() + }) + .collect(); + assert_eq!(output.buckets, expected); + + // And the write itself must land, not just the routing. + table_write.write_arrow_batch(&batch).await.unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + let rows: i64 = messages + .iter() + .flat_map(|m| m.new_files.iter()) + .map(|f| f.row_count) + .sum(); + assert_eq!(rows, 2); + } + + /// TIME as a partition key. `partition_utils` already renders TIME partition + /// values, but that code was unreachable from the write path while the batch + /// encoder rejected the column, so this is the first test that exercises the + /// two together — hence the assertion on the rendered path, not just the row. + #[tokio::test] + async fn test_time_partition_key_writes_formatted_partition() { + let file_io = test_file_io(); + let table_path = "memory:/test_time_partition_key"; + setup_dirs(&file_io, table_path).await; + + let schema = Schema::builder() + .column("tm", DataType::Time(TimeType::new(3).unwrap())) + .column("value", DataType::Int(IntType::new())) + .partition_keys(["tm"]) + .build() + .unwrap(); + let table = Table::new( + file_io.clone(), + Identifier::new("default", "test_time_partition_table"), + table_path.to_string(), + TableSchema::new(0, &schema), + None, + ); + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("tm", ArrowDataType::Time32(TimeUnit::Millisecond), true), + ArrowField::new("value", ArrowDataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + arrow_schema, + vec![ + Arc::new(Time32MillisecondArray::from(vec![Some(45_296_123)])), + Arc::new(Int32Array::from(vec![Some(10)])), + ], + ) + .unwrap(); + + table_write.write_arrow_batch(&batch).await.unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + assert_eq!(messages.len(), 1); + let partition = BinaryRow::from_serialized_bytes(&messages[0].partition).unwrap(); + assert_eq!(partition.get_int(0).unwrap(), 45_296_123); + + let computer = PartitionComputer::new( + table.schema().partition_keys(), + table.schema().fields(), + "__DEFAULT_PARTITION__", + false, + ) + .unwrap(); + assert_eq!( + computer.generate_partition_path(&partition).unwrap(), + "tm=12%3A34%3A56.123/" + ); + } + fn test_bucketed_schema() -> TableSchema { let schema = Schema::builder() .column("id", DataType::Int(IntType::new()))