risingwave_common/array/arrow/
arrow_iceberg.rs

1// Copyright 2024 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::cell::RefCell;
16use std::ops::Div;
17use std::sync::{Arc, LazyLock};
18
19use arrow_array::ArrayRef;
20use arrow_array::cast::AsArray;
21use arrow_schema::extension::ExtensionType;
22use num_traits::abs;
23use parquet_variant_compute::{VariantArray as ParquetVariantArray, VariantType};
24use risingwave_common_log::LogSuppressor;
25use thiserror_ext::AsReport;
26
27pub use super::arrow_58::{
28    FromArrow, ToArrow, arrow_array, arrow_buffer, arrow_cast, arrow_schema,
29    is_parquet_field_match_source_schema, is_parquet_schema_match_source_schema,
30};
31use crate::array::{
32    Array, ArrayBuilder, ArrayError, ArrayImpl, DataChunk, DataType, DecimalArray, IntervalArray,
33    VariantArray as RwVariantArray, VariantArrayBuilder,
34};
35use crate::types::{Scalar, StructType, VariantVal};
36
37pub struct IcebergArrowConvert;
38
39struct DefaultIcebergFromArrow;
40
41impl FromArrow for DefaultIcebergFromArrow {}
42
43// Arrow Decimal128 supports up to 38 decimal digits. We use precision=38, scale=10:
44// - Integer range: up to 10^28 - 1 (28 digits)
45// - Fractional precision: 10 digits
46// - Covers all RisingWave decimal values (MAX_PRECISION=28)
47//
48// Note: When reading Arrow decimals that exceed RisingWave's 96-bit / 28-digit
49// storage limit, the conversion code in arrow_impl.rs will reduce scale and
50// truncate the mantissa (via truncated_i128_and_scale) to make them fit.
51pub const ICEBERG_DECIMAL_PRECISION: u8 = 38;
52pub const ICEBERG_DECIMAL_SCALE: i8 = 10;
53
54impl IcebergArrowConvert {
55    pub fn to_record_batch(
56        &self,
57        schema: arrow_schema::SchemaRef,
58        chunk: &DataChunk,
59    ) -> Result<arrow_array::RecordBatch, ArrayError> {
60        ToArrow::to_record_batch(self, schema, chunk)
61    }
62
63    pub fn chunk_from_record_batch(
64        &self,
65        batch: &arrow_array::RecordBatch,
66    ) -> Result<DataChunk, ArrayError> {
67        FromArrow::from_record_batch(self, batch)
68    }
69
70    pub fn type_from_field(&self, field: &arrow_schema::Field) -> Result<DataType, ArrayError> {
71        FromArrow::from_field(self, field)
72    }
73
74    pub fn to_arrow_field(
75        &self,
76        name: &str,
77        data_type: &DataType,
78    ) -> Result<arrow_schema::Field, ArrayError> {
79        ToArrow::to_arrow_field(self, name, data_type)
80    }
81
82    pub fn struct_from_fields(
83        &self,
84        fields: &arrow_schema::Fields,
85    ) -> Result<StructType, ArrayError> {
86        FromArrow::from_fields(self, fields)
87    }
88
89    pub fn to_arrow_array(
90        &self,
91        data_type: &arrow_schema::DataType,
92        array: &ArrayImpl,
93    ) -> Result<arrow_array::ArrayRef, ArrayError> {
94        ToArrow::to_array(self, data_type, array)
95    }
96
97    pub fn array_from_arrow_array(
98        &self,
99        field: &arrow_schema::Field,
100        array: &arrow_array::ArrayRef,
101    ) -> Result<ArrayImpl, ArrayError> {
102        FromArrow::from_array(self, field, array)
103    }
104
105    /// A helper function to convert an Arrow array to RisingWave array without knowing the field.
106    /// It will use the datatype from arrow array to infer the RisingWave data type.
107    ///
108    /// The difference between this function and `array_from_arrow_array` is that `array_from_arrow_array` will try using `ARROW:extension:name` field metadata to determine the RisingWave data type for extension types.
109    pub fn array_from_arrow_array_raw(
110        &self,
111        array: &arrow_array::ArrayRef,
112    ) -> Result<ArrayImpl, ArrayError> {
113        static FIELD_DUMMY: LazyLock<arrow_schema::Field> =
114            LazyLock::new(|| arrow_schema::Field::new("dummy", arrow_schema::DataType::Null, true));
115        FromArrow::from_array(self, &FIELD_DUMMY, array)
116    }
117}
118
119impl ToArrow for IcebergArrowConvert {
120    fn to_arrow_field(
121        &self,
122        name: &str,
123        data_type: &DataType,
124    ) -> Result<arrow_schema::Field, ArrayError> {
125        let data_type = match data_type {
126            DataType::Boolean => self.bool_type_to_arrow(),
127            DataType::Int16 => self.int32_type_to_arrow(),
128            DataType::Int32 => self.int32_type_to_arrow(),
129            DataType::Int64 => self.int64_type_to_arrow(),
130            DataType::Int256 => self.int256_type_to_arrow(),
131            DataType::Float32 => self.float32_type_to_arrow(),
132            DataType::Float64 => self.float64_type_to_arrow(),
133            DataType::Date => self.date_type_to_arrow(),
134            DataType::Time => self.time_type_to_arrow(),
135            DataType::Timestamp => self.timestamp_type_to_arrow(),
136            DataType::Timestamptz => self.timestamptz_type_to_arrow(),
137            DataType::Interval => self.interval_type_to_arrow(),
138            DataType::Varchar => self.varchar_type_to_arrow(),
139            DataType::Bytea => self.bytea_type_to_arrow(),
140            DataType::Serial => self.serial_type_to_arrow(),
141            DataType::Decimal => return Ok(self.decimal_type_to_arrow(name)),
142            DataType::Jsonb => self.varchar_type_to_arrow(),
143            DataType::Variant => return Ok(variant_arrow_field(name)),
144            DataType::Struct(fields) => self.struct_type_to_arrow(fields)?,
145            DataType::List(list) => self.list_type_to_arrow(list)?,
146            DataType::Map(map) => self.map_type_to_arrow(map)?,
147            DataType::Vector(_) => self.vector_type_to_arrow()?,
148        };
149        Ok(arrow_schema::Field::new(name, data_type, true))
150    }
151
152    #[inline]
153    fn interval_type_to_arrow(&self) -> arrow_schema::DataType {
154        arrow_schema::DataType::Utf8
155    }
156
157    #[inline]
158    fn decimal_type_to_arrow(&self, name: &str) -> arrow_schema::Field {
159        // Fixed-point decimal; precision P, scale S Scale is fixed, precision must be less than 38.
160        let data_type =
161            arrow_schema::DataType::Decimal128(ICEBERG_DECIMAL_PRECISION, ICEBERG_DECIMAL_SCALE);
162        arrow_schema::Field::new(name, data_type, true)
163    }
164
165    fn decimal_to_arrow(
166        &self,
167        data_type: &arrow_schema::DataType,
168        array: &DecimalArray,
169    ) -> Result<arrow_array::ArrayRef, ArrayError> {
170        let (precision, max_scale) = match data_type {
171            arrow_schema::DataType::Decimal128(precision, scale) => (*precision, *scale),
172            _ => return Err(ArrayError::to_arrow("Invalid decimal type")),
173        };
174
175        // Convert Decimal to i128:
176        let max_value = 10_i128.pow(precision as u32) - 1;
177        let values: Vec<Option<i128>> = array
178            .iter()
179            .map(|e| {
180                e.and_then(|e| match e {
181                    crate::array::Decimal::Normalized(e) => {
182                        let value = e.mantissa();
183                        let scale = e.scale() as i8;
184                        let diff_scale = abs(max_scale - scale);
185                        let value = match scale {
186                            _ if scale < max_scale => value
187                                .checked_mul(10_i128.pow(diff_scale as u32))
188                                .filter(|&v| abs(v) <= max_value)
189                                .unwrap_or_else(|| {
190                                    tracing::warn!(
191                                        "Decimal overflow when converting to arrow decimal with precision {} and scale {}. It will be replaced with inf/-inf.",
192                                        precision, max_scale
193                                    );
194                                    if value >= 0 { max_value } else { -max_value }
195                                }),
196                            _ if scale > max_scale => value.div(10_i128.pow(diff_scale as u32)),
197                            _ => value,
198                        };
199                        Some(value)
200                    }
201                    // For Inf, we replace them with the max/min value within the precision.
202                    crate::array::Decimal::PositiveInf => {
203                        Some(max_value)
204                    }
205                    crate::array::Decimal::NegativeInf => {
206                        Some(-max_value)
207                    }
208                    crate::array::Decimal::NaN => None,
209                })
210            })
211            .collect();
212
213        let array = arrow_array::Decimal128Array::from(values)
214            .with_precision_and_scale(precision, max_scale)
215            .map_err(ArrayError::from_arrow)?;
216        Ok(Arc::new(array) as ArrayRef)
217    }
218
219    fn interval_to_arrow(
220        &self,
221        array: &IntervalArray,
222    ) -> Result<arrow_array::ArrayRef, ArrayError> {
223        Ok(Arc::new(arrow_array::StringArray::from(array)))
224    }
225
226    fn variant_to_arrow(
227        &self,
228        array: &RwVariantArray,
229    ) -> Result<arrow_array::ArrayRef, ArrayError> {
230        // For a SQL NULL row the children's bytes are ignored by the parent null bitmap,
231        // so `raw_iter`'s valid variant-null placeholder keeps both child arrays non-null
232        // without a special branch per row.
233        let metadata = Arc::new(arrow_array::BinaryArray::from_iter_values(
234            array.raw_iter().map(|variant| variant.metadata()),
235        )) as ArrayRef;
236        let value = Arc::new(arrow_array::BinaryArray::from_iter_values(
237            array.raw_iter().map(|variant| variant.value()),
238        )) as ArrayRef;
239        let nulls = (!array.null_bitmap().all()).then(|| array.null_bitmap().into());
240
241        Ok(Arc::new(arrow_array::StructArray::new(
242            variant_arrow_fields(),
243            vec![metadata, value],
244            nulls,
245        )))
246    }
247}
248
249impl FromArrow for IcebergArrowConvert {
250    fn from_extension_type(
251        &self,
252        type_name: &str,
253        physical_type: &arrow_schema::DataType,
254    ) -> Result<DataType, ArrayError> {
255        match (type_name, physical_type) {
256            (VariantType::NAME, arrow_schema::DataType::Struct(_)) => Ok(DataType::Variant),
257            (VariantType::NAME, _) => Err(ArrayError::from_arrow(format!(
258                "variant extension type requires a struct physical type, got: {physical_type}"
259            ))),
260            _ => DefaultIcebergFromArrow.from_extension_type(type_name, physical_type),
261        }
262    }
263
264    fn from_extension_array(
265        &self,
266        type_name: &str,
267        array: &arrow_array::ArrayRef,
268    ) -> Result<ArrayImpl, ArrayError> {
269        match type_name {
270            VariantType::NAME => variant_array_to_variant(array),
271            _ => DefaultIcebergFromArrow.from_extension_array(type_name, array),
272        }
273    }
274}
275
276/// The Arrow field layout of an unshredded variant column, tagged with the
277/// `arrow.parquet.variant` extension.
278fn variant_arrow_field(name: &str) -> arrow_schema::Field {
279    arrow_schema::Field::new(
280        name,
281        arrow_schema::DataType::Struct(variant_arrow_fields()),
282        true,
283    )
284    .with_extension_type(VariantType)
285}
286
287fn variant_arrow_fields() -> arrow_schema::Fields {
288    [
289        Arc::new(arrow_schema::Field::new(
290            "metadata",
291            arrow_schema::DataType::Binary,
292            false,
293        )),
294        Arc::new(arrow_schema::Field::new(
295            "value",
296            arrow_schema::DataType::Binary,
297            false,
298        )),
299    ]
300    .into()
301}
302
303fn variant_array_to_variant(array: &arrow_array::ArrayRef) -> Result<ArrayImpl, ArrayError> {
304    let variant_array =
305        ParquetVariantArray::try_new(array.as_ref()).map_err(ArrayError::from_arrow)?;
306    // The shredded encoding cannot be reconstructed yet, and decoding only metadata/value
307    // would yield silently partial objects, so the whole column reads as NULL.
308    if variant_array.typed_value_field().is_some() {
309        static LOG_SUPPRESSOR: LazyLock<LogSuppressor> = LazyLock::new(LogSuppressor::default);
310        if let Ok(suppressed_count) = LOG_SUPPRESSOR.check() {
311            tracing::warn!(
312                suppressed_count,
313                "shredded variant column (with a `typed_value` field) is not supported yet; \
314                 reading it as NULL",
315            );
316        }
317        let mut builder = VariantArrayBuilder::new(variant_array.len());
318        for _ in 0..variant_array.len() {
319            builder.append_null();
320        }
321        return Ok(ArrayImpl::Variant(builder.finish()));
322    }
323    // The physical binary layout is constant across the batch, so resolve the typed
324    // accessors once here rather than re-dispatching and downcasting per row.
325    let metadata_accessor = BinaryArrayAccessor::resolve(variant_array.metadata_field())?;
326    let value_accessor = variant_array
327        .value_field()
328        .map(BinaryArrayAccessor::resolve)
329        .transpose()?;
330    let mut builder = VariantArrayBuilder::new(variant_array.len());
331
332    for idx in 0..variant_array.len() {
333        if variant_array.is_null(idx) {
334            builder.append_null();
335            continue;
336        }
337
338        // `from_parts` fully validates the untrusted bytes and re-encodes them into
339        // RW's canonical form, which the byte-wise `Eq`/`Ord` of VARIANT relies on.
340        // TODO: rows of a batch virtually always share one metadata dictionary;
341        // cache the validated metadata instead of re-validating it per row.
342        let variant_result = match &value_accessor {
343            Some(value) if value.is_valid(idx) => {
344                VariantVal::from_parts(metadata_accessor.value(idx), value.value(idx))
345            }
346            // Mirrors the upstream `VariantArray::try_value` fallback. Shredded arrays are
347            // rejected above, so reaching this arm means the file is malformed for the
348            // unshredded encoding — `value` is required there — and the row silently becomes
349            // variant null rather than SQL NULL.
350            _ => Ok(VariantVal::null()),
351        };
352
353        match variant_result {
354            Ok(variant) => builder.append(Some(variant.as_scalar_ref())),
355            Err(err) => {
356                // A systematically corrupt file would otherwise warn once per row.
357                static LOG_SUPPRESSOR: LazyLock<LogSuppressor> =
358                    LazyLock::new(LogSuppressor::default);
359                if let Ok(suppressed_count) = LOG_SUPPRESSOR.check() {
360                    tracing::warn!(
361                        error = %err.as_report(),
362                        suppressed_count,
363                        "failed to decode iceberg variant value at index {}. It will be replaced with null.",
364                        idx,
365                    );
366                }
367                builder.append_null();
368            }
369        }
370    }
371
372    Ok(ArrayImpl::Variant(builder.finish()))
373}
374
375/// A binary-like Arrow array with its physical layout resolved once per batch, so
376/// the per-row hot path just indexes the already-downcast array.
377///
378/// `VariantArray::try_new` validates the `metadata` and `value` fields as `Binary`,
379/// `LargeBinary`, or `BinaryView` before we get here, so [`resolve`](Self::resolve)
380/// covers every layout a variant field can carry and its error arm is unreachable
381/// for arrays taken from a `VariantArray`.
382enum BinaryArrayAccessor<'a> {
383    Binary(&'a arrow_array::BinaryArray),
384    LargeBinary(&'a arrow_array::LargeBinaryArray),
385    BinaryView(&'a arrow_array::BinaryViewArray),
386}
387
388impl<'a> BinaryArrayAccessor<'a> {
389    fn resolve(array: &'a ArrayRef) -> Result<Self, ArrayError> {
390        use arrow_schema::DataType;
391        match array.data_type() {
392            DataType::Binary => Ok(Self::Binary(array.as_binary::<i32>())),
393            DataType::LargeBinary => Ok(Self::LargeBinary(array.as_binary::<i64>())),
394            DataType::BinaryView => Ok(Self::BinaryView(array.as_binary_view())),
395            other => Err(ArrayError::from_arrow(format!(
396                "variant metadata/value has unsupported binary layout: {other}"
397            ))),
398        }
399    }
400
401    #[inline]
402    fn is_valid(&self, index: usize) -> bool {
403        use arrow_array::Array as _;
404        match self {
405            Self::Binary(array) => array.is_valid(index),
406            Self::LargeBinary(array) => array.is_valid(index),
407            Self::BinaryView(array) => array.is_valid(index),
408        }
409    }
410
411    #[inline]
412    fn value(&self, index: usize) -> &[u8] {
413        match self {
414            Self::Binary(array) => array.value(index),
415            Self::LargeBinary(array) => array.value(index),
416            Self::BinaryView(array) => array.value(index),
417        }
418    }
419}
420
421/// Iceberg sink with `create_table_if_not_exists` option will use this struct to convert the
422/// iceberg data type to arrow data type.
423///
424/// Specifically, it will add the field id to the arrow field metadata, because iceberg-rust need the field id to be set.
425///
426/// Note: this is different from [`IcebergArrowConvert`], which is used to read from/write to
427/// an _existing_ iceberg table. In that case, we just need to make sure the data is compatible to the existing schema.
428/// But to _create a new table_, we need to meet more requirements of iceberg.
429#[derive(Default)]
430pub struct IcebergCreateTableArrowConvert {
431    next_field_id: RefCell<u32>,
432}
433
434impl IcebergCreateTableArrowConvert {
435    pub fn to_arrow_field(
436        &self,
437        name: &str,
438        data_type: &DataType,
439    ) -> Result<arrow_schema::Field, ArrayError> {
440        ToArrow::to_arrow_field(self, name, data_type)
441    }
442
443    fn add_field_id(&self, arrow_field: &mut arrow_schema::Field) {
444        *self.next_field_id.borrow_mut() += 1;
445        let field_id = *self.next_field_id.borrow();
446
447        // Preserve extension metadata such as `arrow.parquet.variant` while adding the
448        // Iceberg field id required by `arrow_schema_to_schema`.
449        let mut metadata = arrow_field.metadata().clone();
450        // for iceberg-rust
451        metadata.insert("PARQUET:field_id".to_owned(), field_id.to_string());
452        arrow_field.set_metadata(metadata);
453    }
454}
455
456impl ToArrow for IcebergCreateTableArrowConvert {
457    #[inline]
458    fn decimal_type_to_arrow(&self, name: &str) -> arrow_schema::Field {
459        // To create a iceberg table, we need a decimal type with precision and scale to be set
460        // We choose 28 here
461        // The decimal type finally will be converted to an iceberg decimal type.
462        // Iceberg decimal(P,S)
463        // Fixed-point decimal; precision P, scale S Scale is fixed, precision must be less than 38.
464        let data_type =
465            arrow_schema::DataType::Decimal128(ICEBERG_DECIMAL_PRECISION, ICEBERG_DECIMAL_SCALE);
466
467        let mut arrow_field = arrow_schema::Field::new(name, data_type, true);
468        self.add_field_id(&mut arrow_field);
469        arrow_field
470    }
471
472    #[inline]
473    fn interval_type_to_arrow(&self) -> arrow_schema::DataType {
474        arrow_schema::DataType::Utf8
475    }
476
477    fn jsonb_type_to_arrow(&self, name: &str) -> arrow_schema::Field {
478        let data_type = arrow_schema::DataType::Utf8;
479
480        let mut arrow_field = arrow_schema::Field::new(name, data_type, true);
481        self.add_field_id(&mut arrow_field);
482        arrow_field
483    }
484
485    /// Convert RisingWave data type to Arrow data type.
486    ///
487    /// This function returns a `Field` instead of `DataType` because some may be converted to
488    /// extension types which require additional metadata in the field.
489    fn to_arrow_field(
490        &self,
491        name: &str,
492        value: &DataType,
493    ) -> Result<arrow_schema::Field, ArrayError> {
494        let data_type = match value {
495            // using the inline function
496            DataType::Boolean => self.bool_type_to_arrow(),
497            DataType::Int16 => self.int32_type_to_arrow(),
498            DataType::Int32 => self.int32_type_to_arrow(),
499            DataType::Int64 => self.int64_type_to_arrow(),
500            DataType::Int256 => self.varchar_type_to_arrow(),
501            DataType::Float32 => self.float32_type_to_arrow(),
502            DataType::Float64 => self.float64_type_to_arrow(),
503            DataType::Date => self.date_type_to_arrow(),
504            DataType::Time => self.time_type_to_arrow(),
505            DataType::Timestamp => self.timestamp_type_to_arrow(),
506            DataType::Timestamptz => self.timestamptz_type_to_arrow(),
507            DataType::Interval => self.interval_type_to_arrow(),
508            DataType::Varchar => self.varchar_type_to_arrow(),
509            DataType::Bytea => self.bytea_type_to_arrow(),
510            DataType::Serial => self.serial_type_to_arrow(),
511            DataType::Decimal => return Ok(self.decimal_type_to_arrow(name)),
512            DataType::Jsonb => self.varchar_type_to_arrow(),
513            DataType::Variant => {
514                let mut arrow_field = variant_arrow_field(name);
515                self.add_field_id(&mut arrow_field);
516                return Ok(arrow_field);
517            }
518            DataType::Struct(fields) => self.struct_type_to_arrow(fields)?,
519            DataType::List(list) => self.list_type_to_arrow(list)?,
520            DataType::Map(map) => self.map_type_to_arrow(map)?,
521            DataType::Vector(_) => self.vector_type_to_arrow()?,
522        };
523
524        let mut arrow_field = arrow_schema::Field::new(name, data_type, true);
525        self.add_field_id(&mut arrow_field);
526        Ok(arrow_field)
527    }
528}
529
530#[cfg(test)]
531mod test {
532    use std::sync::Arc;
533
534    use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime, Utc};
535    use parquet_variant::{
536        ShortString, Variant, VariantDecimal4, VariantDecimal8, VariantDecimal16,
537    };
538    use parquet_variant_compute::{VariantArrayBuilder, json_to_variant};
539    use uuid::Uuid;
540
541    use super::arrow_array::{ArrayRef, Decimal128Array};
542    use super::arrow_schema::DataType as ArrowDataType;
543    use super::*;
544    use crate::array::{Decimal, DecimalArray};
545    use crate::types::{MapType, ToText};
546
547    #[test]
548    fn decimal() {
549        let array = DecimalArray::from_iter([
550            None,
551            Some(Decimal::NaN),
552            Some(Decimal::PositiveInf),
553            Some(Decimal::NegativeInf),
554            Some(Decimal::Normalized("123.4".parse().unwrap())),
555            Some(Decimal::Normalized("123.456".parse().unwrap())),
556        ]);
557        let ty = ArrowDataType::Decimal128(6, 3);
558        let arrow_array = IcebergArrowConvert.decimal_to_arrow(&ty, &array).unwrap();
559        let expect_array = Arc::new(
560            Decimal128Array::from(vec![
561                None,
562                None,
563                Some(999999),
564                Some(-999999),
565                Some(123400),
566                Some(123456),
567            ])
568            .with_data_type(ty),
569        ) as ArrayRef;
570        assert_eq!(&arrow_array, &expect_array);
571    }
572
573    #[test]
574    fn decimal_with_large_scale() {
575        let array = DecimalArray::from_iter([
576            None,
577            Some(Decimal::NaN),
578            Some(Decimal::PositiveInf),
579            Some(Decimal::NegativeInf),
580            Some(Decimal::Normalized("123.4".parse().unwrap())),
581            Some(Decimal::Normalized("123.456".parse().unwrap())),
582        ]);
583        let ty = ArrowDataType::Decimal128(ICEBERG_DECIMAL_PRECISION, ICEBERG_DECIMAL_SCALE);
584        let arrow_array = IcebergArrowConvert.decimal_to_arrow(&ty, &array).unwrap();
585        let expect_array = Arc::new(
586            Decimal128Array::from(vec![
587                None,
588                None,
589                // With precision=38, max value is 10^38 - 1
590                Some(99999999999999999999999999999999999999),
591                Some(-99999999999999999999999999999999999999),
592                Some(1234000000000),
593                Some(1234560000000),
594            ])
595            .with_data_type(ty),
596        ) as ArrayRef;
597        assert_eq!(&arrow_array, &expect_array);
598    }
599
600    #[test]
601    fn decimal_edge_cases_risingwave_precision() {
602        // Test edge cases between RisingWave decimal precision (28 digits) and Arrow Decimal128(38,10)
603        let array = DecimalArray::from_iter([
604            // Large 27-digit integer (previously would overflow with precision=28, scale=10)
605            Some(Decimal::Normalized(
606                "999999999999999999999999999".parse().unwrap(),
607            )),
608            // RisingWave MAX_PRECISION: 28-digit integer
609            Some(Decimal::Normalized(
610                "9999999999999999999999999999".parse().unwrap(),
611            )),
612            // Large integer with fractional part
613            Some(Decimal::Normalized(
614                "999999999999999999.9999999999".parse().unwrap(),
615            )),
616            // Small value with maximum fractional digits
617            Some(Decimal::Normalized(
618                "0.9999999999999999999999999999".parse().unwrap(),
619            )),
620            // Negative large integer
621            Some(Decimal::Normalized(
622                "-999999999999999999999999999".parse().unwrap(),
623            )),
624            // Edge case: exactly 10^18 (18 digits) - boundary for old precision=28,scale=10
625            Some(Decimal::Normalized("1000000000000000000".parse().unwrap())),
626            // Very small decimal
627            Some(Decimal::Normalized("0.0000000001".parse().unwrap())),
628            // Zero with fractional representation
629            Some(Decimal::Normalized("0.0000000000".parse().unwrap())),
630        ]);
631
632        let ty = ArrowDataType::Decimal128(ICEBERG_DECIMAL_PRECISION, ICEBERG_DECIMAL_SCALE);
633        let arrow_array = IcebergArrowConvert.decimal_to_arrow(&ty, &array).unwrap();
634
635        let expect_array = Arc::new(
636            Decimal128Array::from(vec![
637                // 999999999999999999999999999 * 10^10 (scale 0 → 10)
638                Some(9999999999999999999999999990000000000),
639                // 9999999999999999999999999999 * 10^10
640                Some(99999999999999999999999999990000000000),
641                // 999999999999999999.9999999999 already at scale 10
642                Some(9999999999999999999999999999),
643                // 0.9999999999999999999999999999: scale 28 → 10, truncates to 0.9999999999
644                Some(9999999999),
645                // -999999999999999999999999999 * 10^10
646                Some(-9999999999999999999999999990000000000),
647                // 1000000000000000000 * 10^10
648                Some(10000000000000000000000000000),
649                // 0.0000000001 already at scale 10
650                Some(1),
651                // 0.0000000000 (scale 10)
652                Some(0),
653            ])
654            .with_data_type(ty),
655        ) as ArrayRef;
656
657        assert_eq!(&arrow_array, &expect_array);
658    }
659
660    #[test]
661    fn decimal_special_values_roundtrip() {
662        // Test that special decimal values (inf, -inf, nan) can be written and read back correctly
663        use crate::array::Array;
664
665        let original_array = DecimalArray::from_iter([
666            Some(Decimal::PositiveInf),
667            Some(Decimal::NegativeInf),
668            Some(Decimal::NaN),
669            Some(Decimal::Normalized("123.45".parse().unwrap())),
670            None,
671        ]);
672
673        // Convert to Arrow
674        let ty = ArrowDataType::Decimal128(ICEBERG_DECIMAL_PRECISION, ICEBERG_DECIMAL_SCALE);
675        let arrow_array = IcebergArrowConvert
676            .decimal_to_arrow(&ty, &original_array)
677            .unwrap();
678
679        // Convert back to RisingWave
680        let arrow_decimal: &arrow_array::Decimal128Array = arrow_array
681            .as_any()
682            .downcast_ref()
683            .expect("should be Decimal128Array");
684
685        let roundtrip_array: DecimalArray = arrow_decimal.try_into().unwrap();
686
687        // Verify special values roundtrip correctly
688        assert_eq!(original_array.len(), roundtrip_array.len());
689
690        // PositiveInf -> max value -> PositiveInf
691        assert_eq!(roundtrip_array.value_at(0), Some(Decimal::PositiveInf));
692
693        // NegativeInf -> min value -> NegativeInf
694        assert_eq!(roundtrip_array.value_at(1), Some(Decimal::NegativeInf));
695
696        // NaN -> NULL -> None (NaN cannot roundtrip, becomes NULL in Arrow)
697        assert_eq!(roundtrip_array.value_at(2), None);
698
699        // Normal value roundtrips correctly (scale may be adjusted)
700        assert!(matches!(
701            roundtrip_array.value_at(3),
702            Some(Decimal::Normalized(_))
703        ));
704
705        // NULL -> NULL -> None
706        assert_eq!(roundtrip_array.value_at(4), None);
707    }
708
709    #[test]
710    fn all_variant_internal_types_convert_to_variant() {
711        let long_string = "x".repeat(64);
712        let binary_bytes = [0x0a_u8, 0x0b, 0x0c, 0x0d];
713        let object_and_list_json = Arc::new(arrow_array::StringArray::from(vec![
714            Some(r#"{"a":1,"b":[true,null]}"#),
715            Some(r#"[1,{"x":2},"tail"]"#),
716        ])) as ArrayRef;
717        let object_and_list = json_to_variant(&object_and_list_json).unwrap();
718
719        let timestamp_micros = DateTime::parse_from_rfc3339("2024-11-07T12:33:54.123456+00:00")
720            .unwrap()
721            .with_timezone(&Utc);
722        let timestamp_ntz_micros =
723            NaiveDateTime::parse_from_str("2024-11-07 12:33:54.123456", "%Y-%m-%d %H:%M:%S%.f")
724                .unwrap();
725        let timestamp_nanos = DateTime::parse_from_rfc3339("2024-11-07T12:33:54.123456789+00:00")
726            .unwrap()
727            .with_timezone(&Utc);
728        let timestamp_ntz_nanos =
729            NaiveDateTime::parse_from_str("2024-11-07 12:33:54.123456789", "%Y-%m-%d %H:%M:%S%.f")
730                .unwrap();
731        let time = NaiveTime::from_hms_micro_opt(12, 33, 54, 123_456).unwrap();
732        let uuid = Uuid::parse_str("123e4567-e89b-12d3-a456-426614174000").unwrap();
733
734        let mut builder = VariantArrayBuilder::new(25);
735        builder.append_null(); // null Arrow row (not variant null)
736        builder.append_variant(Variant::Null);
737        builder.append_variant(Variant::BooleanTrue);
738        builder.append_variant(Variant::BooleanFalse);
739        builder.append_variant(Variant::Int8(34));
740        builder.append_variant(Variant::Int16(1234));
741        builder.append_variant(Variant::Int32(100_000));
742        builder.append_variant(Variant::Int64(5_000_000_000));
743        builder.append_variant(Variant::Float(3.5));
744        builder.append_variant(Variant::Double(14.25));
745        // NaN normalizes to the canonical bit pattern and renders as the string "NaN".
746        builder.append_variant(Variant::Double(f64::NAN));
747        builder.append_variant(Variant::Decimal4(
748            VariantDecimal4::try_new(12_345, 0).unwrap(),
749        ));
750        builder.append_variant(Variant::Decimal8(
751            VariantDecimal8::try_new(1_234_567_890_123, 0).unwrap(),
752        ));
753        builder.append_variant(Variant::Decimal16(
754            VariantDecimal16::try_new(1_234_567_890_123_456_789, 0).unwrap(),
755        ));
756        builder.append_variant(Variant::Date(NaiveDate::from_ymd_opt(2024, 11, 7).unwrap()));
757        builder.append_variant(Variant::TimestampMicros(timestamp_micros));
758        builder.append_variant(Variant::TimestampNtzMicros(timestamp_ntz_micros));
759        builder.append_variant(Variant::TimestampNanos(timestamp_nanos));
760        builder.append_variant(Variant::TimestampNtzNanos(timestamp_ntz_nanos));
761        builder.append_variant(Variant::Time(time));
762        builder.append_variant(Variant::Binary(&binary_bytes));
763        builder.append_variant(Variant::ShortString(
764            ShortString::try_new("iceberg").unwrap(),
765        ));
766        builder.append_variant(Variant::from(long_string.as_str()));
767        builder.append_variant(Variant::Uuid(uuid));
768        builder.append_variant(object_and_list.value(0));
769        builder.append_variant(object_and_list.value(1));
770
771        let variant_array = builder.build();
772        let field = variant_array.field("variant_col");
773
774        assert_eq!(
775            IcebergArrowConvert.type_from_field(&field).unwrap(),
776            DataType::Variant
777        );
778
779        let array = Arc::new(variant_array.into_inner()) as ArrayRef;
780
781        let converted = IcebergArrowConvert
782            .array_from_arrow_array(&field, &array)
783            .unwrap();
784        let values = converted
785            .into_variant()
786            .iter()
787            .map(|value| value.map(|value| value.to_text()))
788            .collect::<Vec<_>>();
789
790        let timestamp_micros_json =
791            serde_json::to_string("2024-11-07T12:33:54.123456+00:00").unwrap();
792        let timestamp_ntz_micros_json =
793            serde_json::to_string("2024-11-07T12:33:54.123456").unwrap();
794        let timestamp_nanos_json =
795            serde_json::to_string("2024-11-07T12:33:54.123456789+00:00").unwrap();
796        let timestamp_ntz_nanos_json =
797            serde_json::to_string("2024-11-07T12:33:54.123456789").unwrap();
798        let time_json = serde_json::to_string("12:33:54.123456").unwrap();
799        let binary_json = serde_json::to_string("CgsMDQ==").unwrap();
800        let short_string_json = serde_json::to_string("iceberg").unwrap();
801        let long_string_json = serde_json::to_string(&long_string).unwrap();
802        let uuid_json = serde_json::to_string("123e4567-e89b-12d3-a456-426614174000").unwrap();
803
804        assert_eq!(
805            values,
806            vec![
807                None, // null Arrow row
808                Some("null".to_owned()),
809                Some("true".to_owned()),
810                Some("false".to_owned()),
811                Some("34".to_owned()),
812                Some("1234".to_owned()),
813                Some("100000".to_owned()),
814                Some("5000000000".to_owned()),
815                Some("3.5".to_owned()),
816                Some("14.25".to_owned()),
817                Some(r#""NaN""#.to_owned()),
818                Some("12345".to_owned()),
819                Some("1234567890123".to_owned()),
820                Some("1234567890123456789".to_owned()),
821                Some(r#""2024-11-07""#.to_owned()),
822                Some(timestamp_micros_json),
823                Some(timestamp_ntz_micros_json),
824                Some(timestamp_nanos_json),
825                Some(timestamp_ntz_nanos_json),
826                Some(time_json),
827                Some(binary_json),
828                Some(short_string_json),
829                Some(long_string_json),
830                Some(uuid_json),
831                Some(r#"{"a":1,"b":[true,null]}"#.to_owned()),
832                Some(r#"[1,{"x":2},"tail"]"#.to_owned()),
833            ],
834        );
835    }
836
837    #[test]
838    fn variant_type_recurses_in_nested_types() {
839        let payload_field = arrow_schema::Field::new(
840            "payload",
841            ArrowDataType::Struct(
842                vec![
843                    Arc::new(variant_arrow_field("top_variant")),
844                    Arc::new(arrow_schema::Field::new(
845                        "variant_list",
846                        ArrowDataType::List(Arc::new(variant_arrow_field("element"))),
847                        true,
848                    )),
849                    Arc::new(arrow_schema::Field::new(
850                        "variant_map",
851                        ArrowDataType::Map(
852                            Arc::new(arrow_schema::Field::new(
853                                "entries",
854                                ArrowDataType::Struct(
855                                    vec![
856                                        Arc::new(arrow_schema::Field::new(
857                                            "key",
858                                            ArrowDataType::Utf8,
859                                            false,
860                                        )),
861                                        Arc::new(variant_arrow_field("value")),
862                                    ]
863                                    .into(),
864                                ),
865                                false,
866                            )),
867                            false,
868                        ),
869                        true,
870                    )),
871                ]
872                .into(),
873            ),
874            true,
875        );
876
877        assert_eq!(
878            IcebergArrowConvert.type_from_field(&payload_field).unwrap(),
879            DataType::Struct(StructType::new(vec![
880                ("top_variant", DataType::Variant),
881                ("variant_list", DataType::list(DataType::Variant)),
882                (
883                    "variant_map",
884                    DataType::Map(MapType::from_kv(DataType::Varchar, DataType::Variant)),
885                ),
886            ])),
887        );
888    }
889
890    #[test]
891    fn variant_map_key_is_rejected() {
892        let field = arrow_schema::Field::new(
893            "variant_map_key",
894            ArrowDataType::Map(
895                Arc::new(arrow_schema::Field::new(
896                    "entries",
897                    ArrowDataType::Struct(
898                        vec![
899                            Arc::new(variant_arrow_field("key")),
900                            Arc::new(arrow_schema::Field::new("value", ArrowDataType::Utf8, true)),
901                        ]
902                        .into(),
903                    ),
904                    false,
905                )),
906                false,
907            ),
908            true,
909        );
910
911        let err = IcebergArrowConvert.type_from_field(&field).unwrap_err();
912        assert!(err.to_string().contains("invalid map key type: variant"));
913    }
914
915    #[test]
916    fn invalid_variant_value_decodes_to_null() {
917        let field = variant_arrow_field("variant_col");
918        // Row 0 is a valid variant string "HI"; row 1 has a corrupt value byte.
919        let array = Arc::new(arrow_array::StructArray::from(vec![
920            (
921                Arc::new(arrow_schema::Field::new(
922                    "metadata",
923                    ArrowDataType::Binary,
924                    false,
925                )),
926                Arc::new(arrow_array::BinaryArray::from_iter_values([
927                    &[1_u8, 0, 0][..],
928                    &[1_u8, 0, 0][..],
929                ])) as ArrayRef,
930            ),
931            (
932                Arc::new(arrow_schema::Field::new(
933                    "value",
934                    ArrowDataType::Binary,
935                    true,
936                )),
937                Arc::new(arrow_array::BinaryArray::from_iter_values([
938                    &[0x09_u8, b'H', b'I'][..],
939                    &[255_u8][..],
940                ])) as ArrayRef,
941            ),
942        ])) as ArrayRef;
943
944        let converted = IcebergArrowConvert
945            .array_from_arrow_array(&field, &array)
946            .unwrap();
947        let values = converted
948            .into_variant()
949            .iter()
950            .map(|value| value.map(|value| value.to_text()))
951            .collect::<Vec<_>>();
952
953        assert_eq!(values, vec![Some(r#""HI""#.to_owned()), None]);
954    }
955
956    #[test]
957    fn shredded_variant_decodes_as_null() {
958        let field = variant_arrow_field("variant_col");
959        let array = Arc::new(arrow_array::StructArray::from(vec![
960            (
961                Arc::new(arrow_schema::Field::new(
962                    "metadata",
963                    ArrowDataType::Binary,
964                    false,
965                )),
966                Arc::new(arrow_array::BinaryArray::from_iter_values([
967                    &[1_u8, 0, 0][..]
968                ])) as ArrayRef,
969            ),
970            (
971                Arc::new(arrow_schema::Field::new(
972                    "typed_value",
973                    ArrowDataType::Int64,
974                    true,
975                )),
976                Arc::new(arrow_array::Int64Array::from(vec![7_i64])) as ArrayRef,
977            ),
978        ])) as ArrayRef;
979
980        let decoded = IcebergArrowConvert
981            .array_from_arrow_array(&field, &array)
982            .unwrap();
983        assert_eq!(decoded.len(), 1);
984        assert!(decoded.as_variant().is_null(0));
985    }
986
987    #[test]
988    fn variant_in_projected_struct_decodes_by_declared_type() {
989        let mut builder = VariantArrayBuilder::new(1);
990        builder.append_variant(Variant::from(7_i64));
991        let variant_array = builder.build();
992        let variant_field = variant_array.field("v");
993        let variant_child = Arc::new(variant_array.into_inner()) as ArrayRef;
994        let extra_child = Arc::new(arrow_array::Int64Array::from(vec![1_i64])) as ArrayRef;
995        let actual = Arc::new(arrow_array::StructArray::from(vec![
996            (
997                Arc::new(arrow_schema::Field::new(
998                    "extra",
999                    ArrowDataType::Int64,
1000                    true,
1001                )),
1002                extra_child,
1003            ),
1004            (Arc::new(variant_field), variant_child),
1005        ])) as ArrayRef;
1006
1007        // Declared as struct<v variant>: the parquet struct is a superset, so conversion
1008        // goes through the projected path, which must decode the variant child.
1009        let declared_field = arrow_schema::Field::new(
1010            "s",
1011            ArrowDataType::Struct(vec![Arc::new(variant_arrow_field("v"))].into()),
1012            true,
1013        );
1014        let converted = IcebergArrowConvert
1015            .array_from_arrow_array(&declared_field, &actual)
1016            .unwrap();
1017        assert_eq!(
1018            converted.data_type(),
1019            DataType::Struct(StructType::new(vec![("v", DataType::Variant)])),
1020        );
1021        let ArrayImpl::Struct(struct_array) = &converted else {
1022            panic!("expected struct array");
1023        };
1024        let values = struct_array.field_at(0).as_variant();
1025        assert_eq!(
1026            values.iter().next().unwrap().map(|v| v.to_text()),
1027            Some("7".to_owned())
1028        );
1029    }
1030}
risingwave_common/array/arrow/arrow_iceberg.rs

risingwave_common/array/arrow/
arrow_iceberg.rs
