risingwave_expr_impl/scalar/external/
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
15//! This module contains the expression for computing the iceberg partition value.
16//! spec ref: <https://iceberg.apache.org/spec/#partition-transforms>
17use std::fmt::Formatter;
18use std::str::FromStr;
19use std::sync::Arc;
20
21use anyhow::anyhow;
22use iceberg::spec::{PrimitiveType, Transform, Type as IcebergType};
23use iceberg::transform::{BoxedTransformFunction, create_transform_function};
24use risingwave_common::array::arrow::{IcebergArrowConvert, arrow_schema_iceberg};
25use risingwave_common::array::{ArrayRef, DataChunk};
26use risingwave_common::ensure;
27use risingwave_common::row::OwnedRow;
28use risingwave_common::types::{DataType, Datum};
29use risingwave_expr::expr::{
30    BoxedExpression, ExpressionInfo, SyncExpression, SyncExpressionBoxExt,
31};
32use risingwave_expr::{ExprError, Result, build_function};
33use thiserror_ext::AsReport;
34
35pub struct IcebergTransform {
36    child: Arc<dyn SyncExpression>,
37    transform: BoxedTransformFunction,
38    input_arrow_type: arrow_schema_iceberg::DataType,
39    output_arrow_field: arrow_schema_iceberg::Field,
40    return_type: DataType,
41}
42
43impl std::fmt::Debug for IcebergTransform {
44    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
45        f.debug_struct("IcebergTransform")
46            .field("child", &self.child)
47            .field("return_type", &self.return_type)
48            .finish()
49    }
50}
51
52impl ExpressionInfo for IcebergTransform {
53    fn return_type(&self) -> DataType {
54        self.return_type.clone()
55    }
56}
57
58impl SyncExpression for IcebergTransform {
59    fn eval(&self, data_chunk: &DataChunk) -> Result<ArrayRef> {
60        let array = self.child.eval(data_chunk)?;
61        // Convert to arrow array
62        let arrow_array = IcebergArrowConvert.to_arrow_array(&self.input_arrow_type, &array)?;
63        // Transform
64        let res_array = self.transform.transform(arrow_array).unwrap();
65        // Convert back to array ref and return it
66        Ok(Arc::new(IcebergArrowConvert.array_from_arrow_array(
67            &self.output_arrow_field,
68            &res_array,
69        )?))
70    }
71
72    fn eval_row(&self, _row: &OwnedRow) -> Result<Datum> {
73        Err(ExprError::Internal(anyhow!(
74            "eval_row in iceberg_transform is not supported yet"
75        )))
76    }
77}
78
79#[build_function("iceberg_transform(varchar, any) -> any", type_infer = "unreachable")]
80fn build(return_type: DataType, mut children: Vec<BoxedExpression>) -> Result<BoxedExpression> {
81    let transform_type = {
82        let datum = children[0].eval_const()?.unwrap();
83        let str = datum.as_utf8();
84        Transform::from_str(str).map_err(|_| ExprError::InvalidParam {
85            name: "transform type in iceberg_transform",
86            reason: format!("Fail to parse {str} as iceberg transform type").into(),
87        })?
88    };
89
90    // For Identity and Void transform, we will use `InputRef` and const null in frontend,
91    // so it should not reach here.
92    assert!(!matches!(
93        transform_type,
94        Transform::Identity | Transform::Void
95    ));
96
97    // Check type:
98    // 1. input type can be transform successfully
99    // 2. return type is the same as the result type
100    let input_arrow_type = IcebergArrowConvert
101        .to_arrow_field("", &children[1].return_type())?
102        .data_type()
103        .clone();
104    let output_arrow_field = IcebergArrowConvert.to_arrow_field("", &return_type)?;
105    let input_type = iceberg::arrow::arrow_type_to_type(&input_arrow_type).map_err(|err| {
106        ExprError::InvalidParam {
107            name: "input type in iceberg_transform",
108            reason: format!(
109                "Failed to convert input type to iceberg type, got error: {}",
110                err.as_report()
111            )
112            .into(),
113        }
114    })?;
115    let expect_res_type = transform_type.result_type(&input_type).map_err(
116        |err| ExprError::InvalidParam {
117            name: "input type in iceberg_transform",
118            reason: format!(
119                "Failed to get result type for transform type {:?} and input type {:?}, got error: {}",
120                transform_type, input_type, err.as_report()
121            )
122            .into()
123        })?;
124    let actual_res_type = iceberg::arrow::arrow_type_to_type(
125        &IcebergArrowConvert
126            .to_arrow_field("", &return_type)?
127            .data_type()
128            .clone(),
129    )
130    .map_err(|err| ExprError::InvalidParam {
131        name: "return type in iceberg_transform",
132        reason: format!(
133            "Failed to convert return type to iceberg type, got error: {}",
134            err.as_report()
135        )
136        .into(),
137    })?;
138
139    ensure!(
140        (expect_res_type == actual_res_type)
141            ||
142            // This is a confusing stuff.<https://github.com/apache/iceberg/pull/11749>
143            (matches!(transform_type, Transform::Day) && matches!(actual_res_type, IcebergType::Primitive(PrimitiveType::Int))),
144        ExprError::InvalidParam {
145            name: "return type in iceberg_transform",
146            reason: format!(
147                "Expect return type {:?} but got {:?}, RisingWave return type is {:?}, input type is {:?}, transform type is {:?}",
148                expect_res_type,
149                actual_res_type,
150                return_type,
151                (input_type, input_arrow_type),
152                transform_type
153            )
154            .into()
155        }
156    );
157
158    let child = match children.remove(1) {
159        BoxedExpression::Sync(child) => child,
160        BoxedExpression::Async(_) => {
161            return Err(ExprError::Internal(anyhow!(
162                "async child in iceberg_transform is not supported"
163            )));
164        }
165    };
166    let transform = create_transform_function(&transform_type)
167        .map_err(|err| ExprError::Internal(err.into()))?;
168    Ok(IcebergTransform {
169        child,
170        transform,
171        input_arrow_type,
172        output_arrow_field,
173        return_type,
174    }
175    .boxed())
176}
177
178#[cfg(test)]
179mod test {
180    use risingwave_common::array::{DataChunk, DataChunkTestExt};
181    use risingwave_expr::expr::build_from_pretty;
182
183    #[tokio::test]
184    async fn test_bucket() {
185        let (input, expected) = DataChunk::from_pretty(
186            "i   i
187             34  1373",
188        )
189        .split_column_at(1);
190        let expr = build_from_pretty("(iceberg_transform:int4 bucket[2017]:varchar $0:int)");
191        let res = expr.eval(&input).await.unwrap();
192        assert_eq!(res, *expected.column_at(0));
193    }
194
195    #[tokio::test]
196    async fn test_truncate() {
197        let (input, expected) = DataChunk::from_pretty(
198            "T         T
199            iceberg   ice
200            risingwave ris
201            delta     del",
202        )
203        .split_column_at(1);
204        let expr = build_from_pretty("(iceberg_transform:varchar truncate[3]:varchar $0:varchar)");
205        let res = expr.eval(&input).await.unwrap();
206        assert_eq!(res, *expected.column_at(0));
207    }
208
209    #[tokio::test]
210    async fn test_year_month_day_hour() {
211        let (input, expected) = DataChunk::from_pretty(
212            "TZ                                  i i D i
213            1970-01-01T00:00:00.000000000+00:00  0 0 1970-01-01 0
214            1971-02-01T01:00:00.000000000+00:00  1 13 1971-02-01 9505
215            1972-03-01T02:00:00.000000000+00:00  2 26 1972-03-01 18962
216            1970-05-01T06:00:00.000000000+00:00  0 4 1970-05-01 2886
217            1970-06-01T07:00:00.000000000+00:00  0 5 1970-06-01 3631",
218        )
219        .split_column_at(1);
220
221        // year
222        let expr = build_from_pretty("(iceberg_transform:int4 year:varchar $0:timestamptz)");
223        let res = expr.eval(&input).await.unwrap();
224        assert_eq!(res, *expected.column_at(0));
225
226        // month
227        let expr = build_from_pretty("(iceberg_transform:int4 month:varchar $0:timestamptz)");
228        let res = expr.eval(&input).await.unwrap();
229        assert_eq!(res, *expected.column_at(1));
230
231        // day
232        let expr = build_from_pretty("(iceberg_transform:int4 day:varchar $0:timestamptz)");
233        let res = expr.eval(&input).await.unwrap();
234        assert_eq!(res, *expected.column_at(2));
235
236        // hour
237        let expr = build_from_pretty("(iceberg_transform:int4 hour:varchar $0:timestamptz)");
238        let res = expr.eval(&input).await.unwrap();
239        assert_eq!(res, *expected.column_at(3));
240    }
241}
risingwave_expr_impl/scalar/external/iceberg.rs

risingwave_expr_impl/scalar/external/
iceberg.rs
