paimon_datafusion/
blob_view.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::collections::HashMap;
19use std::hash::{Hash, Hasher};
20use std::sync::Arc;
21
22use datafusion::arrow::array::{
23    Array, BinaryBuilder, Int16Array, Int32Array, Int64Array, Int8Array, LargeStringArray,
24    StringArray, StringViewArray, UInt16Array, UInt32Array, UInt64Array, UInt8Array,
25};
26use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, FieldRef};
27use datafusion::common::{DataFusionError, Result as DFResult};
28use datafusion::logical_expr::{
29    ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature,
30    Volatility,
31};
32use datafusion::prelude::SessionContext;
33use paimon::catalog::Catalog;
34use paimon::spec::BlobViewStruct;
35
36use crate::error::to_datafusion_error;
37use crate::runtime::block_on_with_runtime;
38use crate::table_function_args::parse_table_identifier;
39use crate::table_loader::load_data_table_for_read;
40
41const FUNCTION_NAME: &str = "blob_view";
42
43pub fn register_blob_view(ctx: &SessionContext, catalog: Arc<dyn Catalog>, default_database: &str) {
44    ctx.register_udf(ScalarUDF::from(BlobViewFunc::new(
45        catalog,
46        default_database,
47    )));
48}
49
50#[derive(Clone)]
51struct BlobViewFunc {
52    catalog: Arc<dyn Catalog>,
53    default_database: String,
54    signature: Signature,
55    aliases: Vec<String>,
56}
57
58impl BlobViewFunc {
59    fn new(catalog: Arc<dyn Catalog>, default_database: &str) -> Self {
60        Self {
61            catalog,
62            default_database: default_database.to_string(),
63            signature: Signature::any(3, Volatility::Immutable),
64            aliases: vec!["sys.blob_view".to_string()],
65        }
66    }
67
68    fn field_id(&self, table_name: &str, field_name: &str) -> DFResult<i32> {
69        let identifier = parse_table_identifier(FUNCTION_NAME, table_name, &self.default_database)?;
70        let catalog = Arc::clone(&self.catalog);
71        let table = block_on_with_runtime(
72            async move { load_data_table_for_read(&catalog, &identifier, FUNCTION_NAME).await },
73            "blob_view: catalog access thread panicked",
74        )?;
75
76        let field = table
77            .schema()
78            .fields()
79            .iter()
80            .find(|field| field.name() == field_name)
81            .ok_or_else(|| {
82                DataFusionError::Plan(format!(
83                    "blob_view: cannot find blob field {field_name} in upstream table {table_name}"
84                ))
85            })?;
86        if !field.data_type().is_blob_type() {
87            return Err(DataFusionError::Plan(format!(
88                "blob_view: field {field_name} in upstream table {table_name} is not a BLOB field"
89            )));
90        }
91        Ok(field.id())
92    }
93}
94
95impl std::fmt::Debug for BlobViewFunc {
96    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
97        f.debug_struct("BlobViewFunc")
98            .field("default_database", &self.default_database)
99            .field("aliases", &self.aliases)
100            .finish()
101    }
102}
103
104impl PartialEq for BlobViewFunc {
105    fn eq(&self, other: &Self) -> bool {
106        self.default_database == other.default_database
107            && Arc::ptr_eq(&self.catalog, &other.catalog)
108    }
109}
110
111impl Eq for BlobViewFunc {}
112
113impl Hash for BlobViewFunc {
114    fn hash<H: Hasher>(&self, state: &mut H) {
115        FUNCTION_NAME.hash(state);
116        self.default_database.hash(state);
117        let ptr = Arc::as_ptr(&self.catalog) as *const () as usize;
118        ptr.hash(state);
119    }
120}
121
122impl ScalarUDFImpl for BlobViewFunc {
123    fn name(&self) -> &str {
124        FUNCTION_NAME
125    }
126
127    fn aliases(&self) -> &[String] {
128        &self.aliases
129    }
130
131    fn signature(&self) -> &Signature {
132        &self.signature
133    }
134
135    fn return_type(&self, _arg_types: &[ArrowDataType]) -> DFResult<ArrowDataType> {
136        Ok(ArrowDataType::Binary)
137    }
138
139    fn return_field_from_args(&self, _args: ReturnFieldArgs) -> DFResult<FieldRef> {
140        Ok(Arc::new(Field::new(
141            FUNCTION_NAME,
142            ArrowDataType::Binary,
143            true,
144        )))
145    }
146
147    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DFResult<ColumnarValue> {
148        if args.args.len() != 3 {
149            return Err(DataFusionError::Plan(
150                "blob_view expects 3 arguments: (table_name, field_name_or_id, row_id)".to_string(),
151            ));
152        }
153
154        let arrays = ColumnarValue::values_to_arrays(&args.args)?;
155        let table_names = arrays[0].as_ref();
156        let fields = arrays[1].as_ref();
157        let row_ids = arrays[2].as_ref();
158        let mut field_id_cache = HashMap::new();
159        let mut builder = BinaryBuilder::new();
160
161        for row in 0..table_names.len() {
162            let Some(table_name) = string_at(table_names, row)? else {
163                builder.append_null();
164                continue;
165            };
166            let Some(row_id) = int64_at(row_ids, row, "row_id")? else {
167                builder.append_null();
168                continue;
169            };
170
171            let field_id = if is_string_array(fields) {
172                let Some(field_name) = string_at(fields, row)? else {
173                    builder.append_null();
174                    continue;
175                };
176                let cache_key = (table_name.clone(), field_name.clone());
177                match field_id_cache.get(&cache_key) {
178                    Some(field_id) => *field_id,
179                    None => {
180                        let field_id = self.field_id(&table_name, &field_name)?;
181                        field_id_cache.insert(cache_key, field_id);
182                        field_id
183                    }
184                }
185            } else {
186                match int64_at(fields, row, "field_id")? {
187                    Some(field_id) => i32::try_from(field_id).map_err(|_| {
188                        DataFusionError::Plan(format!(
189                            "blob_view: field_id {field_id} is outside i32 range"
190                        ))
191                    })?,
192                    None => {
193                        builder.append_null();
194                        continue;
195                    }
196                }
197            };
198
199            let identifier =
200                parse_table_identifier(FUNCTION_NAME, &table_name, &self.default_database)?;
201            let value = BlobViewStruct::new(identifier, field_id, row_id)
202                .serialize()
203                .map_err(to_datafusion_error)?;
204            builder.append_value(value);
205        }
206
207        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
208    }
209}
210
211fn is_string_array(array: &dyn Array) -> bool {
212    matches!(
213        array.data_type(),
214        ArrowDataType::Utf8 | ArrowDataType::LargeUtf8 | ArrowDataType::Utf8View
215    )
216}
217
218fn string_at(array: &dyn Array, row: usize) -> DFResult<Option<String>> {
219    if array.is_null(row) {
220        return Ok(None);
221    }
222    if let Some(values) = array.as_any().downcast_ref::<StringArray>() {
223        return Ok(Some(values.value(row).to_string()));
224    }
225    if let Some(values) = array.as_any().downcast_ref::<LargeStringArray>() {
226        return Ok(Some(values.value(row).to_string()));
227    }
228    if let Some(values) = array.as_any().downcast_ref::<StringViewArray>() {
229        return Ok(Some(values.value(row).to_string()));
230    }
231    Err(DataFusionError::Plan(format!(
232        "blob_view: expected string argument, got {:?}",
233        array.data_type()
234    )))
235}
236
237fn int64_at(array: &dyn Array, row: usize, name: &str) -> DFResult<Option<i64>> {
238    if array.is_null(row) {
239        return Ok(None);
240    }
241    macro_rules! downcast_int {
242        ($ty:ty) => {
243            if let Some(values) = array.as_any().downcast_ref::<$ty>() {
244                return Ok(Some(values.value(row) as i64));
245            }
246        };
247    }
248    downcast_int!(Int8Array);
249    downcast_int!(Int16Array);
250    downcast_int!(Int32Array);
251    downcast_int!(Int64Array);
252    downcast_int!(UInt8Array);
253    downcast_int!(UInt16Array);
254    downcast_int!(UInt32Array);
255    if let Some(values) = array.as_any().downcast_ref::<UInt64Array>() {
256        return i64::try_from(values.value(row)).map(Some).map_err(|_| {
257            DataFusionError::Plan(format!(
258                "blob_view: {name} {} is outside i64 range",
259                values.value(row)
260            ))
261        });
262    }
263    Err(DataFusionError::Plan(format!(
264        "blob_view: expected integer {name}, got {:?}",
265        array.data_type()
266    )))
267}
paimon_datafusion/blob_view.rs

paimon_datafusion/
blob_view.rs
