tobixdev commented on code in PR #20312: URL: https://github.com/apache/datafusion/pull/20312#discussion_r2969273500
########## datafusion-examples/examples/extension_types/temperature.rs: ########## @@ -0,0 +1,268 @@ +// 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 arrow::array::{Array, Float64Array, RecordBatch, StringArray}; +use arrow::util::display::{ArrayFormatter, DisplayIndex, FormatOptions, FormatResult}; +use arrow_schema::extension::ExtensionType; +use arrow_schema::{ArrowError, DataType, Field, Schema, SchemaRef}; +use datafusion::dataframe::DataFrame; +use datafusion::error::Result; +use datafusion::execution::SessionStateBuilder; +use datafusion::prelude::SessionContext; +use datafusion_common::internal_err; +use datafusion_common::types::DFExtensionType; +use datafusion_expr::registry::{ + DefaultExtensionTypeRegistration, ExtensionTypeRegistry, MemoryExtensionTypeRegistry, +}; +use std::fmt::Write; +use std::sync::Arc; + +/// This example demonstrates using DataFusion's extension type API to create a custom +/// semantic type [`TemperatureExtensionType`]. Review Comment: TODO Probably remove semantic and just say extension type. I think it's confusing. ########## datafusion-examples/examples/extension_types/temperature.rs: ########## @@ -0,0 +1,268 @@ +// 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 arrow::array::{Array, Float64Array, RecordBatch, StringArray}; +use arrow::util::display::{ArrayFormatter, DisplayIndex, FormatOptions, FormatResult}; +use arrow_schema::extension::ExtensionType; +use arrow_schema::{ArrowError, DataType, Field, Schema, SchemaRef}; +use datafusion::dataframe::DataFrame; +use datafusion::error::Result; +use datafusion::execution::SessionStateBuilder; +use datafusion::prelude::SessionContext; +use datafusion_common::internal_err; +use datafusion_common::types::DFExtensionType; +use datafusion_expr::registry::{ + DefaultExtensionTypeRegistration, ExtensionTypeRegistry, MemoryExtensionTypeRegistry, +}; +use std::fmt::Write; +use std::sync::Arc; + +/// This example demonstrates using DataFusion's extension type API to create a custom +/// semantic type [`TemperatureExtensionType`]. +pub async fn temperature_example() -> Result<()> { + let ctx = create_session_context()?; + register_temperature_table(&ctx).await?; + + // Print the example table with the custom pretty-printer. + ctx.table("example").await?.show().await +} + +/// Creates the DataFusion session context with the custom extension type implementation. +fn create_session_context() -> Result<SessionContext> { + let registry = MemoryExtensionTypeRegistry::new_empty(); + + // The registration creates a new instance of the extension type with the deserialized metadata. + let temp_registration = DefaultExtensionTypeRegistration::new_arc(|metadata| { + Ok(TemperatureExtensionType(metadata)) + }); + registry.add_extension_type_registration(temp_registration)?; + + let state = SessionStateBuilder::default() + .with_extension_type_registry(Arc::new(registry)) + .build(); + Ok(SessionContext::new_with_state(state)) +} + +/// Registers the example table and returns the data frame. +async fn register_temperature_table(ctx: &SessionContext) -> Result<DataFrame> { + let schema = example_schema(); + + let city_names = Arc::new(StringArray::from(vec![ + "Vienna", "Tokyo", "New York", "Sydney", + ])); + + // We'll use the same raw float values across columns to show how the Review Comment: TODO: That's not accurate anymore. ########## datafusion-examples/examples/extension_types/temperature.rs: ########## @@ -0,0 +1,268 @@ +// 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 arrow::array::{Array, Float64Array, RecordBatch, StringArray}; +use arrow::util::display::{ArrayFormatter, DisplayIndex, FormatOptions, FormatResult}; +use arrow_schema::extension::ExtensionType; +use arrow_schema::{ArrowError, DataType, Field, Schema, SchemaRef}; +use datafusion::dataframe::DataFrame; +use datafusion::error::Result; +use datafusion::execution::SessionStateBuilder; +use datafusion::prelude::SessionContext; +use datafusion_common::internal_err; +use datafusion_common::types::DFExtensionType; +use datafusion_expr::registry::{ + DefaultExtensionTypeRegistration, ExtensionTypeRegistry, MemoryExtensionTypeRegistry, +}; +use std::fmt::Write; +use std::sync::Arc; + +/// This example demonstrates using DataFusion's extension type API to create a custom +/// semantic type [`TemperatureExtensionType`]. +pub async fn temperature_example() -> Result<()> { + let ctx = create_session_context()?; + register_temperature_table(&ctx).await?; + + // Print the example table with the custom pretty-printer. + ctx.table("example").await?.show().await +} + +/// Creates the DataFusion session context with the custom extension type implementation. +fn create_session_context() -> Result<SessionContext> { + let registry = MemoryExtensionTypeRegistry::new_empty(); + + // The registration creates a new instance of the extension type with the deserialized metadata. + let temp_registration = DefaultExtensionTypeRegistration::new_arc(|metadata| { + Ok(TemperatureExtensionType(metadata)) + }); + registry.add_extension_type_registration(temp_registration)?; + + let state = SessionStateBuilder::default() + .with_extension_type_registry(Arc::new(registry)) + .build(); + Ok(SessionContext::new_with_state(state)) +} + +/// Registers the example table and returns the data frame. +async fn register_temperature_table(ctx: &SessionContext) -> Result<DataFrame> { + let schema = example_schema(); + + let city_names = Arc::new(StringArray::from(vec![ + "Vienna", "Tokyo", "New York", "Sydney", + ])); + + // We'll use the same raw float values across columns to show how the + // extension type changes the formatting behavior. + let celsius_temps = vec![15.1, 22.5, 18.98, 25.0]; + let fahrenheit_temps = vec![59.18, 72.5, 66.164, 77.0]; + let kelvin_temps = vec![288.25, 295.65, 292.13, 298.15]; + + let batch = RecordBatch::try_new( + schema, + vec![ + city_names, + Arc::new(Float64Array::from(celsius_temps)), + Arc::new(Float64Array::from(fahrenheit_temps)), + Arc::new(Float64Array::from(kelvin_temps)), + ], + )?; + + ctx.register_batch("example", batch)?; + ctx.table("example").await +} + +/// The schema of the example table. +fn example_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("city", DataType::Utf8, false), + Field::new("celsius", DataType::Float64, false) + .with_extension_type(TemperatureExtensionType(TemperatureUnit::Celsius)), + Field::new("fahrenheit", DataType::Float64, false) + .with_extension_type(TemperatureExtensionType(TemperatureUnit::Fahrenheit)), + Field::new("kelvin", DataType::Float64, false) + .with_extension_type(TemperatureExtensionType(TemperatureUnit::Kelvin)), + ])) +} + +/// Represents the unit of a temperature reading. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TemperatureUnit { + Celsius, + Fahrenheit, + Kelvin, +} + +/// Represents a float that semantically represents a temperature. The temperature can be one of +/// the supported [`TemperatureUnit`]s. +/// +/// The unit is realized as an additional extension type metadata and is stored alongside the +/// extension type name in the Arrow field metadata. This metadata can also be stored within files, +/// allowing DataFusion to read temperature data from, for example, Parquet files. +/// +/// The field metadata for a Celsius temperature field will look like this (serialized as JSON): +/// ```json +/// { +/// "ARROW:extension:name": "custom.temperature", +/// "ARROW:extension:metadata": "celsius" +/// } +/// ``` +/// +/// See [the official Arrow documentation](https://arrow.apache.org/docs/format/Columnar.html#extension-types) +/// for more details on the extension type mechanism. +#[derive(Debug)] +pub struct TemperatureExtensionType(TemperatureUnit); + +/// Implementation of [`ExtensionType`] for [`TemperatureExtensionType`]. +impl ExtensionType for TemperatureExtensionType { + /// Arrow extension type name that is stored in the `ARROW:extension:name` field. + const NAME: &'static str = "custom.temperature"; + type Metadata = TemperatureUnit; + + fn metadata(&self) -> &Self::Metadata { + &self.0 + } + + /// Arrow extension type metadata is encoded as a string and stored in the + /// `ARROW:extension:metadata` field. As we only store the name of the unit, a simple string Review Comment: TODO find correct word instead of Field as the term has a concrete meaning in arrow-rs ########## datafusion/expr/src/registry.rs: ########## @@ -215,3 +218,281 @@ impl FunctionRegistry for MemoryFunctionRegistry { self.udwfs.keys().cloned().collect() } } + +/// A cheaply cloneable pointer to an [ExtensionTypeRegistration]. +pub type ExtensionTypeRegistrationRef = Arc<dyn ExtensionTypeRegistration>; + +/// The registration of an extension type. Implementations of this trait are responsible for +/// *creating* instances of [`DFExtensionType`] that represent the entire semantics of an extension +/// type. +/// +/// # Why do we need a Registration? +/// +/// A good question is why this trait is even necessary. Why not directly register the +/// [`DFExtensionType`] in a registration? +/// +/// While this works for extension types without parameters (e.g., `arrow.uuid`), it does not work +/// for more complex extension types that may have another extension type as a parameter. For Review Comment: TODO Have additional metadata ########## datafusion-examples/examples/extension_types/temperature.rs: ########## @@ -0,0 +1,268 @@ +// 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 arrow::array::{Array, Float64Array, RecordBatch, StringArray}; +use arrow::util::display::{ArrayFormatter, DisplayIndex, FormatOptions, FormatResult}; +use arrow_schema::extension::ExtensionType; +use arrow_schema::{ArrowError, DataType, Field, Schema, SchemaRef}; +use datafusion::dataframe::DataFrame; +use datafusion::error::Result; +use datafusion::execution::SessionStateBuilder; +use datafusion::prelude::SessionContext; +use datafusion_common::internal_err; +use datafusion_common::types::DFExtensionType; +use datafusion_expr::registry::{ + DefaultExtensionTypeRegistration, ExtensionTypeRegistry, MemoryExtensionTypeRegistry, +}; +use std::fmt::Write; +use std::sync::Arc; + +/// This example demonstrates using DataFusion's extension type API to create a custom +/// semantic type [`TemperatureExtensionType`]. +pub async fn temperature_example() -> Result<()> { + let ctx = create_session_context()?; + register_temperature_table(&ctx).await?; + + // Print the example table with the custom pretty-printer. + ctx.table("example").await?.show().await +} + +/// Creates the DataFusion session context with the custom extension type implementation. +fn create_session_context() -> Result<SessionContext> { + let registry = MemoryExtensionTypeRegistry::new_empty(); + + // The registration creates a new instance of the extension type with the deserialized metadata. + let temp_registration = DefaultExtensionTypeRegistration::new_arc(|metadata| { + Ok(TemperatureExtensionType(metadata)) + }); + registry.add_extension_type_registration(temp_registration)?; + + let state = SessionStateBuilder::default() + .with_extension_type_registry(Arc::new(registry)) + .build(); + Ok(SessionContext::new_with_state(state)) +} + +/// Registers the example table and returns the data frame. +async fn register_temperature_table(ctx: &SessionContext) -> Result<DataFrame> { + let schema = example_schema(); + + let city_names = Arc::new(StringArray::from(vec![ + "Vienna", "Tokyo", "New York", "Sydney", + ])); + + // We'll use the same raw float values across columns to show how the + // extension type changes the formatting behavior. + let celsius_temps = vec![15.1, 22.5, 18.98, 25.0]; + let fahrenheit_temps = vec![59.18, 72.5, 66.164, 77.0]; + let kelvin_temps = vec![288.25, 295.65, 292.13, 298.15]; + + let batch = RecordBatch::try_new( + schema, + vec![ + city_names, + Arc::new(Float64Array::from(celsius_temps)), + Arc::new(Float64Array::from(fahrenheit_temps)), + Arc::new(Float64Array::from(kelvin_temps)), + ], + )?; + + ctx.register_batch("example", batch)?; + ctx.table("example").await +} + +/// The schema of the example table. +fn example_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("city", DataType::Utf8, false), + Field::new("celsius", DataType::Float64, false) + .with_extension_type(TemperatureExtensionType(TemperatureUnit::Celsius)), + Field::new("fahrenheit", DataType::Float64, false) + .with_extension_type(TemperatureExtensionType(TemperatureUnit::Fahrenheit)), + Field::new("kelvin", DataType::Float64, false) + .with_extension_type(TemperatureExtensionType(TemperatureUnit::Kelvin)), + ])) +} + +/// Represents the unit of a temperature reading. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TemperatureUnit { + Celsius, + Fahrenheit, + Kelvin, +} + +/// Represents a float that semantically represents a temperature. The temperature can be one of +/// the supported [`TemperatureUnit`]s. +/// +/// The unit is realized as an additional extension type metadata and is stored alongside the +/// extension type name in the Arrow field metadata. This metadata can also be stored within files, +/// allowing DataFusion to read temperature data from, for example, Parquet files. +/// +/// The field metadata for a Celsius temperature field will look like this (serialized as JSON): +/// ```json +/// { +/// "ARROW:extension:name": "custom.temperature", +/// "ARROW:extension:metadata": "celsius" +/// } +/// ``` +/// +/// See [the official Arrow documentation](https://arrow.apache.org/docs/format/Columnar.html#extension-types) +/// for more details on the extension type mechanism. +#[derive(Debug)] +pub struct TemperatureExtensionType(TemperatureUnit); + +/// Implementation of [`ExtensionType`] for [`TemperatureExtensionType`]. +impl ExtensionType for TemperatureExtensionType { + /// Arrow extension type name that is stored in the `ARROW:extension:name` field. + const NAME: &'static str = "custom.temperature"; + type Metadata = TemperatureUnit; + + fn metadata(&self) -> &Self::Metadata { + &self.0 + } + + /// Arrow extension type metadata is encoded as a string and stored in the + /// `ARROW:extension:metadata` field. As we only store the name of the unit, a simple string + /// suffices. Extension types can store more complex metadata using serialization formats like + /// JSON. + fn serialize_metadata(&self) -> Option<String> { + let s = match self.0 { + TemperatureUnit::Celsius => "celsius", + TemperatureUnit::Fahrenheit => "fahrenheit", + TemperatureUnit::Kelvin => "kelvin", + }; + Some(s.to_string()) + } + + fn deserialize_metadata( + metadata: Option<&str>, + ) -> std::result::Result<Self::Metadata, ArrowError> { + match metadata { + Some("celsius") => Ok(TemperatureUnit::Celsius), + Some("fahrenheit") => Ok(TemperatureUnit::Fahrenheit), + Some("kelvin") => Ok(TemperatureUnit::Kelvin), + Some(other) => Err(ArrowError::InvalidArgumentError(format!( + "Invalid metadata for temperature type: {other}" + ))), + None => Err(ArrowError::InvalidArgumentError( + "Temperature type requires metadata (unit)".to_owned(), + )), + } + } + + fn supports_data_type( + &self, + data_type: &DataType, + ) -> std::result::Result<(), ArrowError> { + match data_type { + DataType::Float64 => Ok(()), + _ => Err(ArrowError::InvalidArgumentError(format!( + "Invalid data type: {data_type} for temperature type, expected Float64", + ))), + } + } + + fn try_new( + data_type: &DataType, + metadata: Self::Metadata, + ) -> std::result::Result<Self, ArrowError> { + let instance = Self(metadata); + instance.supports_data_type(data_type)?; + Ok(instance) + } +} + +/// Implementation of [`DFExtensionType`] for [`TemperatureExtensionType`]. Review Comment: TODO: 1-2 sentences that this is the DataFusion API ########## datafusion/expr/src/registry.rs: ########## @@ -215,3 +218,281 @@ impl FunctionRegistry for MemoryFunctionRegistry { self.udwfs.keys().cloned().collect() } } + +/// A cheaply cloneable pointer to an [ExtensionTypeRegistration]. +pub type ExtensionTypeRegistrationRef = Arc<dyn ExtensionTypeRegistration>; + +/// The registration of an extension type. Implementations of this trait are responsible for +/// *creating* instances of [`DFExtensionType`] that represent the entire semantics of an extension +/// type. +/// +/// # Why do we need a Registration? +/// +/// A good question is why this trait is even necessary. Why not directly register the +/// [`DFExtensionType`] in a registration? +/// +/// While this works for extension types without parameters (e.g., `arrow.uuid`), it does not work +/// for more complex extension types that may have another extension type as a parameter. For +/// example, consider an extension type `custom.shortened(n)` that aims to short the pretty-printing +/// string to `n` characters. Here, `n` is a parameter of the extension type and should be a field +/// in the concrete struct that implements the [`DFExtensionType`]. The job of the registration is +/// to read the metadata from the field and create the corresponding [`DFExtensionType`] instance +/// with the correct `n` set. Review Comment: TODO: Add Self vs self analogy ########## datafusion/common/src/types/canonical_extensions/uuid.rs: ########## @@ -0,0 +1,95 @@ +// 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 crate::error::_internal_err; +use crate::types::extension::DFExtensionType; +use arrow::array::{Array, FixedSizeBinaryArray}; +use arrow::datatypes::DataType; +use arrow::util::display::{ArrayFormatter, DisplayIndex, FormatOptions, FormatResult}; +use std::fmt::Write; +use uuid::{Bytes, Uuid}; + +/// Defines the extension type logic for the canonical `arrow.uuid` extension type. +/// +/// See [`DFExtensionType`] for information on DataFusion's extension type mechanism. +impl DFExtensionType for arrow_schema::extension::Uuid { Review Comment: TODO: Create follow-up issue for adding all canonical extension types -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
