You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@paimon.apache.org by lz...@apache.org on 2026/09/25 02:37:59 UTC
(paimon-rust) branch main updated: Support Avro writes, defaults, and custom partition locations for native Format Tables (#950)
This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-rust.git
The following commit(s) were added to refs/heads/main by this push:
new b2f1eb3c Support Avro writes, defaults, and custom partition locations for native Format Tables (#950)
b2f1eb3c is described below
commit b2f1eb3cb03843520f614adf4c4663d26fbcbba3
Author: Jingsong Lee <ji...@gmail.com>
AuthorDate: Fri Sep 25 10:37:54 2026 +0800
Support Avro writes, defaults, and custom partition locations for native Format Tables (#950)
---
.../datafusion/src/partition_count_pushdown.rs | 7 +-
.../datafusion/tests/rest_format_table_scan.rs | 131 +-
crates/paimon/src/arrow/format/avro.rs | 81 +-
crates/paimon/src/arrow/format/avro_write.rs | 1034 ++++++++++++++
crates/paimon/src/arrow/format/mod.rs | 18 +
crates/paimon/src/spec/core_options.rs | 29 +
.../paimon/src/table/format_partition_location.rs | 544 ++++++++
crates/paimon/src/table/format_table_defaults.rs | 1404 ++++++++++++++++++++
crates/paimon/src/table/format_table_scan.rs | 51 +-
.../paimon/src/table/format_table_write_tests.rs | 32 +-
crates/paimon/src/table/format_table_writer.rs | 125 +-
crates/paimon/src/table/mod.rs | 2 +
12 files changed, 3394 insertions(+), 64 deletions(-)
diff --git a/crates/integrations/datafusion/src/partition_count_pushdown.rs b/crates/integrations/datafusion/src/partition_count_pushdown.rs
index 0956c29a..58aad534 100644
--- a/crates/integrations/datafusion/src/partition_count_pushdown.rs
+++ b/crates/integrations/datafusion/src/partition_count_pushdown.rs
@@ -161,11 +161,10 @@ fn rewrite_aggregate(aggregate: &Aggregate) -> DFResult<Option<LogicalPlan>> {
return Ok(None);
}
match &analysis.pushed_predicate {
- Some(predicate) => {
- if !table.new_read_builder().is_exact_filter_pushdown(predicate) {
- return Ok(None);
- }
+ Some(predicate) if !table.new_read_builder().is_exact_filter_pushdown(predicate) => {
+ return Ok(None);
}
+ Some(_) => {}
None if !scan.filters.is_empty() => return Ok(None),
None => {}
}
diff --git a/crates/integrations/datafusion/tests/rest_format_table_scan.rs b/crates/integrations/datafusion/tests/rest_format_table_scan.rs
index a446147f..7a2c7b4f 100644
--- a/crates/integrations/datafusion/tests/rest_format_table_scan.rs
+++ b/crates/integrations/datafusion/tests/rest_format_table_scan.rs
@@ -396,7 +396,7 @@ async fn test_managed_scan_filter_on_a_partition_column_reads_the_registered_val
#[cfg(not(windows))]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
-async fn test_managed_scan_refuses_a_partition_at_a_custom_location() {
+async fn test_managed_scan_reads_a_partition_at_a_custom_location() {
let temp_dir = tempfile::tempdir().unwrap();
let external_dir = tempfile::tempdir().unwrap();
let table = ManagedTable::new(&temp_dir, &[("dt", varchar())]).await;
@@ -405,6 +405,7 @@ async fn test_managed_scan_refuses_a_partition_at_a_custom_location() {
// Another engine registered dt=b somewhere else; the table directory still has a stale copy.
write_ids(&temp_dir.path().join("dt=b"), &[2]);
write_ids(external_dir.path(), &[3]);
+ write_ids(&external_dir.path().join("_temporary"), &[999]);
table.server.set_table_partition_options(
DATABASE,
TABLE,
@@ -415,13 +416,131 @@ async fn test_managed_scan_refuses_a_partition_at_a_custom_location() {
)]),
);
- // Reading the default directory would return the stale row, so a scan that reaches the
- // partition fails instead. One that does not reach it is unaffected.
- for predicate in ["dt = 'b'", "TRUE"] {
- let error = table.error(predicate).await;
- assert!(error.contains("custom location"), "{predicate}: {error}");
+ // A scan must use the catalog's location and ignore the stale default directory.
+ assert_eq!(table.ids("dt = 'b'").await, vec![3]);
+ assert_eq!(table.ids("TRUE").await, vec![1, 3]);
+ assert_eq!(table.ids("dt = 'a'").await, vec![1]);
+}
+
+#[cfg(not(windows))]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn test_managed_scan_rejects_overlapping_external_partition_locations() {
+ let temp_dir = tempfile::tempdir().unwrap();
+ let external_dir = tempfile::tempdir().unwrap();
+ let table = ManagedTable::new(&temp_dir, &[("dt", varchar())]).await;
+ table.register(&[&[("dt", "a")], &[("dt", "b")]]).await;
+ let shared = format!("file://{}", external_dir.path().display());
+ for dt in ["a", "b"] {
+ table.server.set_table_partition_options(
+ DATABASE,
+ TABLE,
+ &spec(&[("dt", dt)]),
+ HashMap::from([("path".to_string(), shared.clone())]),
+ );
+ }
+ let error = table.error("TRUE").await;
+ assert!(error.contains("overlapping locations"), "{error}");
+}
+
+#[cfg(not(windows))]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn test_managed_scan_rejects_external_location_inside_table_directory() {
+ let temp_dir = tempfile::tempdir().unwrap();
+ let table = ManagedTable::new(&temp_dir, &[("dt", varchar())]).await;
+ table.register(&[&[("dt", "a")]]).await;
+ table.server.set_table_partition_options(
+ DATABASE,
+ TABLE,
+ &spec(&[("dt", "a")]),
+ HashMap::from([(
+ "path".to_string(),
+ format!("file://{}/other-partition", temp_dir.path().display()),
+ )]),
+ );
+ let error = table.error("dt = 'a'").await;
+ assert!(
+ error.contains("overlaps the Format Table directory"),
+ "{error}"
+ );
+}
+
+#[cfg(not(windows))]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn test_managed_scan_rejects_encoded_traversal_without_reading_external_data() {
+ let temp_dir = tempfile::tempdir().unwrap();
+ let table = ManagedTable::new(&temp_dir, &[("dt", varchar())]).await;
+ table.register(&[&[("dt", "a")]]).await;
+ write_ids(&temp_dir.path().join("dt=a"), &[1]);
+ for invalid in [
+ "file:/external/%2e%2e/secret",
+ "file:/external/%252e%252e/secret",
+ "file:/external/a?query=1",
+ "file:/external/%ZZ",
+ ] {
+ table.server.set_table_partition_options(
+ DATABASE,
+ TABLE,
+ &spec(&[("dt", "a")]),
+ HashMap::from([("path".to_string(), invalid.to_string())]),
+ );
+ let error = table.error("dt = 'a'").await;
+ assert!(
+ error.contains("invalid custom location"),
+ "{invalid}: {error}"
+ );
}
+}
+
+#[cfg(not(windows))]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn test_managed_scan_prunes_unrelated_invalid_custom_location() {
+ let temp_dir = tempfile::tempdir().unwrap();
+ let table = ManagedTable::new(&temp_dir, &[("dt", varchar())]).await;
+ table.register(&[&[("dt", "a")], &[("dt", "b")]]).await;
+ write_ids(&temp_dir.path().join("dt=a"), &[1]);
+ table.server.set_table_partition_options(
+ DATABASE,
+ TABLE,
+ &spec(&[("dt", "b")]),
+ HashMap::from([("path".to_string(), "file:/external/../secret".to_string())]),
+ );
+
+ // Only the reached partition is opened. A full-table scan still fails on
+ // the bad catalog metadata before it can return a partial answer.
assert_eq!(table.ids("dt = 'a'").await, vec![1]);
+ assert!(table
+ .error("TRUE")
+ .await
+ .contains("invalid custom location"));
+}
+
+#[cfg(not(windows))]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn test_managed_scan_uses_catalog_values_for_external_multi_key_partition() {
+ let temp_dir = tempfile::tempdir().unwrap();
+ let external_dir = tempfile::tempdir().unwrap();
+ let table = ManagedTable::new(
+ &temp_dir,
+ &[
+ ("dt", varchar()),
+ ("active", DataType::Boolean(BooleanType::new())),
+ ],
+ )
+ .await;
+ table.register(&[&[("dt", "a"), ("active", "true")]]).await;
+ write_ids(external_dir.path(), &[10, 11]);
+ table.server.set_table_partition_options(
+ DATABASE,
+ TABLE,
+ &spec(&[("dt", "a"), ("active", "true")]),
+ HashMap::from([(
+ "path".to_string(),
+ format!("file://{}", external_dir.path().display()),
+ )]),
+ );
+
+ assert_eq!(table.ids("active").await, vec![10, 11]);
+ assert_eq!(table.ids("NOT active").await, Vec::<i64>::new());
}
#[cfg(not(windows))]
diff --git a/crates/paimon/src/arrow/format/avro.rs b/crates/paimon/src/arrow/format/avro.rs
index 6c17cbca..37f01c51 100644
--- a/crates/paimon/src/arrow/format/avro.rs
+++ b/crates/paimon/src/arrow/format/avro.rs
@@ -24,10 +24,11 @@ use crate::Error;
use apache_avro::types::Value;
use apache_avro::Reader;
use arrow_array::{
- BinaryArray, BooleanArray, Date32Array, Decimal128Array, Float32Array, Float64Array,
- Int16Array, Int32Array, Int64Array, Int8Array, LargeBinaryArray, ListArray, MapArray,
- RecordBatch, StringArray, StructArray, Time32MillisecondArray, TimestampMicrosecondArray,
- TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray,
+ BinaryArray, BooleanArray, Date32Array, Decimal128Array, FixedSizeListArray, Float32Array,
+ Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, LargeBinaryArray, ListArray,
+ MapArray, RecordBatch, StringArray, StructArray, Time32MillisecondArray,
+ TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray,
+ TimestampSecondArray,
};
use arrow_buffer::{BooleanBuffer, NullBuffer, OffsetBuffer, ScalarBuffer};
use arrow_schema::SchemaRef;
@@ -398,6 +399,13 @@ fn build_column(
DataType::Array(arr_type) => {
build_array_column(records, name, arr_type.element_type(), num_rows)?
}
+ DataType::Vector(vector_type) => build_vector_column(
+ records,
+ name,
+ vector_type.element_type(),
+ vector_type.length() as usize,
+ num_rows,
+ )?,
DataType::Map(map_type) => build_map_column(records, name, map_type, num_rows)?,
// Java encodes MULTISET<T> as a map from the element to an INT count,
// sharing the MAP path (`AvroSchemaConverter#extractValueTypeToAvroMap`
@@ -504,6 +512,71 @@ fn build_array_column(
Ok(Arc::new(list_arr))
}
+fn build_vector_column(
+ records: &[Value],
+ name: &str,
+ element_type: &DataType,
+ length: usize,
+ num_rows: usize,
+) -> crate::Result<Arc<dyn arrow_array::Array>> {
+ let idx = field_index(records, name);
+ let mut elements = Vec::with_capacity(num_rows.saturating_mul(length));
+ let mut valid = Vec::with_capacity(num_rows);
+ for record in records.iter().take(num_rows) {
+ match get_field_at(record, idx) {
+ Some(Value::Array(values)) if values.len() == length => {
+ valid.push(true);
+ elements.extend(values.iter().cloned());
+ }
+ Some(Value::Array(values)) => {
+ return Err(Error::DataInvalid {
+ message: format!(
+ "Avro vector '{name}' has {} elements, expected {length}",
+ values.len()
+ ),
+ source: None,
+ });
+ }
+ Some(other) => {
+ return Err(Error::DataInvalid {
+ message: format!("Avro vector '{name}' is not an array: {other:?}"),
+ source: None,
+ });
+ }
+ None => {
+ valid.push(false);
+ elements.extend(std::iter::repeat_n(Value::Null, length));
+ }
+ }
+ }
+ let element_records = elements
+ .into_iter()
+ .map(|element| Value::Record(vec![("element".to_string(), element)]))
+ .collect::<Vec<_>>();
+ let values = build_column(
+ &element_records,
+ "element",
+ element_type,
+ element_records.len(),
+ )?;
+ let field = Arc::new(arrow_schema::Field::new(
+ "element",
+ crate::arrow::paimon_type_to_arrow(element_type)?,
+ element_type.is_nullable(),
+ ));
+ let array = FixedSizeListArray::try_new(
+ field,
+ length as i32,
+ values,
+ Some(NullBuffer::new(BooleanBuffer::from(valid))),
+ )
+ .map_err(|source| Error::DataInvalid {
+ message: format!("Cannot build Avro vector '{name}'"),
+ source: Some(Box::new(source)),
+ })?;
+ Ok(Arc::new(array))
+}
+
fn build_map_column(
records: &[Value],
name: &str,
diff --git a/crates/paimon/src/arrow/format/avro_write.rs b/crates/paimon/src/arrow/format/avro_write.rs
new file mode 100644
index 00000000..d52ef600
--- /dev/null
+++ b/crates/paimon/src/arrow/format/avro_write.rs
@@ -0,0 +1,1034 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Avro OCF writer compatible with Java's `AvroSchemaConverter`. Each Arrow
+//! batch becomes an OCF block. The schema and sync marker are written once;
+//! later blocks stream to storage without retaining batches or a whole file.
+
+use std::collections::HashMap;
+
+use apache_avro::types::Value;
+use apache_avro::{Codec, Schema, Writer};
+use arrow_array::{
+ Array, BinaryArray, BooleanArray, Date32Array, Decimal128Array, FixedSizeListArray,
+ Float32Array, Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, LargeBinaryArray,
+ ListArray, MapArray, RecordBatch, StringArray, StructArray, Time32MillisecondArray,
+ TimestampMicrosecondArray, TimestampMillisecondArray, TimestampSecondArray,
+};
+use arrow_schema::SchemaRef;
+use async_trait::async_trait;
+use bytes::Bytes;
+use serde_json::{json, Value as JsonValue};
+
+use super::{FormatFileWriter, FormatWriteResult};
+use crate::io::{FileWrite, OutputFile};
+use crate::spec::{BlobDescriptor, BlobViewStruct, CoreOptions, DataField, DataType};
+use crate::{Error, Result};
+
+pub(crate) struct AvroFormatWriter {
+ writer: Box<dyn FileWrite>,
+ schema: Schema,
+ arrow_schema: SchemaRef,
+ fields: Vec<DataField>,
+ codec: Codec,
+ marker: [u8; 16],
+ block_size: usize,
+ bytes_written: usize,
+}
+
+impl AvroFormatWriter {
+ pub(crate) async fn new(
+ output: &OutputFile,
+ arrow_schema: SchemaRef,
+ fields: Vec<DataField>,
+ compression: &str,
+ zstd_level: i32,
+ options: Option<&HashMap<String, String>>,
+ ) -> Result<Self> {
+ let compression = options
+ .and_then(|options| options.get("avro.codec"))
+ .map(String::as_str)
+ .unwrap_or(compression);
+ let codec = match compression.to_ascii_lowercase().as_str() {
+ "none" | "null" | "uncompressed" => Codec::Null,
+ "snappy" => Codec::Snappy,
+ "zstd" | "zstandard" => Codec::Zstandard(apache_avro::ZstandardSettings::new(
+ zstd_level.clamp(0, 22) as u8,
+ )),
+ "deflate" => Codec::Deflate(Default::default()),
+ other => {
+ return Err(Error::Unsupported {
+ message: format!("Unsupported Avro compression codec '{other}'"),
+ });
+ }
+ };
+ let block_size = match options {
+ Some(options) => CoreOptions::new(options)
+ .file_block_size()?
+ .unwrap_or(64 * 1024),
+ None => 64 * 1024,
+ };
+ if !(32..=1024 * 1024 * 1024).contains(&block_size) {
+ return Err(Error::ConfigInvalid {
+ message: format!(
+ "file.block-size for avro must be between 32 bytes and 1 gb, but was {block_size} bytes"
+ ),
+ });
+ }
+ let schema_json = row_schema(&fields)?;
+ let schema = Schema::parse_str(&schema_json.to_string()).map_err(avro_error)?;
+ let header = Writer::with_codec(&schema, Vec::new(), codec)
+ .into_inner()
+ .map_err(avro_error)?;
+ let marker: [u8; 16] = header[header.len() - 16..]
+ .try_into()
+ .expect("Avro OCF header ends with a 16-byte sync marker");
+ let mut writer = output.writer().await?;
+ writer.write(Bytes::from(header.clone())).await?;
+ Ok(Self {
+ writer,
+ schema,
+ arrow_schema,
+ fields,
+ codec,
+ marker,
+ block_size: block_size as usize,
+ bytes_written: header.len(),
+ })
+ }
+}
+
+#[async_trait]
+impl FormatFileWriter for AvroFormatWriter {
+ async fn write(&mut self, batch: &RecordBatch) -> Result<()> {
+ if batch.schema() != self.arrow_schema {
+ return Err(Error::DataInvalid {
+ message: "Avro writer input schema differs from its file schema".into(),
+ source: None,
+ });
+ }
+ if batch.num_rows() == 0 {
+ return Ok(());
+ }
+ let mut block = Writer::builder()
+ .schema(&self.schema)
+ .writer(Vec::new())
+ .codec(self.codec)
+ .marker(self.marker)
+ .has_header(true)
+ .block_size(self.block_size)
+ .build();
+ for row in 0..batch.num_rows() {
+ let record = self
+ .fields
+ .iter()
+ .enumerate()
+ .map(|(column, field)| {
+ Ok((
+ field.name().to_string(),
+ value_at(batch.column(column).as_ref(), row, field.data_type())?,
+ ))
+ })
+ .collect::<Result<Vec<_>>>()?;
+ block.append(Value::Record(record)).map_err(avro_error)?;
+ }
+ let bytes = block.into_inner().map_err(avro_error)?;
+ self.bytes_written += bytes.len();
+ self.writer.write(Bytes::from(bytes)).await?;
+ Ok(())
+ }
+
+ fn num_bytes(&self) -> usize {
+ self.bytes_written
+ }
+
+ fn in_progress_size(&self) -> usize {
+ 0
+ }
+
+ async fn flush(&mut self) -> Result<()> {
+ Ok(())
+ }
+
+ async fn close(mut self: Box<Self>) -> Result<FormatWriteResult> {
+ self.writer.close().await?;
+ Ok(FormatWriteResult::new(self.bytes_written as u64))
+ }
+}
+
+fn avro_error(source: apache_avro::Error) -> Error {
+ Error::DataInvalid {
+ message: format!("Cannot write Avro data: {source}"),
+ source: Some(Box::new(source)),
+ }
+}
+
+fn row_schema(fields: &[DataField]) -> Result<JsonValue> {
+ let fields = fields
+ .iter()
+ .map(|field| {
+ let typ = avro_type(field.data_type(), &format!("record_{}", field.name()))?;
+ let mut entry = json!({"name": field.name(), "type": typ});
+ if field.data_type().is_nullable() {
+ entry["default"] = JsonValue::Null;
+ }
+ Ok(entry)
+ })
+ .collect::<Result<Vec<_>>>()?;
+ Ok(json!({
+ "type": "record",
+ "name": "record",
+ "namespace": "org.apache.paimon.avro.generated",
+ "fields": fields,
+ }))
+}
+
+fn avro_type(typ: &DataType, name: &str) -> Result<JsonValue> {
+ let nonnull = match typ {
+ DataType::Boolean(_) => json!("boolean"),
+ DataType::TinyInt(_) | DataType::SmallInt(_) | DataType::Int(_) => json!("int"),
+ DataType::BigInt(_) => json!("long"),
+ DataType::Float(_) => json!("float"),
+ DataType::Double(_) => json!("double"),
+ DataType::Char(_) | DataType::VarChar(_) => json!("string"),
+ DataType::Binary(_) | DataType::VarBinary(_) | DataType::Blob(_) => json!("bytes"),
+ DataType::Date(_) => json!({"type": "int", "logicalType": "date"}),
+ DataType::Time(time) if time.precision() <= 3 => {
+ json!({"type": "int", "logicalType": "time-millis"})
+ }
+ DataType::Timestamp(timestamp) if timestamp.precision() <= 6 => {
+ let logical = if timestamp.precision() <= 3 {
+ "timestamp-millis"
+ } else {
+ "timestamp-micros"
+ };
+ json!({"type": "long", "logicalType": logical})
+ }
+ DataType::LocalZonedTimestamp(timestamp) if timestamp.precision() <= 6 => {
+ let logical = if timestamp.precision() <= 3 {
+ "local-timestamp-millis"
+ } else {
+ "local-timestamp-micros"
+ };
+ json!({"type": "long", "logicalType": logical})
+ }
+ DataType::Decimal(decimal) => json!({
+ "type": "bytes", "logicalType": "decimal",
+ "precision": decimal.precision(), "scale": decimal.scale(),
+ }),
+ DataType::Array(array) => {
+ json!({"type": "array", "items": avro_type(array.element_type(), name)?})
+ }
+ DataType::Vector(vector) => {
+ json!({"type": "array", "items": avro_type(vector.element_type(), name)?})
+ }
+ DataType::Row(row) => {
+ let fields = row
+ .fields()
+ .iter()
+ .map(|field| {
+ let child_name = format!("{name}_{}", field.name());
+ let mut entry = json!({
+ "name": field.name(),
+ "type": avro_type(field.data_type(), &child_name)?,
+ });
+ if field.data_type().is_nullable() {
+ entry["default"] = JsonValue::Null;
+ }
+ Ok(entry)
+ })
+ .collect::<Result<Vec<_>>>()?;
+ json!({"type": "record", "name": name, "fields": fields})
+ }
+ DataType::Map(map) => avro_map_type(map.key_type(), map.value_type(), name)?,
+ DataType::Multiset(multiset) => avro_map_type(
+ multiset.element_type(),
+ &DataType::Int(crate::spec::IntType::new()),
+ name,
+ )?,
+ other => {
+ return Err(Error::Unsupported {
+ message: format!("Avro writer does not support data type {other:?}"),
+ });
+ }
+ };
+ if typ.is_nullable() {
+ Ok(json!(["null", nonnull]))
+ } else {
+ Ok(nonnull)
+ }
+}
+
+fn avro_map_type(key: &DataType, value: &DataType, name: &str) -> Result<JsonValue> {
+ if matches!(key, DataType::Char(_) | DataType::VarChar(_)) {
+ Ok(json!({"type": "map", "values": avro_type(value, name)?}))
+ } else {
+ let key_name = format!("{name}_key");
+ let value_name = format!("{name}_value");
+ Ok(json!({
+ "type": "array",
+ "logicalType": "map",
+ "items": {
+ "type": "record", "name": name,
+ "fields": [
+ {"name": "key", "type": avro_type(key, &key_name)?},
+ {"name": "value", "type": avro_type(value, &value_name)?},
+ ],
+ },
+ }))
+ }
+}
+
+fn value_at(array: &dyn Array, row: usize, typ: &DataType) -> Result<Value> {
+ if array.is_null(row) {
+ if !typ.is_nullable() {
+ return Err(Error::DataInvalid {
+ message: format!("Cannot write null to non-null Avro type {typ:?}"),
+ source: None,
+ });
+ }
+ return Ok(Value::Union(0, Box::new(Value::Null)));
+ }
+ let value = nonnull_value_at(array, row, typ)?;
+ if typ.is_nullable() {
+ Ok(Value::Union(1, Box::new(value)))
+ } else {
+ Ok(value)
+ }
+}
+
+macro_rules! primitive {
+ ($array:expr, $row:expr, $ty:ty, $variant:ident) => {
+ Value::$variant(
+ $array
+ .as_any()
+ .downcast_ref::<$ty>()
+ .ok_or_else(|| Error::DataInvalid {
+ message: format!("Avro input has wrong Arrow type for {}", stringify!($ty)),
+ source: None,
+ })?
+ .value($row)
+ .into(),
+ )
+ };
+}
+
+fn nonnull_value_at(array: &dyn Array, row: usize, typ: &DataType) -> Result<Value> {
+ let value = match typ {
+ DataType::Boolean(_) => primitive!(array, row, BooleanArray, Boolean),
+ DataType::TinyInt(_) => primitive!(array, row, Int8Array, Int),
+ DataType::SmallInt(_) => primitive!(array, row, Int16Array, Int),
+ DataType::Int(_) => primitive!(array, row, Int32Array, Int),
+ DataType::BigInt(_) => primitive!(array, row, Int64Array, Long),
+ DataType::Float(_) => primitive!(array, row, Float32Array, Float),
+ DataType::Double(_) => primitive!(array, row, Float64Array, Double),
+ DataType::Char(_) | DataType::VarChar(_) => {
+ Value::String(downcast::<StringArray>(array)?.value(row).to_string())
+ }
+ DataType::Binary(_) | DataType::VarBinary(_) => {
+ Value::Bytes(downcast::<BinaryArray>(array)?.value(row).to_vec())
+ }
+ DataType::Blob(_) => {
+ let bytes = downcast::<LargeBinaryArray>(array)?.value(row);
+ let serialized = if BlobViewStruct::is_blob_view_struct(bytes) {
+ BlobViewStruct::deserialize(bytes)?.serialize()?
+ } else {
+ BlobDescriptor::deserialize(bytes)
+ .map_err(|source| Error::DataInvalid {
+ message: "Avro BLOB requires a BlobDescriptor or BlobViewStruct".into(),
+ source: Some(Box::new(source)),
+ })?
+ .serialize()
+ };
+ Value::Bytes(serialized)
+ }
+ DataType::Date(_) => Value::Date(downcast::<Date32Array>(array)?.value(row)),
+ DataType::Time(_) => {
+ Value::TimeMillis(downcast::<Time32MillisecondArray>(array)?.value(row))
+ }
+ DataType::Timestamp(timestamp) if timestamp.precision() == 0 => Value::TimestampMillis(
+ seconds_to_millis(downcast::<TimestampSecondArray>(array)?.value(row))?,
+ ),
+ DataType::Timestamp(timestamp) if timestamp.precision() <= 3 => {
+ Value::TimestampMillis(downcast::<TimestampMillisecondArray>(array)?.value(row))
+ }
+ DataType::Timestamp(_) => {
+ Value::TimestampMicros(downcast::<TimestampMicrosecondArray>(array)?.value(row))
+ }
+ DataType::LocalZonedTimestamp(timestamp) if timestamp.precision() == 0 => {
+ Value::LocalTimestampMillis(seconds_to_millis(
+ downcast::<TimestampSecondArray>(array)?.value(row),
+ )?)
+ }
+ DataType::LocalZonedTimestamp(timestamp) if timestamp.precision() <= 3 => {
+ Value::LocalTimestampMillis(downcast::<TimestampMillisecondArray>(array)?.value(row))
+ }
+ DataType::LocalZonedTimestamp(_) => {
+ Value::LocalTimestampMicros(downcast::<TimestampMicrosecondArray>(array)?.value(row))
+ }
+ DataType::Decimal(_) => {
+ let raw = downcast::<Decimal128Array>(array)?.value(row);
+ Value::Decimal(apache_avro::Decimal::from(decimal_bytes(raw)))
+ }
+ DataType::Array(array_type) => {
+ let list = downcast::<ListArray>(array)?;
+ let values = list.value(row);
+ Value::Array(
+ (0..values.len())
+ .map(|index| value_at(values.as_ref(), index, array_type.element_type()))
+ .collect::<Result<Vec<_>>>()?,
+ )
+ }
+ DataType::Vector(vector_type) => {
+ let list = downcast::<FixedSizeListArray>(array)?;
+ let values = list.value(row);
+ if values.len() != vector_type.length() as usize {
+ return Err(Error::DataInvalid {
+ message: format!(
+ "Avro vector has {} elements, expected {}",
+ values.len(),
+ vector_type.length()
+ ),
+ source: None,
+ });
+ }
+ Value::Array(
+ (0..values.len())
+ .map(|index| value_at(values.as_ref(), index, vector_type.element_type()))
+ .collect::<Result<Vec<_>>>()?,
+ )
+ }
+ DataType::Row(row_type) => {
+ let structure = downcast::<StructArray>(array)?;
+ Value::Record(
+ row_type
+ .fields()
+ .iter()
+ .enumerate()
+ .map(|(index, field)| {
+ Ok((
+ field.name().to_string(),
+ value_at(structure.column(index).as_ref(), row, field.data_type())?,
+ ))
+ })
+ .collect::<Result<Vec<_>>>()?,
+ )
+ }
+ DataType::Map(map_type) => {
+ map_value(array, row, map_type.key_type(), map_type.value_type())?
+ }
+ DataType::Multiset(multiset) => map_value(
+ array,
+ row,
+ multiset.element_type(),
+ &DataType::Int(crate::spec::IntType::new()),
+ )?,
+ other => {
+ return Err(Error::Unsupported {
+ message: format!("Avro writer does not support data type {other:?}"),
+ });
+ }
+ };
+ Ok(value)
+}
+
+fn map_value(array: &dyn Array, row: usize, key: &DataType, value: &DataType) -> Result<Value> {
+ let map = downcast::<MapArray>(array)?;
+ let start = map.value_offsets()[row] as usize;
+ let end = map.value_offsets()[row + 1] as usize;
+ if matches!(key, DataType::Char(_) | DataType::VarChar(_)) {
+ let mut entries = HashMap::with_capacity(end - start);
+ for index in start..end {
+ let map_key = downcast::<StringArray>(map.keys())?.value(index);
+ entries.insert(
+ map_key.to_string(),
+ value_at(map.values().as_ref(), index, value)?,
+ );
+ }
+ Ok(Value::Map(entries))
+ } else {
+ let mut entries = Vec::with_capacity(end - start);
+ for index in start..end {
+ entries.push(Value::Record(vec![
+ ("key".into(), value_at(map.keys(), index, key)?),
+ (
+ "value".into(),
+ value_at(map.values().as_ref(), index, value)?,
+ ),
+ ]));
+ }
+ Ok(Value::Array(entries))
+ }
+}
+
+fn decimal_bytes(value: i128) -> Vec<u8> {
+ let bytes = value.to_be_bytes();
+ let mut first = 0;
+ while first < bytes.len() - 1
+ && ((bytes[first] == 0 && bytes[first + 1] & 0x80 == 0)
+ || (bytes[first] == 0xff && bytes[first + 1] & 0x80 != 0))
+ {
+ first += 1;
+ }
+ bytes[first..].to_vec()
+}
+
+fn seconds_to_millis(seconds: i64) -> Result<i64> {
+ seconds.checked_mul(1000).ok_or_else(|| Error::DataInvalid {
+ message: "Avro timestamp seconds overflow milliseconds".into(),
+ source: None,
+ })
+}
+
+fn downcast<T: Array + 'static>(array: &dyn Array) -> Result<&T> {
+ array
+ .as_any()
+ .downcast_ref::<T>()
+ .ok_or_else(|| Error::DataInvalid {
+ message: format!(
+ "Avro writer expected Arrow array {}",
+ std::any::type_name::<T>()
+ ),
+ source: None,
+ })
+}
+
+#[cfg(test)]
+mod tests {
+ use std::sync::Arc;
+
+ use arrow_array::builder::{Int32Builder, MapBuilder, MapFieldNames, StringBuilder};
+ use arrow_array::types::Int32Type;
+ use arrow_array::{
+ ArrayRef, Float32Array, Int32Array, Int64Array, LargeBinaryArray, StringArray,
+ };
+ use arrow_buffer::{BooleanBuffer, NullBuffer};
+ use futures::TryStreamExt;
+
+ use super::*;
+ use crate::arrow::build_target_arrow_schema;
+ use crate::arrow::format::{avro::AvroFormatReader, FormatFileReader};
+ use crate::btree::test_util::BytesFileRead;
+ use crate::catalog::Identifier;
+ use crate::io::FileIOBuilder;
+ use crate::spec::{
+ ArrayType, BigIntType, BlobType, FloatType, IntType, MapType, RowType, VarCharType,
+ VectorType,
+ };
+
+ fn fields() -> Vec<DataField> {
+ vec![
+ DataField::new(0, "id".into(), DataType::BigInt(BigIntType::new())),
+ DataField::new(
+ 1,
+ "name".into(),
+ DataType::VarChar(VarCharType::string_type()),
+ ),
+ DataField::new(2, "score".into(), DataType::Int(IntType::new())),
+ ]
+ }
+
+ #[tokio::test]
+ async fn streams_multiple_nullable_batches_as_one_java_readable_ocf() {
+ let fields = fields();
+ let schema = build_target_arrow_schema(&fields).unwrap();
+ let file_io = FileIOBuilder::new("memory").build().unwrap();
+ let path = "memory:/avro-writer/part.avro";
+ let output = file_io.new_output(path).unwrap();
+ let mut writer =
+ AvroFormatWriter::new(&output, schema.clone(), fields.clone(), "zstd", 1, None)
+ .await
+ .unwrap();
+ let first = RecordBatch::try_new(
+ schema.clone(),
+ vec![
+ Arc::new(Int64Array::from(vec![Some(1), None])) as ArrayRef,
+ Arc::new(StringArray::from(vec![Some("a"), None])),
+ Arc::new(Int32Array::from(vec![Some(10), Some(-5)])),
+ ],
+ )
+ .unwrap();
+ let second = RecordBatch::try_new(
+ schema,
+ vec![
+ Arc::new(Int64Array::from(vec![Some(3)])) as ArrayRef,
+ Arc::new(StringArray::from(vec![Some("c")])),
+ Arc::new(Int32Array::from(vec![None])),
+ ],
+ )
+ .unwrap();
+ writer.write(&first).await.unwrap();
+ let first_size = writer.num_bytes();
+ writer.write(&second).await.unwrap();
+ assert!(writer.num_bytes() > first_size);
+ assert_eq!(writer.in_progress_size(), 0);
+ let size = Box::new(writer).close().await.unwrap().file_size;
+
+ let bytes = file_io.new_input(path).unwrap().read().await.unwrap();
+ assert_eq!(size, bytes.len() as u64);
+ let reader = apache_avro::Reader::new(bytes.as_ref()).unwrap();
+ let rows = reader.collect::<std::result::Result<Vec<_>, _>>().unwrap();
+ assert_eq!(rows.len(), 3);
+ let Value::Record(first) = &rows[0] else {
+ panic!("expected record");
+ };
+ assert_eq!(first[0].1, Value::Union(1, Box::new(Value::Long(1))));
+ assert_eq!(
+ first[1].1,
+ Value::Union(1, Box::new(Value::String("a".into())))
+ );
+ let Value::Record(second) = &rows[1] else {
+ panic!("expected record");
+ };
+ assert_eq!(second[0].1, Value::Union(0, Box::new(Value::Null)));
+ assert_eq!(second[1].1, Value::Union(0, Box::new(Value::Null)));
+
+ let decoded = AvroFormatReader
+ .read_batch_stream(
+ Box::new(BytesFileRead(bytes.clone())),
+ bytes.len() as u64,
+ &fields,
+ None,
+ None,
+ None,
+ )
+ .await
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+ assert_eq!(decoded.iter().map(RecordBatch::num_rows).sum::<usize>(), 3);
+ let ids = decoded[0]
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap();
+ assert_eq!(ids.iter().collect::<Vec<_>>(), vec![Some(1), None, Some(3)]);
+ }
+
+ #[test]
+ fn schema_matches_java_nullable_fields_and_logical_types() {
+ use crate::spec::{DateType, DecimalType, TimeType, TimestampType};
+ let fields = vec![
+ DataField::new(0, "day".into(), DataType::Date(DateType::new())),
+ DataField::new(1, "at".into(), DataType::Time(TimeType::new(3).unwrap())),
+ DataField::new(
+ 2,
+ "created".into(),
+ DataType::Timestamp(TimestampType::new(6).unwrap()),
+ ),
+ DataField::new(
+ 3,
+ "amount".into(),
+ DataType::Decimal(DecimalType::new(12, 2).unwrap()),
+ ),
+ ];
+ let schema = row_schema(&fields).unwrap();
+ let actual = schema.to_string();
+ for logical in ["date", "time-millis", "timestamp-micros", "decimal"] {
+ assert!(actual.contains(logical), "missing {logical}: {actual}");
+ }
+ for field in schema["fields"].as_array().unwrap() {
+ assert_eq!(field["type"][0], "null");
+ assert_eq!(field["default"], JsonValue::Null);
+ }
+ Schema::parse_str(&actual).unwrap();
+ }
+
+ #[test]
+ fn decimal_encoding_keeps_two_complement_sign_and_minimal_length() {
+ assert_eq!(decimal_bytes(0), vec![0]);
+ assert_eq!(decimal_bytes(127), vec![127]);
+ assert_eq!(decimal_bytes(128), vec![0, 128]);
+ assert_eq!(decimal_bytes(-1), vec![255]);
+ assert_eq!(decimal_bytes(-129), vec![255, 127]);
+ }
+
+ #[test]
+ fn blob_fields_require_java_serializable_references() {
+ let typ = DataType::Blob(BlobType::new());
+ let descriptor = BlobDescriptor::new("memory:/blob-data".into(), 2, 4).serialize();
+ let view = BlobViewStruct::new(Identifier::new("db", "source"), 3, 42)
+ .serialize()
+ .unwrap();
+ let valid = LargeBinaryArray::from(vec![
+ Some(descriptor.as_slice()),
+ Some(view.as_slice()),
+ None,
+ ]);
+ assert_eq!(
+ value_at(&valid, 0, &typ).unwrap(),
+ Value::Union(1, Box::new(Value::Bytes(descriptor)))
+ );
+ assert_eq!(
+ value_at(&valid, 1, &typ).unwrap(),
+ Value::Union(1, Box::new(Value::Bytes(view)))
+ );
+ assert_eq!(
+ value_at(&valid, 2, &typ).unwrap(),
+ Value::Union(0, Box::new(Value::Null))
+ );
+ let raw = LargeBinaryArray::from(vec![b"raw blob".as_slice()]);
+ assert!(value_at(&raw, 0, &typ)
+ .unwrap_err()
+ .to_string()
+ .contains("BlobDescriptor or BlobViewStruct"));
+ }
+
+ #[test]
+ fn zero_precision_timestamps_use_java_millisecond_logical_type() {
+ use crate::spec::{LocalZonedTimestampType, TimestampType};
+ let local = DataType::LocalZonedTimestamp(LocalZonedTimestampType::new(0).unwrap());
+ let plain = DataType::Timestamp(TimestampType::new(0).unwrap());
+ let seconds = TimestampSecondArray::from(vec![Some(3), Some(-2), None]);
+ assert_eq!(
+ value_at(&seconds, 0, &plain).unwrap(),
+ Value::Union(1, Box::new(Value::TimestampMillis(3000)))
+ );
+ assert_eq!(
+ value_at(&seconds, 1, &local).unwrap(),
+ Value::Union(1, Box::new(Value::LocalTimestampMillis(-2000)))
+ );
+ assert_eq!(
+ value_at(&seconds, 2, &plain).unwrap(),
+ Value::Union(0, Box::new(Value::Null))
+ );
+ let extreme = TimestampSecondArray::from(vec![i64::MAX]);
+ assert!(value_at(&extreme, 0, &plain)
+ .unwrap_err()
+ .to_string()
+ .contains("overflow"));
+ }
+
+ #[test]
+ fn nested_schema_matches_java_array_map_and_row_layouts() {
+ let fields = vec![
+ DataField::new(
+ 0,
+ "numbers".into(),
+ DataType::Array(ArrayType::new(DataType::Int(IntType::new()))),
+ ),
+ DataField::new(
+ 1,
+ "named".into(),
+ DataType::Map(MapType::new(
+ DataType::VarChar(VarCharType::string_type()),
+ DataType::Int(IntType::new()),
+ )),
+ ),
+ DataField::new(
+ 2,
+ "numbered".into(),
+ DataType::Map(MapType::new(
+ DataType::Int(IntType::new()),
+ DataType::VarChar(VarCharType::string_type()),
+ )),
+ ),
+ DataField::new(
+ 3,
+ "nested".into(),
+ DataType::Row(RowType::new(vec![DataField::new(
+ 4,
+ "label".into(),
+ DataType::VarChar(VarCharType::string_type()),
+ )])),
+ ),
+ ];
+ let schema = row_schema(&fields).unwrap();
+ assert_eq!(schema["fields"][0]["type"][1]["type"], "array");
+ assert_eq!(schema["fields"][1]["type"][1]["type"], "map");
+ assert_eq!(schema["fields"][2]["type"][1]["logicalType"], "map");
+ assert_eq!(schema["fields"][2]["type"][1]["items"]["type"], "record");
+ assert_eq!(schema["fields"][3]["type"][1]["fields"][0]["name"], "label");
+ Schema::parse_str(&schema.to_string()).unwrap();
+ }
+
+ #[test]
+ fn nested_arrow_values_use_java_map_and_array_encodings() {
+ let int = DataType::Int(IntType::new());
+ let text = DataType::VarChar(VarCharType::string_type());
+ let lists = arrow_array::ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
+ Some(vec![Some(1), None, Some(-2)]),
+ None,
+ ]);
+ let list_type = DataType::Array(ArrayType::new(int.clone()));
+ assert_eq!(
+ value_at(&lists, 0, &list_type).unwrap(),
+ Value::Union(
+ 1,
+ Box::new(Value::Array(vec![
+ Value::Union(1, Box::new(Value::Int(1))),
+ Value::Union(0, Box::new(Value::Null)),
+ Value::Union(1, Box::new(Value::Int(-2))),
+ ]))
+ )
+ );
+ assert_eq!(
+ value_at(&lists, 1, &list_type).unwrap(),
+ Value::Union(0, Box::new(Value::Null))
+ );
+
+ let names = MapFieldNames {
+ entry: "entries".into(),
+ key: "key".into(),
+ value: "value".into(),
+ };
+ let mut string_map = MapBuilder::new(
+ Some(names.clone()),
+ StringBuilder::new(),
+ Int32Builder::new(),
+ );
+ string_map.keys().append_value("x");
+ string_map.values().append_value(7);
+ string_map.append(true).unwrap();
+ let string_map = string_map.finish();
+ let string_map_type = DataType::Map(MapType::new(text.clone(), int.clone()));
+ assert_eq!(
+ value_at(&string_map, 0, &string_map_type).unwrap(),
+ Value::Union(
+ 1,
+ Box::new(Value::Map(HashMap::from([(
+ "x".into(),
+ Value::Union(1, Box::new(Value::Int(7))),
+ )])))
+ )
+ );
+
+ let mut int_map = MapBuilder::new(Some(names), Int32Builder::new(), StringBuilder::new());
+ int_map.keys().append_value(3);
+ int_map.values().append_value("three");
+ int_map.append(true).unwrap();
+ let int_map = int_map.finish();
+ let int_map_type = DataType::Map(MapType::new(int, text));
+ assert_eq!(
+ value_at(&int_map, 0, &int_map_type).unwrap(),
+ Value::Union(
+ 1,
+ Box::new(Value::Array(vec![Value::Record(vec![
+ ("key".into(), Value::Union(1, Box::new(Value::Int(3)))),
+ (
+ "value".into(),
+ Value::Union(1, Box::new(Value::String("three".into())))
+ ),
+ ])]))
+ )
+ );
+ }
+
+ #[tokio::test]
+ async fn nested_arrays_and_maps_survive_file_roundtrip() {
+ let number_type = DataType::Int(IntType::new());
+ let text_type = DataType::VarChar(VarCharType::string_type());
+ let fields = vec![
+ DataField::new(
+ 0,
+ "numbers".into(),
+ DataType::Array(ArrayType::new(number_type.clone())),
+ ),
+ DataField::new(
+ 1,
+ "named".into(),
+ DataType::Map(MapType::new(text_type, number_type)),
+ ),
+ ];
+ let numbers = arrow_array::ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
+ Some(vec![Some(1), None, Some(-2)]),
+ Some(vec![]),
+ ]);
+ let mut map = MapBuilder::new(
+ Some(MapFieldNames {
+ entry: "entries".into(),
+ key: "key".into(),
+ value: "value".into(),
+ }),
+ StringBuilder::new(),
+ Int32Builder::new(),
+ );
+ map.keys().append_value("first");
+ map.values().append_value(7);
+ map.append(true).unwrap();
+ map.append(true).unwrap();
+ let map = map.finish();
+ let schema = Arc::new(arrow_schema::Schema::new(vec![
+ arrow_schema::Field::new("numbers", numbers.data_type().clone(), true),
+ arrow_schema::Field::new("named", map.data_type().clone(), true),
+ ]));
+ let batch =
+ RecordBatch::try_new(schema.clone(), vec![Arc::new(numbers), Arc::new(map)]).unwrap();
+ let file_io = FileIOBuilder::new("memory").build().unwrap();
+ let path = "memory:/avro-writer/nested.avro";
+ let output = file_io.new_output(path).unwrap();
+ let mut writer = AvroFormatWriter::new(&output, schema, fields.clone(), "zstd", 1, None)
+ .await
+ .unwrap();
+ writer.write(&batch).await.unwrap();
+ Box::new(writer).close().await.unwrap();
+ let bytes = file_io.new_input(path).unwrap().read().await.unwrap();
+ let decoded = AvroFormatReader
+ .read_batch_stream(
+ Box::new(BytesFileRead(bytes.clone())),
+ bytes.len() as u64,
+ &fields,
+ None,
+ None,
+ None,
+ )
+ .await
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+ assert_eq!(decoded.len(), 1);
+ assert_eq!(decoded[0].num_rows(), 2);
+ let decoded_numbers = decoded[0]
+ .column(0)
+ .as_any()
+ .downcast_ref::<arrow_array::ListArray>()
+ .unwrap();
+ assert_eq!(decoded_numbers.value_length(0), 3);
+ assert_eq!(decoded_numbers.value_length(1), 0);
+ let decoded_map = decoded[0]
+ .column(1)
+ .as_any()
+ .downcast_ref::<arrow_array::MapArray>()
+ .unwrap();
+ assert_eq!(decoded_map.value_length(0), 1);
+ assert_eq!(decoded_map.value_length(1), 0);
+ }
+
+ #[tokio::test]
+ async fn vector_uses_java_array_schema_and_roundtrips_nullable_rows() {
+ let fields = vec![DataField::new(
+ 0,
+ "embedding".into(),
+ DataType::Vector(VectorType::new(3, DataType::Float(FloatType::new())).unwrap()),
+ )];
+ let schema = build_target_arrow_schema(&fields).unwrap();
+ let values: ArrayRef = Arc::new(Float32Array::from(vec![
+ Some(1.0),
+ Some(2.0),
+ Some(3.0),
+ None,
+ None,
+ None,
+ ]));
+ let vector: ArrayRef = Arc::new(
+ FixedSizeListArray::try_new(
+ Arc::new(arrow_schema::Field::new(
+ "element",
+ arrow_schema::DataType::Float32,
+ true,
+ )),
+ 3,
+ values,
+ Some(NullBuffer::new(BooleanBuffer::from(vec![true, false]))),
+ )
+ .unwrap(),
+ );
+ let batch = RecordBatch::try_new(schema.clone(), vec![vector]).unwrap();
+ let file_io = FileIOBuilder::new("memory").build().unwrap();
+ let path = "memory:/avro-writer/vector.avro";
+ let output = file_io.new_output(path).unwrap();
+ let mut writer = AvroFormatWriter::new(&output, schema, fields.clone(), "snappy", 1, None)
+ .await
+ .unwrap();
+ writer.write(&batch).await.unwrap();
+ Box::new(writer).close().await.unwrap();
+
+ let bytes = file_io.new_input(path).unwrap().read().await.unwrap();
+ assert_eq!(
+ row_schema(&fields).unwrap()["fields"][0]["type"][1]["type"],
+ "array"
+ );
+ let decoded = AvroFormatReader
+ .read_batch_stream(
+ Box::new(BytesFileRead(bytes.clone())),
+ bytes.len() as u64,
+ &fields,
+ None,
+ None,
+ None,
+ )
+ .await
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+ assert_eq!(decoded.len(), 1);
+ assert_eq!(decoded[0], batch);
+ }
+
+ #[tokio::test]
+ async fn rejects_unknown_codec_before_opening_output() {
+ let fields = fields();
+ let schema = build_target_arrow_schema(&fields).unwrap();
+ let file_io = FileIOBuilder::new("memory").build().unwrap();
+ let output = file_io.new_output("memory:/avro-writer/bad.avro").unwrap();
+ let result = AvroFormatWriter::new(&output, schema, fields, "brotli", 1, None).await;
+ assert!(matches!(result, Err(Error::Unsupported { .. })));
+ assert!(!output.exists().await.unwrap());
+ }
+
+ #[tokio::test]
+ async fn avro_codec_option_overrides_file_compression() {
+ let fields = fields();
+ let schema = build_target_arrow_schema(&fields).unwrap();
+ let file_io = FileIOBuilder::new("memory").build().unwrap();
+ let output = file_io
+ .new_output("memory:/avro-writer/override.avro")
+ .unwrap();
+ let options = HashMap::from([
+ ("avro.codec".to_string(), "snappy".to_string()),
+ ("file.block-size".to_string(), "32 b".to_string()),
+ ]);
+ let writer = AvroFormatWriter::new(&output, schema, fields, "zstd", 1, Some(&options))
+ .await
+ .unwrap();
+ assert_eq!(writer.codec, Codec::Snappy);
+ assert_eq!(writer.block_size, 32);
+ Box::new(writer).close().await.unwrap();
+ }
+
+ #[tokio::test]
+ async fn invalid_avro_block_size_fails_before_file_creation() {
+ let fields = fields();
+ let schema = build_target_arrow_schema(&fields).unwrap();
+ let file_io = FileIOBuilder::new("memory").build().unwrap();
+ let output = file_io
+ .new_output("memory:/avro-writer/bad-block.avro")
+ .unwrap();
+ for invalid in ["31", "2 gb", "-1", "abc"] {
+ let options = HashMap::from([("file.block-size".to_string(), invalid.to_string())]);
+ let error = AvroFormatWriter::new(
+ &output,
+ schema.clone(),
+ fields.clone(),
+ "zstd",
+ 1,
+ Some(&options),
+ )
+ .await
+ .err()
+ .unwrap();
+ assert!(
+ error.to_string().contains("file.block-size"),
+ "{invalid}: {error}"
+ );
+ }
+ assert!(!output.exists().await.unwrap());
+ }
+}
diff --git a/crates/paimon/src/arrow/format/mod.rs b/crates/paimon/src/arrow/format/mod.rs
index 8f90214f..37607d31 100644
--- a/crates/paimon/src/arrow/format/mod.rs
+++ b/crates/paimon/src/arrow/format/mod.rs
@@ -16,6 +16,7 @@
// under the License.
mod avro;
+mod avro_write;
pub(crate) mod blob;
mod mosaic;
mod orc;
@@ -445,6 +446,7 @@ fn supported_write_formats() -> Vec<&'static str> {
vec![
".parquet",
".blob",
+ ".avro",
".row",
#[cfg(feature = "vortex")]
".vortex",
@@ -482,6 +484,22 @@ pub(crate) async fn create_format_writer(
Ok(Box::new(
blob::BlobFormatWriter::new(output, file_io).await?,
))
+ } else if lower.ends_with(".avro") {
+ let fields = match write_fields {
+ Some(fields) => fields.to_vec(),
+ None => row::row_type_from_arrow_schema(&schema)?,
+ };
+ Ok(Box::new(
+ avro_write::AvroFormatWriter::new(
+ output,
+ schema,
+ fields,
+ compression,
+ zstd_level,
+ format_options,
+ )
+ .await?,
+ ))
} else if lower.ends_with(".row") {
let row_type = match write_fields {
Some(fields) => fields.to_vec(),
diff --git a/crates/paimon/src/spec/core_options.rs b/crates/paimon/src/spec/core_options.rs
index f100c6c3..114b89ab 100644
--- a/crates/paimon/src/spec/core_options.rs
+++ b/crates/paimon/src/spec/core_options.rs
@@ -1292,6 +1292,18 @@ impl<'a> CoreOptions<'a> {
.unwrap_or(DEFAULT_TARGET_FILE_SIZE)
}
+ /// Explicit `file.block-size`, in bytes. Formats choose their own default.
+ pub(crate) fn file_block_size(&self) -> crate::Result<Option<i64>> {
+ self.options
+ .get("file.block-size")
+ .map(|raw| {
+ parse_memory_size(raw).ok_or_else(|| crate::Error::ConfigInvalid {
+ message: format!("Invalid file.block-size: {raw}"),
+ })
+ })
+ .transpose()
+ }
+
pub fn blob_target_file_size(&self) -> i64 {
self.options
.get("blob.target-file-size")
@@ -2413,6 +2425,23 @@ mod tests {
assert_eq!(parse_memory_size("abc"), None);
}
+ #[test]
+ fn file_block_size_reports_invalid_input_and_preserves_format_default() {
+ let empty = HashMap::new();
+ assert_eq!(CoreOptions::new(&empty).file_block_size().unwrap(), None);
+ let valid = HashMap::from([("file.block-size".into(), "64 kb".into())]);
+ assert_eq!(
+ CoreOptions::new(&valid).file_block_size().unwrap(),
+ Some(65_536)
+ );
+ let invalid = HashMap::from([("file.block-size".into(), "nope".into())]);
+ assert!(CoreOptions::new(&invalid)
+ .file_block_size()
+ .unwrap_err()
+ .to_string()
+ .contains("file.block-size"));
+ }
+
#[test]
fn test_parse_memory_size_accepts_every_java_unit_spelling() {
// Every alias in Java `MemorySize.MemoryUnit`, which lists three
diff --git a/crates/paimon/src/table/format_partition_location.rs b/crates/paimon/src/table/format_partition_location.rs
new file mode 100644
index 00000000..907f43c6
--- /dev/null
+++ b/crates/paimon/src/table/format_partition_location.rs
@@ -0,0 +1,544 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Resolve and validate the locations of catalog-managed Format Table partitions.
+//! A catalog location is untrusted metadata: its path must be checked before a
+//! scan opens files, and one partition must never claim another one's data.
+
+use std::collections::HashMap;
+
+use url::Url;
+
+use crate::spec::Partition;
+use crate::Error;
+
+type SpecKey = Vec<(String, String)>;
+type FileSystemKey = (String, String);
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+struct StoragePath {
+ scheme: String,
+ authority: String,
+ path: String,
+ location: String,
+}
+
+impl StoragePath {
+ fn parse(location: &str, custom: bool) -> crate::Result<Self> {
+ validate_location_text(location)?;
+ let decoded = if custom {
+ let decoded = decode_percent_once(location)?;
+ if decoded.contains('%') {
+ return Err(invalid_location());
+ }
+ validate_location_text(&decoded)?;
+ decoded
+ } else {
+ location.to_string()
+ };
+
+ let url = if decoded.starts_with('/') {
+ Url::from_file_path(&decoded).map_err(|_| invalid_location())?
+ } else {
+ Url::parse(&decoded).map_err(|_| invalid_location())?
+ };
+ let abfs = matches!(url.scheme().to_ascii_lowercase().as_str(), "abfs" | "abfss");
+ if url.cannot_be_a_base()
+ || url.query().is_some()
+ || url.fragment().is_some()
+ || url.path().is_empty()
+ || url.path() == "/"
+ || (!url.username().is_empty() && (!abfs || url.username().contains(':')))
+ || url.password().is_some()
+ {
+ return Err(invalid_location());
+ }
+
+ let scheme = match url.scheme().to_ascii_lowercase().as_str() {
+ "s3a" | "s3n" => "s3".to_string(),
+ "abfss" => "abfs".to_string(),
+ other => other.to_string(),
+ };
+ let host = url.host_str().unwrap_or_default().to_ascii_lowercase();
+ let authority = if abfs && !url.username().is_empty() {
+ format!("{}@{host}", url.username().to_ascii_lowercase())
+ } else {
+ host
+ };
+ if custom
+ && ((scheme == "file" && !authority.is_empty())
+ || (scheme != "file" && scheme != "hdfs" && authority.is_empty())
+ || (scheme == "hdfs" && authority.is_empty())
+ || scheme == "viewfs")
+ {
+ return Err(invalid_location());
+ }
+ let authority = match url.port() {
+ Some(port) if scheme == "hdfs" && port == 8020 => authority,
+ Some(port) => format!("{authority}:{port}"),
+ None => authority,
+ };
+ // Hadoop Path collapses redundant separators. Ownership must use the
+ // same directory identity or `a//b` can evade an `a/b` overlap check.
+ let path = format!(
+ "/{}",
+ url.path()
+ .split('/')
+ .filter(|segment| !segment.is_empty())
+ .collect::<Vec<_>>()
+ .join("/")
+ );
+ Ok(Self {
+ scheme,
+ authority,
+ path,
+ location: url.to_string().trim_end_matches('/').to_string(),
+ })
+ }
+
+ fn file_system(&self) -> FileSystemKey {
+ (self.scheme.clone(), self.authority.clone())
+ }
+
+ fn same_directory(&self, other: &Self) -> bool {
+ self.file_system() == other.file_system() && self.path == other.path
+ }
+
+ fn overlaps(&self, other: &Self) -> bool {
+ self.file_system() == other.file_system()
+ && (self.path == other.path
+ || self.path.starts_with(&format!("{}/", other.path))
+ || other.path.starts_with(&format!("{}/", self.path)))
+ }
+}
+
+#[derive(Default)]
+struct OwnershipNode {
+ owned: bool,
+ children: HashMap<String, OwnershipNode>,
+}
+
+/// Mirrors Java's `FormatTablePartitionPathResolver` for paths the Rust
+/// `FileIO` can address. A duplicate catalog row is harmless only if it names
+/// the same spec and the same canonical location.
+pub(super) struct FormatPartitionLocations {
+ table_root: StoragePath,
+ by_spec: HashMap<SpecKey, StoragePath>,
+ roots: HashMap<FileSystemKey, OwnershipNode>,
+}
+
+impl FormatPartitionLocations {
+ pub(super) fn new(table_path: &str) -> crate::Result<Self> {
+ Ok(Self {
+ table_root: StoragePath::parse(table_path, false)?,
+ by_spec: HashMap::new(),
+ roots: HashMap::new(),
+ })
+ }
+
+ /// Return `None` for an identical duplicate registry row, otherwise the
+ /// path to list. An explicit path spelling the default directory is a
+ /// rebind, not an external location.
+ pub(super) fn resolve(
+ &mut self,
+ partition: &Partition,
+ default_path: &str,
+ ) -> crate::Result<Option<String>> {
+ let default = StoragePath::parse(default_path, false)?;
+ let custom = partition
+ .options
+ .as_ref()
+ .and_then(|options| options.get("path"));
+ let selected = if let Some(custom) = custom {
+ // The default path can contain escaped partition values. Compare
+ // that spelling before decoding a user-supplied external location.
+ let default_spelling = StoragePath::parse(custom, false);
+ let parsed = if default_spelling
+ .as_ref()
+ .is_ok_and(|path| path.same_directory(&default))
+ {
+ Ok(default.clone())
+ } else {
+ StoragePath::parse(custom, true)
+ }
+ .map_err(|source| Error::DataInvalid {
+ message: format!(
+ "Catalog returned an invalid custom location for partition {:?}",
+ partition.spec
+ ),
+ source: Some(Box::new(source)),
+ })?;
+ if parsed.same_directory(&default) {
+ default.clone()
+ } else {
+ if parsed.overlaps(&self.table_root) {
+ return Err(Error::DataInvalid {
+ message: format!(
+ "Custom location for partition {:?} overlaps the Format Table directory",
+ partition.spec
+ ),
+ source: None,
+ });
+ }
+ parsed
+ }
+ } else {
+ default.clone()
+ };
+
+ let spec = spec_key(&partition.spec);
+ if let Some(previous) = self.by_spec.get(&spec) {
+ if previous.same_directory(&selected) {
+ return Ok(None);
+ }
+ return Err(overlapping_locations());
+ }
+ let node = self.roots.entry(selected.file_system()).or_default();
+ let mut node = node;
+ for segment in selected.path.trim_start_matches('/').split('/') {
+ if node.owned {
+ return Err(overlapping_locations());
+ }
+ node = node.children.entry(segment.to_string()).or_default();
+ }
+ if node.owned || !node.children.is_empty() {
+ return Err(overlapping_locations());
+ }
+ node.owned = true;
+ let location = if selected.same_directory(&default) {
+ default_path.to_string()
+ } else {
+ selected.location.clone()
+ };
+ self.by_spec.insert(spec, selected);
+ Ok(Some(location))
+ }
+}
+
+fn spec_key(spec: &HashMap<String, String>) -> SpecKey {
+ let mut key = spec
+ .iter()
+ .map(|(name, value)| (name.clone(), value.clone()))
+ .collect::<Vec<_>>();
+ key.sort();
+ key
+}
+
+fn validate_location_text(location: &str) -> crate::Result<()> {
+ if location.is_empty()
+ || location.trim() != location
+ || location
+ .chars()
+ .any(|c| matches!(c, '?' | '#' | '\\') || c.is_control())
+ || location.split('/').any(|part| matches!(part, "." | ".."))
+ {
+ return Err(invalid_location());
+ }
+ Ok(())
+}
+
+fn decode_percent_once(value: &str) -> crate::Result<String> {
+ let bytes = value.as_bytes();
+ let mut decoded = Vec::with_capacity(bytes.len());
+ let mut index = 0;
+ while index < bytes.len() {
+ if bytes[index] == b'%' {
+ if index + 2 >= bytes.len() {
+ return Err(invalid_location());
+ }
+ let high = (bytes[index + 1] as char)
+ .to_digit(16)
+ .ok_or_else(invalid_location)?;
+ let low = (bytes[index + 2] as char)
+ .to_digit(16)
+ .ok_or_else(invalid_location)?;
+ decoded.push(((high << 4) | low) as u8);
+ index += 3;
+ } else {
+ decoded.push(bytes[index]);
+ index += 1;
+ }
+ }
+ String::from_utf8(decoded).map_err(|_| invalid_location())
+}
+
+fn invalid_location() -> Error {
+ Error::DataInvalid {
+ message: "Invalid custom Format Table partition location".into(),
+ source: None,
+ }
+}
+
+fn overlapping_locations() -> Error {
+ Error::DataInvalid {
+ message: "Catalog returned overlapping locations for Format Table partitions".into(),
+ source: None,
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ fn partition(name: &str, path: Option<&str>) -> Partition {
+ Partition {
+ spec: HashMap::from([("dt".to_string(), name.to_string())]),
+ record_count: Partition::UNKNOWN,
+ file_size_in_bytes: Partition::UNKNOWN,
+ file_count: Partition::UNKNOWN,
+ last_file_creation_time: Partition::UNKNOWN,
+ total_buckets: Partition::UNKNOWN_TOTAL_BUCKETS,
+ done: false,
+ created_at: None,
+ created_by: None,
+ updated_at: None,
+ updated_by: None,
+ options: path.map(|path| HashMap::from([("path".to_string(), path.to_string())])),
+ }
+ }
+
+ fn default(name: &str) -> String {
+ format!("file:///warehouse/events/dt={name}")
+ }
+
+ #[test]
+ fn custom_partition_reads_external_location_and_default_rebind() {
+ let mut locations = FormatPartitionLocations::new("file:///warehouse/events").unwrap();
+ assert_eq!(
+ locations
+ .resolve(&partition("a", Some("file:/external/a")), &default("a"))
+ .unwrap(),
+ Some("file:///external/a".to_string())
+ );
+ assert_eq!(
+ locations
+ .resolve(
+ &partition("b", Some("file:/warehouse/events/dt=b")),
+ &default("b")
+ )
+ .unwrap(),
+ Some(default("b"))
+ );
+ assert_eq!(
+ locations
+ .resolve(&partition("c", None), &default("c"))
+ .unwrap(),
+ Some(default("c"))
+ );
+ }
+
+ #[test]
+ fn duplicate_registry_rows_must_have_the_same_spec_and_location() {
+ let mut locations = FormatPartitionLocations::new("file:///warehouse/events").unwrap();
+ let a = partition("a", Some("file:/external/a"));
+ assert!(locations.resolve(&a, &default("a")).unwrap().is_some());
+ assert!(locations.resolve(&a, &default("a")).unwrap().is_none());
+ let conflicting = partition("a", Some("file:/external/other"));
+ assert!(locations
+ .resolve(&conflicting, &default("a"))
+ .unwrap_err()
+ .to_string()
+ .contains("overlapping locations"));
+ }
+
+ #[test]
+ fn custom_location_cannot_claim_table_data_or_another_partition() {
+ for invalid in [
+ "file:/warehouse/events",
+ "file:/warehouse/events/dt=b",
+ "file:/warehouse/events/dt=b/child",
+ "file:/warehouse",
+ ] {
+ let mut locations = FormatPartitionLocations::new("file:///warehouse/events").unwrap();
+ let error = locations
+ .resolve(&partition("a", Some(invalid)), &default("a"))
+ .unwrap_err();
+ assert!(error.to_string().contains("overlaps"), "{invalid}: {error}");
+ }
+
+ let mut locations = FormatPartitionLocations::new("file:///warehouse/events").unwrap();
+ locations
+ .resolve(
+ &partition("a", Some("file:/external/parent")),
+ &default("a"),
+ )
+ .unwrap();
+ let error = locations
+ .resolve(
+ &partition("b", Some("file:/external/parent/child")),
+ &default("b"),
+ )
+ .unwrap_err();
+ assert!(error.to_string().contains("overlapping locations"));
+ // A textual prefix without a path-segment boundary is a different directory.
+ assert!(locations
+ .resolve(
+ &partition("c", Some("file:/external/parent-sibling")),
+ &default("c"),
+ )
+ .unwrap()
+ .is_some());
+ }
+
+ #[test]
+ fn invalid_custom_locations_fail_before_listing() {
+ for invalid in [
+ "file:/external/../secret",
+ "file:/external/%2e%2e/secret",
+ "file:/external/%252e%252e/secret",
+ "file:/external/%2F..%2Fsecret",
+ "file:/external/%ZZ",
+ "file:/external/%FF",
+ "file:/external/a?query=1",
+ "file:/external/a#fragment",
+ "file:/external/\\secret",
+ "file:/external/a\n",
+ " file:/external/a",
+ "file://remote-host/external/a",
+ "relative/a",
+ "viewfs://mount/external/a",
+ "s3:///external/a",
+ ] {
+ let mut locations = FormatPartitionLocations::new("file:///warehouse/events").unwrap();
+ let error = locations
+ .resolve(&partition("a", Some(invalid)), &default("a"))
+ .unwrap_err();
+ assert!(
+ error.to_string().contains("invalid custom location"),
+ "{invalid}: {error}"
+ );
+ }
+ }
+
+ #[test]
+ fn ownership_uses_canonical_file_system_identity() {
+ let mut locations = FormatPartitionLocations::new("s3://warehouse/events").unwrap();
+ locations
+ .resolve(
+ &partition("a", Some("s3a://bucket/external/a")),
+ "s3://warehouse/events/dt=a",
+ )
+ .unwrap();
+ let error = locations
+ .resolve(
+ &partition("b", Some("s3://bucket/external/a/child")),
+ "s3://warehouse/events/dt=b",
+ )
+ .unwrap_err();
+ assert!(error.to_string().contains("overlapping locations"));
+ }
+
+ #[test]
+ fn abfs_container_userinfo_is_part_of_file_system_identity() {
+ let table = "abfss://warehouse@account.dfs.core.windows.net/events";
+ let mut locations = FormatPartitionLocations::new(table).unwrap();
+ let first = "abfs://archive@account.dfs.core.windows.net/external/a";
+ assert!(locations
+ .resolve(
+ &partition("a", Some(first)),
+ "abfss://warehouse@account.dfs.core.windows.net/events/dt=a",
+ )
+ .unwrap()
+ .is_some());
+ let same_container = "abfss://archive@ACCOUNT.dfs.core.windows.net/external/a/child";
+ let error = locations
+ .resolve(
+ &partition("b", Some(same_container)),
+ "abfss://warehouse@account.dfs.core.windows.net/events/dt=b",
+ )
+ .unwrap_err();
+ assert!(error.to_string().contains("overlapping locations"));
+ let another_container = "abfss://other@account.dfs.core.windows.net/external/a/child";
+ assert!(locations
+ .resolve(
+ &partition("c", Some(another_container)),
+ "abfss://warehouse@account.dfs.core.windows.net/events/dt=c",
+ )
+ .unwrap()
+ .is_some());
+ }
+
+ #[test]
+ fn redundant_separators_and_hdfs_default_port_cannot_evade_ownership() {
+ let mut file_locations = FormatPartitionLocations::new("file:///warehouse/events").unwrap();
+ file_locations
+ .resolve(&partition("a", Some("file:/external/a//b")), &default("a"))
+ .unwrap();
+ assert!(file_locations
+ .resolve(
+ &partition("b", Some("file:/external/a/b/child")),
+ &default("b"),
+ )
+ .unwrap_err()
+ .to_string()
+ .contains("overlapping locations"));
+
+ let mut hdfs_locations =
+ FormatPartitionLocations::new("hdfs://NameNode:8020/warehouse/events").unwrap();
+ hdfs_locations
+ .resolve(
+ &partition("a", Some("hdfs://NameNode/external/a")),
+ "hdfs://NameNode/warehouse/events/dt=a",
+ )
+ .unwrap();
+ assert!(hdfs_locations
+ .resolve(
+ &partition("b", Some("hdfs://namenode:8020/external/a/child")),
+ "hdfs://NameNode/warehouse/events/dt=b",
+ )
+ .unwrap_err()
+ .to_string()
+ .contains("overlapping locations"));
+ }
+
+ #[test]
+ fn escaped_default_partition_value_stays_one_path_segment() {
+ let mut locations = FormatPartitionLocations::new("file:///warehouse/events").unwrap();
+ let escaped_default = "file:///warehouse/events/dt=a%2Fb";
+ assert_eq!(
+ locations
+ .resolve(&partition("a/b", Some(escaped_default)), escaped_default)
+ .unwrap(),
+ Some(escaped_default.to_string())
+ );
+ // The external path below the decoded spelling is a different
+ // directory from the single escaped partition segment.
+ assert!(locations
+ .resolve(
+ &partition("other", Some("file:/external/a/b")),
+ &default("other"),
+ )
+ .unwrap()
+ .is_some());
+ }
+
+ #[test]
+ fn rejects_credentials_outside_abfs_and_invalid_abfs_userinfo() {
+ for invalid in [
+ "s3://user@bucket/external/a",
+ "file://user@localhost/external/a",
+ "abfs://container:password@account/external/a",
+ ] {
+ let mut locations = FormatPartitionLocations::new("file:///warehouse/events").unwrap();
+ assert!(locations
+ .resolve(&partition("a", Some(invalid)), &default("a"))
+ .unwrap_err()
+ .to_string()
+ .contains("invalid custom location"));
+ }
+ }
+}
diff --git a/crates/paimon/src/table/format_table_defaults.rs b/crates/paimon/src/table/format_table_defaults.rs
new file mode 100644
index 00000000..68b8fc47
--- /dev/null
+++ b/crates/paimon/src/table/format_table_defaults.rs
@@ -0,0 +1,1404 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Apply schema defaults to nullable Format Table input columns. Java's
+//! `DefaultValueRow` wraps the row after nullability validation and before
+//! partition extraction, so partition and file columns share this behavior.
+
+use std::sync::Arc;
+
+use arrow_array::{
+ Array, ArrayRef, BinaryArray, Date32Array, Int16Array, Int32Array, Int64Array, Int8Array,
+ ListArray, MapArray, RecordBatch, StringArray, StructArray, Time32MillisecondArray,
+ TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray,
+ TimestampSecondArray,
+};
+use arrow_buffer::{OffsetBuffer, ScalarBuffer};
+use arrow_schema::{
+ ArrowError, DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema,
+};
+
+use crate::spec::{DataField, DataType};
+use crate::{Error, Result};
+
+pub(super) struct FormatTableDefaults {
+ values: Vec<Option<ArrayRef>>,
+}
+
+impl FormatTableDefaults {
+ pub(super) fn new(fields: &[DataField], schema: &ArrowSchema) -> Result<Self> {
+ let values = fields
+ .iter()
+ .enumerate()
+ .map(|(index, field)| {
+ field
+ .default_value()
+ .map(|text| {
+ // Java DefaultValueUtils removes one pair of outer
+ // single quotes before casting from VARCHAR.
+ let text = text
+ .strip_prefix('\'')
+ .and_then(|text| text.strip_suffix('\''))
+ .unwrap_or(text);
+ let value = cast_default(field.data_type(), text, schema.field(index))
+ .map_err(|source| Error::DataInvalid {
+ message: format!(
+ "Unsupported default value '{}' for Format Table column '{}'",
+ field.default_value().unwrap_or_default(),
+ field.name()
+ ),
+ source: Some(Box::new(source)),
+ })?;
+ if value.is_null(0) {
+ return Err(Error::DataInvalid {
+ message: format!(
+ "Unsupported default value '{}' for Format Table column '{}'",
+ field.default_value().unwrap_or_default(),
+ field.name()
+ ),
+ source: None,
+ });
+ }
+ Ok(value)
+ })
+ .transpose()
+ })
+ .collect::<Result<Vec<_>>>()?;
+ Ok(Self { values })
+ }
+
+ pub(super) fn apply(&self, batch: &RecordBatch) -> Result<RecordBatch> {
+ if !self
+ .values
+ .iter()
+ .enumerate()
+ .any(|(index, value)| value.is_some() && batch.column(index).null_count() > 0)
+ {
+ return Ok(batch.clone());
+ }
+ let columns = batch
+ .columns()
+ .iter()
+ .enumerate()
+ .map(|(index, input)| {
+ let Some(value) = &self.values[index] else {
+ return Ok(input.clone());
+ };
+ if input.null_count() == 0 {
+ return Ok(input.clone());
+ }
+ let positions = arrow_array::UInt32Array::from(vec![0_u32; batch.num_rows()]);
+ let repeated = arrow_select::take::take(value.as_ref(), &positions, None).map_err(
+ |source| Error::DataInvalid {
+ message: format!("Cannot repeat Format Table default for column {index}"),
+ source: Some(Box::new(source)),
+ },
+ )?;
+ let present = arrow_array::BooleanArray::from_iter(
+ (0..input.len()).map(|row| Some(!input.is_null(row))),
+ );
+ arrow_select::zip::zip(&present, input, &repeated).map_err(|source| {
+ Error::DataInvalid {
+ message: format!("Cannot apply Format Table default for column {index}"),
+ source: Some(Box::new(source)),
+ }
+ })
+ })
+ .collect::<Result<Vec<_>>>()?;
+ RecordBatch::try_new(Arc::clone(&batch.schema()), columns).map_err(|source| {
+ Error::DataInvalid {
+ message: "Cannot assemble Format Table row with defaults".into(),
+ source: Some(Box::new(source)),
+ }
+ })
+ }
+}
+
+fn cast_default(
+ data_type: &DataType,
+ text: &str,
+ arrow_field: &arrow_schema::Field,
+) -> std::result::Result<ArrayRef, arrow_schema::ArrowError> {
+ match data_type {
+ // Java's StringToStringCastRule counts characters, then pads CHAR with
+ // spaces or truncates VARCHAR to its declared length.
+ DataType::Char(typ) => {
+ let mut value = text.chars().take(typ.length()).collect::<String>();
+ value.extend(std::iter::repeat_n(
+ ' ',
+ typ.length().saturating_sub(value.chars().count()),
+ ));
+ Ok(Arc::new(StringArray::from(vec![value])))
+ }
+ DataType::VarChar(typ) => Ok(Arc::new(StringArray::from(vec![text
+ .chars()
+ .take(typ.length() as usize)
+ .collect::<String>()]))),
+ // Java's StringToBinaryCastRule uses UTF-8 bytes and pads only BINARY.
+ DataType::Binary(typ) => {
+ let mut value = text.as_bytes().to_vec();
+ value.resize(typ.length(), 0);
+ value.truncate(typ.length());
+ Ok(Arc::new(BinaryArray::from(vec![value.as_slice()])))
+ }
+ DataType::VarBinary(typ) => {
+ let bytes = text.as_bytes();
+ let bytes = &bytes[..bytes.len().min(typ.length() as usize)];
+ Ok(Arc::new(BinaryArray::from(vec![bytes])))
+ }
+ // Java BinaryStringUtils.toByte/toShort/toInt/toLong accept a decimal
+ // suffix, validate its digits and truncate it toward zero.
+ DataType::TinyInt(_) => Ok(Arc::new(Int8Array::from(vec![parse_java_integer(
+ text,
+ i64::from(i8::MIN),
+ i64::from(i8::MAX),
+ )? as i8]))),
+ DataType::SmallInt(_) => Ok(Arc::new(Int16Array::from(vec![parse_java_integer(
+ text,
+ i64::from(i16::MIN),
+ i64::from(i16::MAX),
+ )? as i16]))),
+ DataType::Int(_) => Ok(Arc::new(Int32Array::from(vec![parse_java_integer(
+ text,
+ i64::from(i32::MIN),
+ i64::from(i32::MAX),
+ )? as i32]))),
+ DataType::BigInt(_) => Ok(Arc::new(Int64Array::from(vec![parse_java_integer(
+ text,
+ i64::MIN,
+ i64::MAX,
+ )?]))),
+ // Java's BinaryStringUtils treats an all-digit DATE/TIME value as the
+ // internal day/millisecond count and a TIMESTAMP value as the count in
+ // the requested precision, rather than parsing it as a calendar string.
+ DataType::Date(_) if numeric_default(text) => Ok(Arc::new(Date32Array::from(vec![text
+ .parse::<i32>()
+ .map_err(|source| arrow_schema::ArrowError::CastError(source.to_string()))?]))),
+ DataType::Time(_) if numeric_default(text) => {
+ Ok(Arc::new(Time32MillisecondArray::from(vec![text
+ .parse::<i32>()
+ .map_err(|source| {
+ arrow_schema::ArrowError::CastError(source.to_string())
+ })?])))
+ }
+ DataType::Timestamp(typ) if numeric_default(text) => {
+ let value = text
+ .parse::<i64>()
+ .map_err(|source| arrow_schema::ArrowError::CastError(source.to_string()))?;
+ match typ.precision() {
+ 0 => Ok(Arc::new(TimestampSecondArray::from(vec![value]))),
+ 3 => Ok(Arc::new(TimestampMillisecondArray::from(vec![value]))),
+ 6 => Ok(Arc::new(TimestampMicrosecondArray::from(vec![value]))),
+ 9 => Ok(Arc::new(TimestampNanosecondArray::from(vec![value]))),
+ precision => Err(arrow_schema::ArrowError::CastError(format!(
+ "Java does not support a numeric TIMESTAMP default at precision {precision}"
+ ))),
+ }
+ }
+ DataType::Timestamp(typ) => {
+ // Arrow stores timestamps at 0, 3, 6 or 9 digits, while Java's
+ // DateTimeUtils.parseTimestampData truncates to the declared
+ // precision before creating the internal timestamp.
+ let value = arrow_cast::cast(&StringArray::from(vec![text]), arrow_field.data_type())?;
+ if value.is_null(0) {
+ return Ok(value);
+ }
+ let precision = typ.precision();
+ let stored_precision = match precision {
+ 0 => 0,
+ 1..=3 => 3,
+ 4..=6 => 6,
+ 7..=9 => 9,
+ _ => {
+ return Err(arrow_schema::ArrowError::CastError(format!(
+ "Unsupported timestamp precision {precision}"
+ )))
+ }
+ };
+ let divisor = 10_i64.pow(stored_precision - precision);
+ if divisor == 1 {
+ return Ok(value);
+ }
+ let raw = match stored_precision {
+ 3 => value
+ .as_any()
+ .downcast_ref::<TimestampMillisecondArray>()
+ .unwrap()
+ .value(0),
+ 6 => value
+ .as_any()
+ .downcast_ref::<TimestampMicrosecondArray>()
+ .unwrap()
+ .value(0),
+ 9 => value
+ .as_any()
+ .downcast_ref::<TimestampNanosecondArray>()
+ .unwrap()
+ .value(0),
+ _ => unreachable!(),
+ };
+ let truncated = raw.div_euclid(divisor) * divisor;
+ match stored_precision {
+ 3 => Ok(Arc::new(TimestampMillisecondArray::from(vec![truncated]))),
+ 6 => Ok(Arc::new(TimestampMicrosecondArray::from(vec![truncated]))),
+ 9 => Ok(Arc::new(TimestampNanosecondArray::from(vec![truncated]))),
+ _ => unreachable!(),
+ }
+ }
+ DataType::LocalZonedTimestamp(typ) => {
+ let datetime = arrow_cast::parse::string_to_datetime(&chrono::Local, text)?;
+ let precision = typ.precision();
+ let quantum = 10_u32.pow(9 - precision);
+ let nanos = datetime.timestamp_subsec_nanos() / quantum * quantum;
+ let (units_per_second, fraction) = match precision {
+ 0 => (1_i128, 0_i128),
+ 1..=3 => (1_000, i128::from(nanos / 1_000_000)),
+ 4..=6 => (1_000_000, i128::from(nanos / 1_000)),
+ 7..=9 => (1_000_000_000, i128::from(nanos)),
+ _ => {
+ return Err(arrow_schema::ArrowError::CastError(format!(
+ "Unsupported local timestamp precision {precision}"
+ )))
+ }
+ };
+ let value =
+ i64::try_from(i128::from(datetime.timestamp()) * units_per_second + fraction)
+ .map_err(|source| arrow_schema::ArrowError::CastError(source.to_string()))?;
+ match precision {
+ 0 => Ok(Arc::new(
+ TimestampSecondArray::from(vec![value]).with_timezone("UTC"),
+ )),
+ 1..=3 => Ok(Arc::new(
+ TimestampMillisecondArray::from(vec![value]).with_timezone("UTC"),
+ )),
+ 4..=6 => Ok(Arc::new(
+ TimestampMicrosecondArray::from(vec![value]).with_timezone("UTC"),
+ )),
+ _ => Ok(Arc::new(
+ TimestampNanosecondArray::from(vec![value]).with_timezone("UTC"),
+ )),
+ }
+ }
+ DataType::Array(typ) => {
+ let ArrowDataType::List(element_field) = arrow_field.data_type() else {
+ return Err(cast_error("ARRAY default has a non-list Arrow field"));
+ };
+ let content = literal_content(text, '[', ']', "ARRAY")?;
+ let elements = split_tokens(content)
+ .iter()
+ .map(|token| cast_token(typ.element_type(), token, element_field))
+ .collect::<std::result::Result<Vec<_>, _>>()?;
+ let values = concat_default_values(&elements, element_field.data_type())?;
+ Ok(Arc::new(ListArray::try_new(
+ Arc::clone(element_field),
+ single_offset(elements.len())?,
+ values,
+ None,
+ )?))
+ }
+ DataType::Row(typ) => {
+ let ArrowDataType::Struct(arrow_fields) = arrow_field.data_type() else {
+ return Err(cast_error("ROW default has a non-struct Arrow field"));
+ };
+ let content = literal_content(text, '{', '}', "STRUCT")?;
+ let tokens = split_tokens(content);
+ if !content.is_empty() && tokens.len() != typ.fields().len() {
+ return Err(cast_error(format!(
+ "ROW default has {} fields, expected {}",
+ tokens.len(),
+ typ.fields().len()
+ )));
+ }
+ let values = typ
+ .fields()
+ .iter()
+ .enumerate()
+ .map(|(index, field)| match tokens.get(index) {
+ Some(token) => cast_token(field.data_type(), token, &arrow_fields[index]),
+ None => Ok(arrow_array::new_null_array(
+ arrow_fields[index].data_type(),
+ 1,
+ )),
+ })
+ .collect::<std::result::Result<Vec<_>, _>>()?;
+ Ok(Arc::new(StructArray::try_new(
+ arrow_fields.clone(),
+ values,
+ None,
+ )?))
+ }
+ DataType::Map(typ) => {
+ let ArrowDataType::Map(entries_field, sorted) = arrow_field.data_type() else {
+ return Err(cast_error("MAP default has a non-map Arrow field"));
+ };
+ let ArrowDataType::Struct(entry_fields) = entries_field.data_type() else {
+ return Err(cast_error("MAP default entries are not a struct"));
+ };
+ let trimmed = text.trim();
+ let pairs: Vec<(DefaultToken, DefaultToken)> = if trimmed.starts_with('{') {
+ let content = literal_content(trimmed, '{', '}', "MAP")?;
+ split_raw(content, ",", 0)
+ .into_iter()
+ .filter(|entry| !entry.is_empty())
+ .map(|entry| {
+ let pair = split_raw(&entry, "->", 2);
+ if pair.len() != 2 {
+ return Err(cast_error(format!("Invalid MAP entry: {entry}")));
+ }
+ Ok((single_token(&pair[0])?, single_token(&pair[1])?))
+ })
+ .collect::<std::result::Result<Vec<_>, _>>()?
+ } else {
+ let content = literal_content(trimmed, '{', '}', "MAP")?;
+ let tokens = split_tokens(content);
+ if !tokens.len().is_multiple_of(2) {
+ return Err(cast_error("MAP default has an odd number of tokens"));
+ }
+ tokens
+ .as_chunks::<2>()
+ .0
+ .iter()
+ .map(|pair| (pair[0].clone(), pair[1].clone()))
+ .collect()
+ };
+ let mut keys: Vec<ArrayRef> = Vec::with_capacity(pairs.len());
+ let mut values: Vec<ArrayRef> = Vec::with_capacity(pairs.len());
+ for (key_token, value_token) in &pairs {
+ let key = cast_token(typ.key_type(), key_token, &entry_fields[0])?;
+ let value = cast_token(typ.value_type(), value_token, &entry_fields[1])?;
+ // Java builds a HashMap, so a later entry replaces an earlier
+ // value with the same cast key.
+ if let Some(index) = keys.iter().position(|old| old.to_data() == key.to_data()) {
+ values[index] = value;
+ } else {
+ keys.push(key);
+ values.push(value);
+ }
+ }
+ let count = keys.len();
+ let entries = StructArray::try_new(
+ entry_fields.clone(),
+ vec![
+ concat_default_values(&keys, entry_fields[0].data_type())?,
+ concat_default_values(&values, entry_fields[1].data_type())?,
+ ],
+ None,
+ )?;
+ Ok(Arc::new(MapArray::try_new(
+ Arc::clone(entries_field),
+ single_offset(count)?,
+ entries,
+ None,
+ *sorted,
+ )?))
+ }
+ DataType::Multiset(_) => Err(cast_error(
+ "Java does not support casting a string default to MULTISET",
+ )),
+ DataType::Vector(_) => Err(cast_error(
+ "Java does not support casting a string default to VECTOR",
+ )),
+ DataType::Variant(_) => Err(cast_error(
+ "Java does not support casting a string default to VARIANT",
+ )),
+ DataType::Blob(_) => Err(arrow_schema::ArrowError::CastError(
+ "Java does not support casting a string default to BLOB".into(),
+ )),
+ _ => arrow_cast::cast(&StringArray::from(vec![text]), arrow_field.data_type()),
+ }
+}
+
+fn numeric_default(text: &str) -> bool {
+ !text.is_empty() && text.bytes().all(|byte| byte.is_ascii_digit())
+}
+
+fn parse_java_integer(text: &str, min: i64, max: i64) -> std::result::Result<i64, ArrowError> {
+ let bytes = text.as_bytes();
+ if bytes.is_empty() {
+ return Err(cast_error("Empty integer default"));
+ }
+ let negative = bytes[0] == b'-';
+ let mut index = usize::from(negative || bytes[0] == b'+');
+ if index == bytes.len() {
+ return Err(cast_error(format!("Invalid integer default: {text}")));
+ }
+
+ // Accumulate a negative value so the signed minimum remains representable.
+ // Like Java, the digits before '.' may be empty; the fractional part does
+ // not affect the result but must contain only decimal digits.
+ let limit = if negative { min } else { -max };
+ let mut value = 0_i64;
+ while index < bytes.len() && bytes[index] != b'.' {
+ let digit = bytes[index];
+ if !digit.is_ascii_digit() {
+ return Err(cast_error(format!("Invalid integer default: {text}")));
+ }
+ value = value
+ .checked_mul(10)
+ .and_then(|value| value.checked_sub(i64::from(digit - b'0')))
+ .filter(|value| *value >= limit)
+ .ok_or_else(|| cast_error(format!("Integer default overflow: {text}")))?;
+ index += 1;
+ }
+ if index < bytes.len() && !bytes[index + 1..].iter().all(u8::is_ascii_digit) {
+ return Err(cast_error(format!("Invalid integer default: {text}")));
+ }
+ Ok(if negative { value } else { -value })
+}
+
+#[derive(Clone)]
+struct DefaultToken {
+ value: String,
+ literal: bool,
+}
+
+fn cast_error(message: impl Into<String>) -> ArrowError {
+ ArrowError::CastError(message.into())
+}
+
+fn cast_token(
+ data_type: &DataType,
+ token: &DefaultToken,
+ arrow_field: &ArrowField,
+) -> std::result::Result<ArrayRef, ArrowError> {
+ if !token.literal && token.value == "null" {
+ return Ok(arrow_array::new_null_array(arrow_field.data_type(), 1));
+ }
+ let value = cast_default(data_type, &token.value, arrow_field)?;
+ if value.is_null(0) {
+ return Err(cast_error(format!(
+ "Cannot cast '{}' to {data_type}",
+ token.value
+ )));
+ }
+ Ok(value)
+}
+
+fn concat_default_values(
+ values: &[ArrayRef],
+ data_type: &ArrowDataType,
+) -> std::result::Result<ArrayRef, ArrowError> {
+ if values.is_empty() {
+ return Ok(arrow_array::new_empty_array(data_type));
+ }
+ let arrays = values
+ .iter()
+ .map(|value| value.as_ref())
+ .collect::<Vec<_>>();
+ arrow_select::concat::concat(&arrays)
+}
+
+fn single_offset(count: usize) -> std::result::Result<OffsetBuffer<i32>, ArrowError> {
+ let count = i32::try_from(count).map_err(|source| cast_error(source.to_string()))?;
+ Ok(OffsetBuffer::new(ScalarBuffer::from(vec![0, count])))
+}
+
+// Java StringToArrayCastRule, StringToMapCastRule and StringToRowCastRule
+// accept both bracket syntax and SQL function syntax.
+fn literal_content<'a>(
+ text: &'a str,
+ open: char,
+ close: char,
+ function: &str,
+) -> std::result::Result<&'a str, ArrowError> {
+ let text = text.trim();
+ if let Some(content) = text.strip_prefix(open).and_then(|v| v.strip_suffix(close)) {
+ return Ok(content.trim());
+ }
+ let prefix = text.get(..function.len());
+ if prefix.is_some_and(|prefix| prefix.eq_ignore_ascii_case(function)) {
+ let tail = text[function.len()..].trim_start();
+ if let Some(content) = tail.strip_prefix('(').and_then(|v| v.strip_suffix(')')) {
+ return Ok(content.trim());
+ }
+ }
+ Err(cast_error(format!("Invalid {function} default: {text}")))
+}
+
+fn single_token(text: &str) -> std::result::Result<DefaultToken, ArrowError> {
+ let tokens = split_tokens(text);
+ if tokens.len() != 1 {
+ return Err(cast_error(format!("Invalid MAP entry token: {text}")));
+ }
+ Ok(tokens.into_iter().next().unwrap())
+}
+
+// Port of Java TokenSplitter.split. At the current level quotes and escapes
+// make a token literal; nested punctuation is retained for recursive casting.
+fn split_tokens(content: &str) -> Vec<DefaultToken> {
+ let mut tokens = Vec::new();
+ let mut current = String::new();
+ let mut brackets = Vec::new();
+ let mut in_quotes = false;
+ let mut escaped = false;
+ let mut literal = false;
+ let mut end = 0;
+ for ch in content.chars() {
+ let nested = !brackets.is_empty();
+ if escaped {
+ escaped = false;
+ current.push(ch);
+ end = current.len();
+ continue;
+ }
+ if ch == '\\' {
+ escaped = true;
+ if nested {
+ current.push(ch);
+ end = current.len();
+ } else {
+ literal = true;
+ }
+ continue;
+ }
+ if ch == '"' {
+ in_quotes = !in_quotes;
+ if nested {
+ current.push(ch);
+ end = current.len();
+ } else {
+ literal = true;
+ }
+ continue;
+ }
+ if !in_quotes {
+ if is_open_bracket(ch) {
+ brackets.push(ch);
+ } else if is_close_bracket(ch) && !brackets.is_empty() {
+ brackets.pop();
+ } else if ch == ',' && brackets.is_empty() {
+ push_token(&mut tokens, ¤t, end, literal);
+ current.clear();
+ end = 0;
+ literal = false;
+ continue;
+ } else if ch.is_whitespace() && end == 0 {
+ continue;
+ }
+ }
+ current.push(ch);
+ if in_quotes || !ch.is_whitespace() {
+ end = current.len();
+ }
+ }
+ push_token(&mut tokens, ¤t, end, literal);
+ tokens
+}
+
+fn push_token(tokens: &mut Vec<DefaultToken>, current: &str, end: usize, literal: bool) {
+ if end > 0 || literal {
+ tokens.push(DefaultToken {
+ value: current[..end].to_string(),
+ literal,
+ });
+ }
+}
+
+// Java TokenSplitter.splitRaw leaves quotes and escapes in place for its
+// second pass, where a MAP entry is split on `->`.
+fn split_raw(content: &str, delimiter: &str, limit: usize) -> Vec<String> {
+ let mut tokens = Vec::new();
+ let mut current = String::new();
+ let mut brackets = Vec::new();
+ let mut in_quotes = false;
+ let mut escaped = false;
+ let mut splits = 0;
+ let mut chars = content.chars().peekable();
+ while let Some(ch) = chars.next() {
+ if escaped {
+ escaped = false;
+ current.push(ch);
+ continue;
+ }
+ if ch == '\\' {
+ escaped = true;
+ current.push(ch);
+ continue;
+ }
+ if ch == '"' {
+ in_quotes = !in_quotes;
+ current.push(ch);
+ continue;
+ }
+ if !in_quotes {
+ if is_open_bracket(ch) {
+ brackets.push(ch);
+ } else if is_close_bracket(ch) && !brackets.is_empty() {
+ brackets.pop();
+ } else if brackets.is_empty()
+ && (limit == 0 || splits < limit - 1)
+ && ch == delimiter.chars().next().unwrap()
+ && chars
+ .clone()
+ .take(delimiter.chars().count() - 1)
+ .collect::<String>()
+ == delimiter.chars().skip(1).collect::<String>()
+ {
+ tokens.push(current.trim().to_string());
+ current.clear();
+ splits += 1;
+ for _ in 1..delimiter.chars().count() {
+ chars.next();
+ }
+ continue;
+ }
+ }
+ current.push(ch);
+ }
+ tokens.push(current.trim().to_string());
+ tokens
+}
+
+fn is_open_bracket(ch: char) -> bool {
+ matches!(ch, '[' | '{' | '(')
+}
+
+fn is_close_bracket(ch: char) -> bool {
+ matches!(ch, ']' | '}' | ')')
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::arrow::build_target_arrow_schema;
+ use crate::spec::{BigIntType, DataType, VarCharType};
+ use arrow_array::{Int64Array, StringArray};
+
+ fn fields() -> Vec<DataField> {
+ vec![
+ DataField::new(
+ 0,
+ "dt".into(),
+ DataType::VarChar(VarCharType::string_type()),
+ )
+ .with_default_value(Some("'new'".into())),
+ DataField::new(1, "id".into(), DataType::BigInt(BigIntType::new()))
+ .with_default_value(Some("42".into())),
+ DataField::new(
+ 2,
+ "note".into(),
+ DataType::VarChar(VarCharType::string_type()),
+ ),
+ ]
+ }
+
+ #[test]
+ fn replaces_only_nulls_in_partition_and_data_columns() {
+ let fields = fields();
+ let schema = build_target_arrow_schema(&fields).unwrap();
+ let defaults = FormatTableDefaults::new(&fields, &schema).unwrap();
+ let input = RecordBatch::try_new(
+ schema,
+ vec![
+ Arc::new(StringArray::from(vec![None, Some("old"), None])),
+ Arc::new(Int64Array::from(vec![Some(7), None, None])),
+ Arc::new(StringArray::from(vec![None, Some("x"), None])),
+ ],
+ )
+ .unwrap();
+
+ let actual = defaults.apply(&input).unwrap();
+ let dt = actual
+ .column(0)
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap();
+ let id = actual
+ .column(1)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap();
+ let note = actual
+ .column(2)
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap();
+ assert_eq!(
+ dt.iter().collect::<Vec<_>>(),
+ vec![Some("new"), Some("old"), Some("new")]
+ );
+ assert_eq!(
+ id.iter().collect::<Vec<_>>(),
+ vec![Some(7), Some(42), Some(42)]
+ );
+ assert_eq!(note.iter().collect::<Vec<_>>(), vec![None, Some("x"), None]);
+ assert_eq!(
+ input.column(0).null_count(),
+ 2,
+ "input must remain unchanged"
+ );
+ }
+
+ #[test]
+ fn rejects_invalid_default_at_writer_creation() {
+ let mut fields = fields();
+ fields[1] = fields[1]
+ .clone()
+ .with_default_value(Some("'not a number'".into()));
+ let schema = build_target_arrow_schema(&fields).unwrap();
+ assert!(FormatTableDefaults::new(&fields, &schema)
+ .err()
+ .unwrap()
+ .to_string()
+ .contains("Unsupported default value"));
+ }
+
+ #[test]
+ fn integer_defaults_truncate_fraction_like_java() {
+ use crate::spec::{IntType, SmallIntType, TinyIntType};
+ use arrow_array::{Int16Array, Int32Array, Int8Array};
+
+ let fields = vec![
+ DataField::new(0, "tiny".into(), DataType::TinyInt(TinyIntType::new()))
+ .with_default_value(Some("42.9".into())),
+ DataField::new(1, "small".into(), DataType::SmallInt(SmallIntType::new()))
+ .with_default_value(Some("-42.9".into())),
+ DataField::new(2, "int".into(), DataType::Int(IntType::new()))
+ .with_default_value(Some("2147483647.9".into())),
+ DataField::new(3, "big".into(), DataType::BigInt(BigIntType::new()))
+ .with_default_value(Some("-9223372036854775808.9".into())),
+ ];
+ let schema = build_target_arrow_schema(&fields).unwrap();
+ let defaults = FormatTableDefaults::new(&fields, &schema).unwrap();
+ let input = RecordBatch::try_new(
+ schema.clone(),
+ schema
+ .fields()
+ .iter()
+ .map(|field| arrow_array::new_null_array(field.data_type(), 1))
+ .collect(),
+ )
+ .unwrap();
+ let actual = defaults.apply(&input).unwrap();
+ assert_eq!(
+ actual
+ .column(0)
+ .as_any()
+ .downcast_ref::<Int8Array>()
+ .unwrap()
+ .value(0),
+ 42
+ );
+ assert_eq!(
+ actual
+ .column(1)
+ .as_any()
+ .downcast_ref::<Int16Array>()
+ .unwrap()
+ .value(0),
+ -42
+ );
+ assert_eq!(
+ actual
+ .column(2)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .value(0),
+ i32::MAX
+ );
+ assert_eq!(
+ actual
+ .column(3)
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .unwrap()
+ .value(0),
+ i64::MIN
+ );
+ }
+
+ #[test]
+ fn integer_defaults_reject_overflow_and_invalid_fraction() {
+ use crate::spec::{IntType, SmallIntType, TinyIntType};
+
+ for (data_type, text) in [
+ (DataType::TinyInt(TinyIntType::new()), "128.1"),
+ (DataType::SmallInt(SmallIntType::new()), "32768.1"),
+ (DataType::Int(IntType::new()), "2147483648.1"),
+ (DataType::BigInt(BigIntType::new()), "9223372036854775808.1"),
+ (DataType::Int(IntType::new()), "42.x"),
+ ] {
+ let fields =
+ [DataField::new(0, "number".into(), data_type)
+ .with_default_value(Some(text.into()))];
+ let schema = build_target_arrow_schema(&fields).unwrap();
+ assert!(
+ FormatTableDefaults::new(&fields, &schema).is_err(),
+ "{text}"
+ );
+ }
+ }
+
+ #[test]
+ fn no_nulls_passes_input_through() {
+ let fields = fields();
+ let schema = build_target_arrow_schema(&fields).unwrap();
+ let defaults = FormatTableDefaults::new(&fields, &schema).unwrap();
+ let input = RecordBatch::try_new(
+ schema,
+ vec![
+ Arc::new(StringArray::from(vec!["old"])),
+ Arc::new(Int64Array::from(vec![7])),
+ Arc::new(StringArray::from(vec!["x"])),
+ ],
+ )
+ .unwrap();
+ let actual = defaults.apply(&input).unwrap();
+ assert_eq!(actual, input);
+ }
+
+ #[test]
+ fn converts_java_quoted_dates_and_numeric_defaults() {
+ use crate::spec::{BooleanType, DateType, DecimalType};
+ use arrow_array::{BooleanArray, Date32Array, Decimal128Array};
+ let fields = vec![
+ DataField::new(0, "day".into(), DataType::Date(DateType::new()))
+ .with_default_value(Some("'2025-01-02'".into())),
+ DataField::new(1, "enabled".into(), DataType::Boolean(BooleanType::new()))
+ .with_default_value(Some("true".into())),
+ DataField::new(
+ 2,
+ "amount".into(),
+ DataType::Decimal(DecimalType::new(10, 2).unwrap()),
+ )
+ .with_default_value(Some("12.34".into())),
+ ];
+ let schema = build_target_arrow_schema(&fields).unwrap();
+ let defaults = FormatTableDefaults::new(&fields, &schema).unwrap();
+ let amount = Decimal128Array::from(vec![None])
+ .with_precision_and_scale(10, 2)
+ .unwrap();
+ let input = RecordBatch::try_new(
+ schema,
+ vec![
+ Arc::new(Date32Array::from(vec![None])),
+ Arc::new(BooleanArray::from(vec![None])),
+ Arc::new(amount),
+ ],
+ )
+ .unwrap();
+ let actual = defaults.apply(&input).unwrap();
+ let day = actual
+ .column(0)
+ .as_any()
+ .downcast_ref::<Date32Array>()
+ .unwrap();
+ let enabled = actual
+ .column(1)
+ .as_any()
+ .downcast_ref::<BooleanArray>()
+ .unwrap();
+ let amount = actual
+ .column(2)
+ .as_any()
+ .downcast_ref::<Decimal128Array>()
+ .unwrap();
+ assert_eq!(day.value(0), 20_090);
+ assert!(enabled.value(0));
+ assert_eq!(amount.value(0), 1234);
+ assert_eq!(input.column(0).null_count(), 1);
+ }
+
+ #[test]
+ fn character_and_binary_defaults_follow_java_length_rules() {
+ use crate::spec::{BinaryType, CharType, VarBinaryType};
+ let fields = vec![
+ DataField::new(
+ 0,
+ "fixed_text".into(),
+ DataType::Char(CharType::new(3).unwrap()),
+ )
+ .with_default_value(Some("'猫'".into())),
+ DataField::new(
+ 1,
+ "short_text".into(),
+ DataType::VarChar(VarCharType::new(2).unwrap()),
+ )
+ .with_default_value(Some("'猫狗鱼'".into())),
+ DataField::new(
+ 2,
+ "fixed_bytes".into(),
+ DataType::Binary(BinaryType::new(4).unwrap()),
+ )
+ .with_default_value(Some("'é'".into())),
+ DataField::new(
+ 3,
+ "short_bytes".into(),
+ DataType::VarBinary(VarBinaryType::new(2).unwrap()),
+ )
+ .with_default_value(Some("'猫'".into())),
+ ];
+ let schema = build_target_arrow_schema(&fields).unwrap();
+ let defaults = FormatTableDefaults::new(&fields, &schema).unwrap();
+ let input = RecordBatch::try_new(
+ schema,
+ vec![
+ Arc::new(StringArray::from(vec![None::<&str>])),
+ Arc::new(StringArray::from(vec![None::<&str>])),
+ Arc::new(BinaryArray::from(vec![None::<&[u8]>])),
+ Arc::new(BinaryArray::from(vec![None::<&[u8]>])),
+ ],
+ )
+ .unwrap();
+ let actual = defaults.apply(&input).unwrap();
+ assert_eq!(
+ actual
+ .column(0)
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap()
+ .value(0),
+ "猫 "
+ );
+ assert_eq!(
+ actual
+ .column(1)
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap()
+ .value(0),
+ "猫狗"
+ );
+ assert_eq!(
+ actual
+ .column(2)
+ .as_any()
+ .downcast_ref::<BinaryArray>()
+ .unwrap()
+ .value(0),
+ &[0xc3, 0xa9, 0, 0]
+ );
+ assert_eq!(
+ actual
+ .column(3)
+ .as_any()
+ .downcast_ref::<BinaryArray>()
+ .unwrap()
+ .value(0),
+ &[0xe7, 0x8c]
+ );
+ }
+
+ #[test]
+ fn numeric_temporal_defaults_use_java_internal_units() {
+ use crate::spec::{DateType, TimeType, TimestampType};
+ let fields = vec![
+ DataField::new(0, "day".into(), DataType::Date(DateType::new()))
+ .with_default_value(Some("42".into())),
+ DataField::new(1, "time".into(), DataType::Time(TimeType::new(3).unwrap()))
+ .with_default_value(Some("12345".into())),
+ DataField::new(
+ 2,
+ "seconds".into(),
+ DataType::Timestamp(TimestampType::new(0).unwrap()),
+ )
+ .with_default_value(Some("123".into())),
+ DataField::new(
+ 3,
+ "micros".into(),
+ DataType::Timestamp(TimestampType::new(6).unwrap()),
+ )
+ .with_default_value(Some("123456".into())),
+ ];
+ let schema = build_target_arrow_schema(&fields).unwrap();
+ let defaults = FormatTableDefaults::new(&fields, &schema).unwrap();
+ let input = RecordBatch::try_new(
+ schema,
+ vec![
+ Arc::new(Date32Array::from(vec![None])),
+ Arc::new(Time32MillisecondArray::from(vec![None])),
+ Arc::new(TimestampSecondArray::from(vec![None])),
+ Arc::new(TimestampMicrosecondArray::from(vec![None])),
+ ],
+ )
+ .unwrap();
+ let actual = defaults.apply(&input).unwrap();
+ assert_eq!(
+ actual
+ .column(0)
+ .as_any()
+ .downcast_ref::<Date32Array>()
+ .unwrap()
+ .value(0),
+ 42
+ );
+ assert_eq!(
+ actual
+ .column(1)
+ .as_any()
+ .downcast_ref::<Time32MillisecondArray>()
+ .unwrap()
+ .value(0),
+ 12345
+ );
+ assert_eq!(
+ actual
+ .column(2)
+ .as_any()
+ .downcast_ref::<TimestampSecondArray>()
+ .unwrap()
+ .value(0),
+ 123
+ );
+ assert_eq!(
+ actual
+ .column(3)
+ .as_any()
+ .downcast_ref::<TimestampMicrosecondArray>()
+ .unwrap()
+ .value(0),
+ 123456
+ );
+ }
+
+ #[test]
+ fn rejects_blob_default_like_java() {
+ use crate::spec::BlobType;
+ let fields = vec![
+ DataField::new(0, "blob".into(), DataType::Blob(BlobType::new()))
+ .with_default_value(Some("'raw bytes'".into())),
+ ];
+ let schema = build_target_arrow_schema(&fields).unwrap();
+ assert!(FormatTableDefaults::new(&fields, &schema).is_err());
+ }
+
+ #[test]
+ fn local_timestamp_default_uses_system_timezone_and_declared_precision() {
+ use crate::spec::LocalZonedTimestampType;
+ use chrono::TimeZone;
+ let fields = vec![DataField::new(
+ 0,
+ "local_time".into(),
+ DataType::LocalZonedTimestamp(LocalZonedTimestampType::new(3).unwrap()),
+ )
+ .with_default_value(Some("'2025-01-02 03:04:05.123456'".into()))];
+ let schema = build_target_arrow_schema(&fields).unwrap();
+ let defaults = FormatTableDefaults::new(&fields, &schema).unwrap();
+ let input = RecordBatch::try_new(
+ schema,
+ vec![Arc::new(
+ TimestampMillisecondArray::from(vec![None]).with_timezone("UTC"),
+ )],
+ )
+ .unwrap();
+ let actual = defaults.apply(&input).unwrap();
+ let local = chrono::NaiveDate::from_ymd_opt(2025, 1, 2)
+ .unwrap()
+ .and_hms_milli_opt(3, 4, 5, 123)
+ .unwrap();
+ let expected = chrono::Local
+ .from_local_datetime(&local)
+ .single()
+ .unwrap()
+ .timestamp_millis();
+ assert_eq!(
+ actual
+ .column(0)
+ .as_any()
+ .downcast_ref::<TimestampMillisecondArray>()
+ .unwrap()
+ .value(0),
+ expected
+ );
+ }
+
+ #[test]
+ fn array_default_parses_each_element_like_java() {
+ use crate::spec::{ArrayType, IntType};
+ use arrow_array::{Int32Array, ListArray};
+
+ let fields = vec![DataField::new(
+ 0,
+ "numbers".into(),
+ DataType::Array(ArrayType::new(DataType::Int(IntType::new()))),
+ )
+ .with_default_value(Some("'[1,2]'".into()))];
+ let schema = build_target_arrow_schema(&fields).unwrap();
+ let defaults = FormatTableDefaults::new(&fields, &schema).unwrap();
+ let empty = arrow_array::new_null_array(schema.field(0).data_type(), 1);
+ let input = RecordBatch::try_new(schema, vec![empty]).unwrap();
+ let actual = defaults.apply(&input).unwrap();
+ let values = actual
+ .column(0)
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .unwrap()
+ .value(0);
+ let values = values.as_any().downcast_ref::<Int32Array>().unwrap();
+ assert_eq!(values.iter().collect::<Vec<_>>(), vec![Some(1), Some(2)]);
+ }
+
+ #[test]
+ fn timestamp_default_truncates_to_declared_precision() {
+ use crate::spec::TimestampType;
+
+ let fields = vec![DataField::new(
+ 0,
+ "time".into(),
+ DataType::Timestamp(TimestampType::new(1).unwrap()),
+ )
+ .with_default_value(Some("'1970-01-01 00:00:00.123456'".into()))];
+ let schema = build_target_arrow_schema(&fields).unwrap();
+ let defaults = FormatTableDefaults::new(&fields, &schema).unwrap();
+ let input = RecordBatch::try_new(
+ schema,
+ vec![Arc::new(TimestampMillisecondArray::from(vec![None]))],
+ )
+ .unwrap();
+ let actual = defaults.apply(&input).unwrap();
+ assert_eq!(
+ actual
+ .column(0)
+ .as_any()
+ .downcast_ref::<TimestampMillisecondArray>()
+ .unwrap()
+ .value(0),
+ 100
+ );
+ }
+
+ #[test]
+ fn intermediate_timestamp_precisions_follow_java() {
+ use crate::spec::TimestampType;
+
+ for (precision, expected) in [
+ (1, 100_i64),
+ (2, 120),
+ (4, 123_400),
+ (5, 123_450),
+ (7, 123_456_700),
+ ] {
+ let field = DataField::new(
+ 0,
+ "time".into(),
+ DataType::Timestamp(TimestampType::new(precision).unwrap()),
+ );
+ let schema = build_target_arrow_schema(std::slice::from_ref(&field)).unwrap();
+ let value = cast_default(
+ field.data_type(),
+ "1970-01-01 00:00:00.123456789",
+ schema.field(0),
+ )
+ .unwrap();
+ let actual = match precision {
+ 1 | 2 => value
+ .as_any()
+ .downcast_ref::<TimestampMillisecondArray>()
+ .unwrap()
+ .value(0),
+ 4 | 5 => value
+ .as_any()
+ .downcast_ref::<TimestampMicrosecondArray>()
+ .unwrap()
+ .value(0),
+ _ => value
+ .as_any()
+ .downcast_ref::<TimestampNanosecondArray>()
+ .unwrap()
+ .value(0),
+ };
+ assert_eq!(actual, expected, "precision {precision}");
+ }
+ }
+
+ #[test]
+ fn nested_array_and_quoted_null_defaults_follow_java() {
+ use crate::spec::{ArrayType, IntType};
+ use arrow_array::{Int32Array, ListArray};
+
+ let nested = DataField::new(
+ 0,
+ "nested".into(),
+ DataType::Array(ArrayType::new(DataType::Array(ArrayType::new(
+ DataType::Int(IntType::new()),
+ )))),
+ )
+ .with_default_value(Some("'[[1, 2], ARRAY(3, null)]'".into()));
+ let labels = DataField::new(
+ 1,
+ "labels".into(),
+ DataType::Array(ArrayType::new(
+ DataType::VarChar(VarCharType::string_type()),
+ )),
+ )
+ .with_default_value(Some("'[null, \"null\", \"a,b\"]'".into()));
+ let fields = [nested, labels];
+ let schema = build_target_arrow_schema(&fields).unwrap();
+ let defaults = FormatTableDefaults::new(&fields, &schema).unwrap();
+ let nested = defaults.values[0]
+ .as_ref()
+ .unwrap()
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .unwrap();
+ let values = nested.value(0);
+ let values = values.as_any().downcast_ref::<ListArray>().unwrap();
+ assert_eq!(
+ values
+ .value(0)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .iter()
+ .collect::<Vec<_>>(),
+ vec![Some(1), Some(2)]
+ );
+ assert_eq!(
+ values
+ .value(1)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .iter()
+ .collect::<Vec<_>>(),
+ vec![Some(3), None]
+ );
+ let labels = defaults.values[1]
+ .as_ref()
+ .unwrap()
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .unwrap();
+ let values = labels.value(0);
+ let values = values.as_any().downcast_ref::<StringArray>().unwrap();
+ assert_eq!(
+ values.iter().collect::<Vec<_>>(),
+ vec![None, Some("null"), Some("a,b")]
+ );
+ }
+
+ #[test]
+ fn map_and_row_defaults_use_java_token_parsing() {
+ use crate::spec::{ArrayType, IntType, MapType, RowType};
+ use arrow_array::{Int32Array, MapArray, StructArray};
+
+ let fields = [
+ DataField::new(
+ 0,
+ "lookup".into(),
+ DataType::Map(MapType::new(
+ DataType::VarChar(VarCharType::string_type()),
+ DataType::Int(IntType::new()),
+ )),
+ )
+ .with_default_value(Some("'{\"a,b\" -> 1, c -> null}'".into())),
+ DataField::new(
+ 1,
+ "record".into(),
+ DataType::Row(RowType::new(vec![
+ DataField::new(
+ 10,
+ "name".into(),
+ DataType::VarChar(VarCharType::string_type()),
+ ),
+ DataField::new(
+ 11,
+ "items".into(),
+ DataType::Array(ArrayType::new(DataType::Int(IntType::new()))),
+ ),
+ ])),
+ )
+ .with_default_value(Some("'STRUCT(\"a,b\", [1,2])'".into())),
+ ];
+ let schema = build_target_arrow_schema(&fields).unwrap();
+ let defaults = FormatTableDefaults::new(&fields, &schema).unwrap();
+ let input = RecordBatch::try_new(
+ schema.clone(),
+ schema
+ .fields()
+ .iter()
+ .map(|field| arrow_array::new_null_array(field.data_type(), 1))
+ .collect(),
+ )
+ .unwrap();
+ let actual = defaults.apply(&input).unwrap();
+ let map = actual
+ .column(0)
+ .as_any()
+ .downcast_ref::<MapArray>()
+ .unwrap();
+ assert_eq!(
+ map.keys()
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap()
+ .iter()
+ .collect::<Vec<_>>(),
+ vec![Some("a,b"), Some("c")]
+ );
+ assert_eq!(
+ map.values()
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .iter()
+ .collect::<Vec<_>>(),
+ vec![Some(1), None]
+ );
+ let row = actual
+ .column(1)
+ .as_any()
+ .downcast_ref::<StructArray>()
+ .unwrap();
+ assert_eq!(
+ row.column(0)
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap()
+ .value(0),
+ "a,b"
+ );
+ let items = row.column(1).as_any().downcast_ref::<ListArray>().unwrap();
+ assert_eq!(
+ items
+ .value(0)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .iter()
+ .collect::<Vec<_>>(),
+ vec![Some(1), Some(2)]
+ );
+
+ let value =
+ cast_default(fields[0].data_type(), "MAP(a, 1, a, 2)", schema.field(0)).unwrap();
+ let map = value.as_any().downcast_ref::<MapArray>().unwrap();
+ assert_eq!(map.keys().len(), 1);
+ assert_eq!(
+ map.values()
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .value(0),
+ 2
+ );
+ }
+
+ #[test]
+ fn malformed_complex_default_fails_before_writing() {
+ use crate::spec::{ArrayType, IntType};
+ let fields = [DataField::new(
+ 0,
+ "numbers".into(),
+ DataType::Array(ArrayType::new(DataType::Int(IntType::new()))),
+ )
+ .with_default_value(Some("'[1, bad]'".into()))];
+ let schema = build_target_arrow_schema(&fields).unwrap();
+ assert!(FormatTableDefaults::new(&fields, &schema).is_err());
+ }
+}
diff --git a/crates/paimon/src/table/format_table_scan.rs b/crates/paimon/src/table/format_table_scan.rs
index 227fbc5b..edb4e7e9 100644
--- a/crates/paimon/src/table/format_table_scan.rs
+++ b/crates/paimon/src/table/format_table_scan.rs
@@ -17,19 +17,20 @@
//! Scan implementation for Java-compatible `type=format-table` metadata.
-use std::collections::{HashMap, HashSet};
+use std::collections::HashMap;
use super::format_partition::{
format_partition_value, is_storage_not_found, parse_format_partition_value,
FormatTablePartitionPaths,
};
+use super::format_partition_location::FormatPartitionLocations;
use super::{Plan, RESTEnv, ScanTrace, Table};
use crate::api::RestError;
use crate::spec::stats::BinaryTableStats;
use crate::spec::{
escape_path_name, extract_datum, unescape_path_name, BinaryRow, BinaryRowBuilder, CoreOptions,
DataField, DataFileMeta, DataType, Datum, Partition, PartitionComputer, Predicate,
- PredicateOperator, PATH_OPTION,
+ PredicateOperator,
};
use crate::table::partition_filter::PartitionFilter;
use crate::table::source::{DataSplitBuilder, RowRange};
@@ -105,9 +106,19 @@ impl<'a> FormatTableScan<'a> {
let listed: Vec<Vec<crate::DataSplit>> = futures::stream::iter(scan_roots)
.map(|scan_root| async move {
let root_segments = path_segments(&scan_root.path);
- let partition_levels_below_root = partition_fields
- .len()
- .saturating_sub(root_segments.len().saturating_sub(table_depth));
+ // Catalog locations may live at any filesystem depth. A root
+ // with a complete catalog partition has no partition levels
+ // below it even when its path is shorter than the table path.
+ let partition_levels_below_root = if scan_root.partition.arity()
+ == partition_fields.len() as i32
+ && !scan_root.partition.is_empty()
+ {
+ 0
+ } else {
+ partition_fields
+ .len()
+ .saturating_sub(root_segments.len().saturating_sub(table_depth))
+ };
let files = list_format_table_data_files(
self.table.file_io(),
&scan_root.path,
@@ -275,16 +286,13 @@ impl<'a> FormatTableScan<'a> {
let partitions = self
.list_catalog_partitions(rest_env, pattern.as_deref(), filter.as_deref())
.await?;
- let mut seen_paths = HashSet::with_capacity(partitions.len());
+ let mut locations = FormatPartitionLocations::new(table_path)?;
let mut roots = Vec::with_capacity(partitions.len());
for partition in partitions {
let partition_path = partition_paths
.relative_path(&partition.spec)
.map_err(|error| self.invalid_catalog_partition_metadata(error))?;
- if !seen_paths.insert(partition_path.clone()) {
- continue;
- }
- let path = join_path(table_path, &partition_path);
+ let default_path = join_path(table_path, &partition_path);
let row = partition_row_from_catalog_spec(
&partition.spec,
partition_fields,
@@ -295,26 +303,9 @@ impl<'a> FormatTableScan<'a> {
if !self.partition_matches(&row)? {
continue;
}
- // An overwrite explicitly rebinds a formerly custom partition to
- // its default directory. That path option is safe to read here;
- // other custom locations still need their own resolver.
- if partition
- .options
- .as_ref()
- .and_then(|options| options.get(PATH_OPTION))
- .is_some_and(|location| {
- location.trim_end_matches('/') != path.trim_end_matches('/')
- })
- {
- return Err(crate::Error::Unsupported {
- message: format!(
- "Partition {:?} of Format Table {} is registered at a custom location, \
- which the Rust reader does not support yet",
- partition.spec,
- self.table.identifier().full_name()
- ),
- });
- }
+ let Some(path) = locations.resolve(&partition, &default_path)? else {
+ continue;
+ };
roots.push(ScanRoot {
path,
partition: row,
diff --git a/crates/paimon/src/table/format_table_write_tests.rs b/crates/paimon/src/table/format_table_write_tests.rs
index 9b4bd883..f941279b 100644
--- a/crates/paimon/src/table/format_table_write_tests.rs
+++ b/crates/paimon/src/table/format_table_write_tests.rs
@@ -229,8 +229,8 @@ async fn format_table_ignores_paimon_bucket_configuration() {
assert_eq!(ids(&table).await, [1]);
}
-#[test]
-fn declared_column_default_is_rejected_before_opening_a_writer() {
+#[tokio::test]
+async fn declared_column_default_routes_null_partition_to_default_directory() {
let table = memory_table("format_column_default", true, &[]);
let mut schema = serde_json::to_value(table.schema()).unwrap();
schema["fields"][0]["defaultValue"] = serde_json::json!("'fallback'");
@@ -242,8 +242,21 @@ fn declared_column_default_is_rejected_before_opening_a_writer() {
schema,
None,
);
- let error = table.new_write_builder().new_write().err().unwrap();
- assert!(error.to_string().contains("column default"));
+ let batch = RecordBatch::try_new(
+ Arc::new(ArrowSchema::new(vec![
+ Field::new("dt", ArrowType::Utf8, true),
+ Field::new("id", ArrowType::Int32, true),
+ ])),
+ vec![
+ Arc::new(StringArray::from(vec![None, Some("given")])),
+ Arc::new(Int32Array::from(vec![Some(1), Some(2)])),
+ ],
+ )
+ .unwrap();
+ append(&table, &batch).await;
+ assert_eq!(ids(&table).await, [1, 2]);
+ assert_eq!(visible_files(&table, "dt=fallback").await.len(), 1);
+ assert_eq!(visible_files(&table, "dt=given").await.len(), 1);
}
#[tokio::test]
@@ -700,7 +713,7 @@ fn unsupported_format_is_rejected_before_a_writer_is_opened() {
#[test]
fn readable_but_unwritable_formats_are_rejected_before_staging() {
- for format in ["avro", "orc", "mosaic"] {
+ for format in ["orc", "mosaic"] {
let table = memory_table(
&format!("format_no_{format}_writer"),
false,
@@ -715,3 +728,12 @@ fn readable_but_unwritable_formats_are_rejected_before_staging() {
);
}
}
+
+#[tokio::test]
+async fn avro_format_table_append_is_readable_through_native_scan() {
+ let table = memory_table("format_avro_append", true, &[("file.format", "avro")]);
+ append(&table, &batch(&[("a", 1), ("a", 2), ("b", 3)])).await;
+ assert_eq!(ids(&table).await, [1, 2, 3]);
+ assert_eq!(visible_files(&table, "dt=a").await.len(), 1);
+ assert_eq!(visible_files(&table, "dt=b").await.len(), 1);
+}
diff --git a/crates/paimon/src/table/format_table_writer.rs b/crates/paimon/src/table/format_table_writer.rs
index 8490a6f9..61b36863 100644
--- a/crates/paimon/src/table/format_table_writer.rs
+++ b/crates/paimon/src/table/format_table_writer.rs
@@ -27,6 +27,7 @@ use arrow_schema::{Field, Schema as ArrowSchema};
use super::commit_message::{CommitMessage, FormatFileCommit};
use super::format_partition::FormatTablePartitionPaths;
+use super::format_table_defaults::FormatTableDefaults;
use super::format_table_scan::supported_format_table_extension;
use super::table_write::take_rows;
use super::Table;
@@ -56,6 +57,7 @@ struct PartitionWriter {
pub(crate) struct FormatTableWriter {
table: Table,
full_schema: Arc<ArrowSchema>,
+ defaults: FormatTableDefaults,
file_schema: Arc<ArrowSchema>,
file_fields: Vec<DataField>,
partition_indices: Vec<usize>,
@@ -84,18 +86,6 @@ impl FormatTableWriter {
pub(crate) fn new(table: &Table, resources: Option<ResourceContext>) -> Result<Self> {
table.ensure_not_branch_reference_for_write()?;
let schema = table.schema();
- if let Some(field) = schema
- .fields()
- .iter()
- .find(|field| field.default_value().is_some())
- {
- return Err(crate::Error::Unsupported {
- message: format!(
- "Format Table column default for '{}' is not supported by the Rust writer",
- field.name()
- ),
- });
- }
let options = CoreOptions::new(schema.options());
let format = options.file_format();
let extension = supported_format_table_extension(&format)?.to_string();
@@ -103,7 +93,7 @@ impl FormatTableWriter {
// unsupported formats before staging any files, rather than failing
// after the first RecordBatch has been routed.
match format.as_str() {
- "parquet" | "row" => {}
+ "parquet" | "row" | "avro" => {}
#[cfg(feature = "vortex")]
"vortex" => {}
_ => {
@@ -115,6 +105,7 @@ impl FormatTableWriter {
}
}
let full_schema = build_target_arrow_schema(schema.fields())?;
+ let defaults = FormatTableDefaults::new(schema.fields(), &full_schema)?;
let partition_indices = schema
.partition_keys()
.iter()
@@ -172,6 +163,7 @@ impl FormatTableWriter {
Ok(Self {
table: table.clone(),
full_schema,
+ defaults,
file_schema,
file_fields,
partition_indices,
@@ -228,6 +220,8 @@ impl FormatTableWriter {
}
}
+ let batch = self.defaults.apply(batch)?;
+
let mut groups: HashMap<Vec<u8>, Vec<usize>> = HashMap::new();
let mut specs = HashMap::new();
let partition_fields = self.table.schema().partition_fields();
@@ -240,7 +234,7 @@ impl FormatTableWriter {
)?;
for row_index in 0..batch.num_rows() {
let partition = BinaryRow::from_arrow(
- batch,
+ &batch,
row_index,
&self.partition_indices,
&partition_fields,
@@ -258,7 +252,7 @@ impl FormatTableWriter {
}
for (key, rows) in groups {
- let input = take_rows(batch, &rows)?;
+ let input = take_rows(&batch, &rows)?;
let file_batch = self.project_data_columns(&input)?;
if !self.writers.contains_key(&key) {
let spec = specs.remove(&key).unwrap();
@@ -473,3 +467,104 @@ fn format_table_compression(options: &HashMap<String, String>, format: &str) ->
.to_string()
})
}
+
+#[cfg(test)]
+mod defaults_tests {
+ use std::sync::Arc;
+
+ use apache_avro::types::Value;
+ use arrow_array::{Int64Array, StringArray};
+
+ use super::*;
+ use crate::catalog::Identifier;
+ use crate::io::FileIOBuilder;
+ use crate::spec::{BigIntType, DataType, Schema, TableSchema, VarCharType};
+
+ fn table_with_defaults(dt_default: &str, id_default: &str) -> Table {
+ let schema = Schema::builder()
+ .column("dt", DataType::VarChar(VarCharType::string_type()))
+ .column("id", DataType::BigInt(BigIntType::new()))
+ .partition_keys(["dt"])
+ .option("type", "format-table")
+ .option("file.format", "avro")
+ .build()
+ .unwrap();
+ let mut json = serde_json::to_value(&schema).unwrap();
+ json["fields"][0]["defaultValue"] = dt_default.into();
+ json["fields"][1]["defaultValue"] = id_default.into();
+ let schema: Schema = serde_json::from_value(json).unwrap();
+ Table::new(
+ FileIOBuilder::new("memory").build().unwrap(),
+ Identifier::new("default", "format_defaults"),
+ "memory:/format_defaults".into(),
+ TableSchema::new(0, &schema),
+ None,
+ )
+ }
+
+ #[tokio::test]
+ async fn defaults_route_rows_and_are_written_to_avro_data_file() {
+ let table = table_with_defaults("'fallback'", "42");
+ let mut writer = FormatTableWriter::new(&table, None).unwrap();
+ let input = RecordBatch::try_new(
+ writer.full_schema.clone(),
+ vec![
+ Arc::new(StringArray::from(vec![None, Some("given")])),
+ Arc::new(Int64Array::from(vec![None, Some(7)])),
+ ],
+ )
+ .unwrap();
+ writer.write(&input).await.unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ assert_eq!(messages.len(), 2);
+ for message in messages {
+ let file = message.format_file.unwrap();
+ let dt = file.partition.get("dt").unwrap();
+ let expected = if dt == "fallback" { 42 } else { 7 };
+ assert_eq!(file.record_count, 1);
+ assert!(file.target_path.contains(&format!("dt={dt}/")));
+ let bytes = table
+ .file_io()
+ .new_input(&file.staged_path)
+ .unwrap()
+ .read()
+ .await
+ .unwrap();
+ let rows = apache_avro::Reader::new(bytes.as_ref())
+ .unwrap()
+ .collect::<std::result::Result<Vec<_>, _>>()
+ .unwrap();
+ assert_eq!(rows.len(), 1);
+ let Value::Record(columns) = &rows[0] else {
+ panic!("expected Avro record");
+ };
+ assert_eq!(columns[0].0, "id");
+ assert_eq!(
+ columns[0].1,
+ Value::Union(1, Box::new(Value::Long(expected)))
+ );
+ table
+ .file_io()
+ .delete_file(&file.staged_path)
+ .await
+ .unwrap();
+ }
+ assert_eq!(input.column(0).null_count(), 1);
+ assert_eq!(input.column(1).null_count(), 1);
+ }
+
+ #[tokio::test]
+ async fn invalid_default_fails_before_any_file_is_staged() {
+ let table = table_with_defaults("'fallback'", "bad-integer");
+ let error = FormatTableWriter::new(&table, None).err().unwrap();
+ assert!(error.to_string().contains("Unsupported default value"));
+ let staging = "memory:/format_defaults/_temporary";
+ assert!(!table
+ .file_io()
+ .new_input(staging)
+ .unwrap()
+ .exists()
+ .await
+ .unwrap());
+ }
+}
diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs
index 231fdd86..329bc1df 100644
--- a/crates/paimon/src/table/mod.rs
+++ b/crates/paimon/src/table/mod.rs
@@ -46,10 +46,12 @@ mod de_vector_read;
mod de_vector_scan;
mod dedicated_format_file_writer;
mod format_partition;
+mod format_partition_location;
mod format_partition_stats;
mod format_partition_truncate;
mod format_read_builder;
mod format_table_commit;
+mod format_table_defaults;
mod format_table_read;
mod format_table_scan;
#[cfg(test)]