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

risingwave_common/array/arrow/
arrow_iceberg.rs
