diff --git a/crates/integrations/datafusion/src/blob_descriptor_functions.rs b/crates/integrations/datafusion/src/blob_descriptor_functions.rs new file mode 100644 index 00000000..de000412 --- /dev/null +++ b/crates/integrations/datafusion/src/blob_descriptor_functions.rs @@ -0,0 +1,222 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use datafusion::arrow::array::{ + Array, BinaryArray, BinaryBuilder, BinaryViewArray, LargeBinaryArray, LargeStringArray, + StringArray, StringBuilder, StringViewArray, +}; +use datafusion::arrow::datatypes::DataType as ArrowDataType; +use datafusion::common::types::logical_binary; +use datafusion::common::utils::take_function_args; +use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue}; +use datafusion::logical_expr::{ + Coercion, ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, + TypeSignatureClass, Volatility, +}; +use datafusion::prelude::SessionContext; +use paimon::spec::BlobDescriptor; + +use crate::error::to_datafusion_error; + +const PATH_TO_DESCRIPTOR: &str = "path_to_descriptor"; +const DESCRIPTOR_TO_STRING: &str = "descriptor_to_string"; + +pub(crate) fn register_blob_descriptor_functions(ctx: &SessionContext) { + ctx.register_udf(ScalarUDF::from(PathToDescriptorFunc::new())); + ctx.register_udf(ScalarUDF::from(DescriptorToStringFunc::new())); +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct PathToDescriptorFunc { + signature: Signature, + aliases: Vec, +} + +impl PathToDescriptorFunc { + fn new() -> Self { + Self { + signature: Signature::string(1, Volatility::Immutable), + aliases: vec!["sys.path_to_descriptor".to_string()], + } + } +} + +impl ScalarUDFImpl for PathToDescriptorFunc { + fn name(&self) -> &str { + PATH_TO_DESCRIPTOR + } + + fn aliases(&self) -> &[String] { + &self.aliases + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[ArrowDataType]) -> DFResult { + Ok(ArrowDataType::Binary) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DFResult { + let [input] = take_function_args(self.name(), args.args)?; + match input { + ColumnarValue::Scalar(value) => path_scalar(value), + ColumnarValue::Array(array) => path_array(array.as_ref()), + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct DescriptorToStringFunc { + signature: Signature, + aliases: Vec, +} + +impl DescriptorToStringFunc { + fn new() -> Self { + Self { + signature: Signature::coercible( + vec![Coercion::new_exact(TypeSignatureClass::Native( + logical_binary(), + ))], + Volatility::Immutable, + ), + aliases: vec!["sys.descriptor_to_string".to_string()], + } + } +} + +impl ScalarUDFImpl for DescriptorToStringFunc { + fn name(&self) -> &str { + DESCRIPTOR_TO_STRING + } + + fn aliases(&self) -> &[String] { + &self.aliases + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[ArrowDataType]) -> DFResult { + Ok(ArrowDataType::Utf8) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DFResult { + let [input] = take_function_args(self.name(), args.args)?; + match input { + ColumnarValue::Scalar(value) => descriptor_scalar(value), + ColumnarValue::Array(array) => descriptor_array(array.as_ref()), + } + } +} + +fn serialize_path(path: &str) -> Vec { + BlobDescriptor::new(path.to_string(), 0, -1).serialize() +} + +fn descriptor_string(bytes: &[u8]) -> DFResult { + BlobDescriptor::deserialize(bytes) + .map(|descriptor| descriptor.to_string()) + .map_err(to_datafusion_error) +} + +fn path_scalar(value: ScalarValue) -> DFResult { + let path = match value { + ScalarValue::Utf8(path) | ScalarValue::LargeUtf8(path) | ScalarValue::Utf8View(path) => { + path + } + ScalarValue::Null => None, + other => return unexpected_type(PATH_TO_DESCRIPTOR, &other.data_type()), + }; + Ok(ColumnarValue::Scalar(ScalarValue::Binary( + path.as_deref().map(serialize_path), + ))) +} + +fn path_array(input: &dyn Array) -> DFResult { + if let Some(values) = input.as_any().downcast_ref::() { + return Ok(descriptor_array_from_paths(values.iter())); + } + if let Some(values) = input.as_any().downcast_ref::() { + return Ok(descriptor_array_from_paths(values.iter())); + } + if let Some(values) = input.as_any().downcast_ref::() { + return Ok(descriptor_array_from_paths(values.iter())); + } + unexpected_type(PATH_TO_DESCRIPTOR, input.data_type()) +} + +fn descriptor_array_from_paths<'a>(paths: impl Iterator>) -> ColumnarValue { + let mut builder = BinaryBuilder::new(); + for path in paths { + match path { + Some(path) => builder.append_value(serialize_path(path)), + None => builder.append_null(), + } + } + ColumnarValue::Array(Arc::new(builder.finish())) +} + +fn descriptor_scalar(value: ScalarValue) -> DFResult { + let bytes = match value { + ScalarValue::Binary(bytes) + | ScalarValue::LargeBinary(bytes) + | ScalarValue::BinaryView(bytes) => bytes, + ScalarValue::Null => None, + other => return unexpected_type(DESCRIPTOR_TO_STRING, &other.data_type()), + }; + Ok(ColumnarValue::Scalar(ScalarValue::Utf8( + bytes.as_deref().map(descriptor_string).transpose()?, + ))) +} + +fn descriptor_array(input: &dyn Array) -> DFResult { + if let Some(values) = input.as_any().downcast_ref::() { + return strings_from_descriptors(values.iter()); + } + if let Some(values) = input.as_any().downcast_ref::() { + return strings_from_descriptors(values.iter()); + } + if let Some(values) = input.as_any().downcast_ref::() { + return strings_from_descriptors(values.iter()); + } + unexpected_type(DESCRIPTOR_TO_STRING, input.data_type()) +} + +fn strings_from_descriptors<'a>( + descriptors: impl Iterator>, +) -> DFResult { + let mut builder = StringBuilder::new(); + for bytes in descriptors { + match bytes { + Some(bytes) => builder.append_value(descriptor_string(bytes)?), + None => builder.append_null(), + } + } + Ok(ColumnarValue::Array(Arc::new(builder.finish()))) +} + +fn unexpected_type(function: &str, data_type: &ArrowDataType) -> DFResult { + Err(DataFusionError::Execution(format!( + "{function} received unexpected argument type {data_type}" + ))) +} diff --git a/crates/integrations/datafusion/src/lib.rs b/crates/integrations/datafusion/src/lib.rs index a60d52b2..05e4037a 100644 --- a/crates/integrations/datafusion/src/lib.rs +++ b/crates/integrations/datafusion/src/lib.rs @@ -36,6 +36,7 @@ //! This version supports partition predicate pushdown by extracting //! translatable partition-only conjuncts from DataFusion filters. +mod blob_descriptor_functions; mod blob_reader; mod blob_view; mod catalog; diff --git a/crates/integrations/datafusion/src/sql_context.rs b/crates/integrations/datafusion/src/sql_context.rs index 55232294..f9d81c86 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -117,6 +117,7 @@ impl SQLContext { )) .build(); let ctx = SessionContext::new_with_state(state); + crate::blob_descriptor_functions::register_blob_descriptor_functions(&ctx); crate::variant_functions::register_variant_functions(&ctx); Self { ctx, diff --git a/crates/integrations/datafusion/tests/blob_descriptor_functions.rs b/crates/integrations/datafusion/tests/blob_descriptor_functions.rs new file mode 100644 index 00000000..029dcc4b --- /dev/null +++ b/crates/integrations/datafusion/tests/blob_descriptor_functions.rs @@ -0,0 +1,183 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use datafusion::arrow::array::{Array, BinaryArray, StringArray}; +use paimon_datafusion::SQLContext; + +const JAVA_V2_HEX: &str = + "0243534544424f4c420d00000066696c653a2f2f2f746d702f610000000000000000ffffffffffffffff"; +const JAVA_V2_STRING: &str = "BlobDescriptor{version=2, uri='file:///tmp/a', offset=0, length=-1}"; +const JAVA_V1_HEX: &str = "010a0000002f746573742f706174686400000000000000c800000000000000"; +const JAVA_V1_STRING: &str = "BlobDescriptor{version=1, uri='/test/path', offset=100, length=200}"; + +fn to_hex(bytes: &[u8]) -> String { + bytes.iter().map(|byte| format!("{byte:02x}")).collect() +} + +async fn query_error(ctx: &SQLContext, sql: &str) -> String { + match ctx.sql(sql).await { + Err(error) => error.to_string(), + Ok(dataframe) => dataframe + .collect() + .await + .expect_err("query should fail") + .to_string(), + } +} + +#[tokio::test] +async fn test_blob_descriptor_functions_are_registered_with_aliases_and_java_format() { + let ctx = SQLContext::new(); + let batches = ctx + .sql( + "SELECT \ + path_to_descriptor('file:///tmp/a'), \ + sys.path_to_descriptor('file:///tmp/a'), \ + descriptor_to_string(path_to_descriptor('file:///tmp/a')), \ + sys.descriptor_to_string(sys.path_to_descriptor('file:///tmp/a'))", + ) + .await + .unwrap() + .collect() + .await + .unwrap(); + + let batch = &batches[0]; + for column in 0..2 { + let descriptors = batch + .column(column) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(to_hex(descriptors.value(0)), JAVA_V2_HEX); + } + for column in 2..4 { + let strings = batch + .column(column) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(strings.value(0), JAVA_V2_STRING); + } +} + +#[tokio::test] +async fn test_blob_descriptor_functions_propagate_nulls() { + let ctx = SQLContext::new(); + let batches = ctx + .sql( + "SELECT id, \ + path_to_descriptor(path), \ + descriptor_to_string(path_to_descriptor(path)) \ + FROM (VALUES \ + (1, 'file:///tmp/a'), \ + (2, CAST(NULL AS VARCHAR)), \ + (3, 'file:///tmp/b') \ + ) AS inputs(id, path) \ + ORDER BY id", + ) + .await + .unwrap() + .collect() + .await + .unwrap(); + + let batch = &batches[0]; + let descriptors = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let strings = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(to_hex(descriptors.value(0)), JAVA_V2_HEX); + assert!(descriptors.is_null(1)); + assert!(strings.is_null(1)); + assert_eq!( + strings.value(2), + "BlobDescriptor{version=2, uri='file:///tmp/b', offset=0, length=-1}" + ); + + let alias_nulls = ctx + .sql( + "SELECT \ + path_to_descriptor(NULL), \ + sys.path_to_descriptor(NULL), \ + descriptor_to_string(NULL), \ + sys.descriptor_to_string(NULL)", + ) + .await + .unwrap() + .collect() + .await + .unwrap(); + let batch = &alias_nulls[0]; + for column in 0..4 { + assert!(batch.column(column).is_null(0)); + } +} + +#[tokio::test] +async fn test_descriptor_to_string_supports_java_v1() { + let ctx = SQLContext::new(); + let sql = format!( + "SELECT descriptor_to_string(X'{JAVA_V1_HEX}'), \ + sys.descriptor_to_string(X'{JAVA_V1_HEX}')" + ); + let batches = ctx.sql(&sql).await.unwrap().collect().await.unwrap(); + let batch = &batches[0]; + for column in 0..2 { + let strings = batch + .column(column) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(strings.value(0), JAVA_V1_STRING); + } +} + +#[tokio::test] +async fn test_blob_descriptor_functions_reject_invalid_arguments() { + let ctx = SQLContext::new(); + for (sql, function) in [ + ("SELECT path_to_descriptor()", "path_to_descriptor"), + ("SELECT path_to_descriptor(1)", "path_to_descriptor"), + ( + "SELECT descriptor_to_string('not binary')", + "descriptor_to_string", + ), + ( + "SELECT descriptor_to_string(X'00', X'01')", + "descriptor_to_string", + ), + ] { + let error = query_error(&ctx, sql).await; + assert!( + error.contains(function), + "expected error for {function}, got: {error}" + ); + } + + let error = query_error(&ctx, "SELECT descriptor_to_string(X'00')").await; + assert!( + error.contains("BlobDescriptor bytes too short"), + "unexpected malformed descriptor error: {error}" + ); +} diff --git a/crates/paimon/src/spec/blob_descriptor.rs b/crates/paimon/src/spec/blob_descriptor.rs index 93c84491..f65b293d 100644 --- a/crates/paimon/src/spec/blob_descriptor.rs +++ b/crates/paimon/src/spec/blob_descriptor.rs @@ -15,6 +15,8 @@ // specific language governing permissions and limitations // under the License. +use std::fmt::{Display, Formatter}; + use crate::Error; const CURRENT_VERSION: u8 = 2; @@ -28,6 +30,16 @@ pub struct BlobDescriptor { length: i64, } +impl Display for BlobDescriptor { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "BlobDescriptor{{version={}, uri='{}', offset={}, length={}}}", + self.version, self.uri, self.offset, self.length + ) + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) struct BlobRangeSpec { offset: u64, @@ -221,6 +233,15 @@ mod tests { assert_eq!(desc, deserialized); } + #[test] + fn test_display_matches_java() { + let desc = BlobDescriptor::new("file:///tmp/a".to_string(), 0, -1); + assert_eq!( + desc.to_string(), + "BlobDescriptor{version=2, uri='file:///tmp/a', offset=0, length=-1}" + ); + } + #[test] fn test_is_blob_descriptor() { let desc = BlobDescriptor::new("file:///tmp/test.blob".to_string(), 0, 1024); diff --git a/docs/src/sql.md b/docs/src/sql.md index d550c734..01635ba6 100644 --- a/docs/src/sql.md +++ b/docs/src/sql.md @@ -68,13 +68,14 @@ async fn example() -> Result<(), Box> { } ``` -`SQLContext::new` creates a session context with the Paimon relation planner -pre-registered. Use `register_catalog(...).await` to add one or more Paimon -catalogs; registering a catalog also registers the built-in scalar function -`blob_view` (alias `sys.blob_view`) and the built-in table-valued functions -(`vector_search`, `hybrid_search`, and `full_text_search` when the `fulltext` -feature is enabled) against it. It also manages session-scoped dynamic options -internally for `SET`/`RESET` support. +`SQLContext::new` creates a session context with the Paimon relation planner and +the catalog-independent `path_to_descriptor` and `descriptor_to_string` scalar +functions pre-registered. Use `register_catalog(...).await` to add one or more +Paimon catalogs; registering a catalog also registers the built-in scalar +function `blob_view` (alias `sys.blob_view`) and the built-in table-valued +functions (`vector_search`, `hybrid_search`, and `full_text_search` when the +`fulltext` feature is enabled) against it. It also manages session-scoped +dynamic options internally for `SET`/`RESET` support. ### REST Catalog Views and SQL Functions @@ -300,6 +301,25 @@ The offset must be non-negative, and lengths below `-1` are invalid. The same directives are supported by `ALTER TABLE ... ADD COLUMN`. +### Blob Descriptor Functions + +`path_to_descriptor(path)` converts a string path into Java-compatible +`BlobDescriptor` bytes with offset `0` and length `-1`. Its alias is +`sys.path_to_descriptor(path)`. The function only serializes the path; it does +not access the referenced object or validate that it exists. + +`descriptor_to_string(descriptor)` converts serialized descriptor bytes to the +same string representation used by Java Paimon. Its alias is +`sys.descriptor_to_string(descriptor)`. Invalid descriptor bytes return an +error. Both functions return `NULL` for `NULL` input. + +```sql +SELECT sys.descriptor_to_string( + sys.path_to_descriptor('file:///tmp/image.png') +); +-- BlobDescriptor{version=2, uri='file:///tmp/image.png', offset=0, length=-1} +``` + ### Blob View Blob View stores an inline reference to a BLOB value in another table, using a