feat: add fp16 support to Substrait (#18086)

## Which issue does this PR close?

- Closes #16298

## Rationale for this change

Float16 is an Arrow type. Substrait serialization for the type is
defined in
https://github.com/apache/arrow/blame/main/format/substrait/extension_types.yaml
as part of Arrow. We should support it.

This picks up where #16793
leaves off.

## What changes are included in this PR?

Support for converting DataType::Float16 to/from Substrait.
Support for converting ScalarValue::Float16 to/from Substrait.

## Are these changes tested?

Yes

## Are there any user-facing changes?

Yes.

The `SubstraitProducer` trait received a new method (`register_type`)
which downstream implementors will need to provide an implementation
for. The example custom producer has been updated with a default
implementation.

One public method that changed is
[`datafusion_substrait::logical_plan::producer::from_empty_relation`](https://docs.rs/datafusion-substrait/50.2.0/datafusion_substrait/logical_plan/producer/fn.from_empty_relation.html).
I'm not sure if that is meant to be part of the public API (for one
thing, it is undocumented, though maybe this is because it serves an
obvious purpose. It also returns a `Rel` which is a pretty internal
structure).

Author: westonpace

Cargo.lock

@@ -35,6 +35,7 @@ async-recursion = "1.0"
 async-trait = { workspace = true }
 chrono = { workspace = true }
 datafusion = { workspace = true, features = ["sql"] }
+half = { workspace = true }
 itertools = { workspace = true }
 object_store = { workspace = true }
 pbjson-types = { workspace = true }

datafusion/substrait/Cargo.toml

@@ -18,6 +18,7 @@
 use crate::logical_plan::consumer::types::from_substrait_type;
 use crate::logical_plan::consumer::utils::{next_struct_field_name, DEFAULT_TIMEZONE};
 use crate::logical_plan::consumer::SubstraitConsumer;
+use crate::variation_const::FLOAT_16_TYPE_NAME;
 #[allow(deprecated)]
 use crate::variation_const::{
     DEFAULT_CONTAINER_TYPE_VARIATION_REF, DEFAULT_TYPE_VARIATION_REF,
@@ -38,6 +39,7 @@ use datafusion::common::{
     not_impl_err, plan_err, substrait_datafusion_err, substrait_err, ScalarValue,
 };
 use datafusion::logical_expr::Expr;
+use prost::Message;
 use std::sync::Arc;
 use substrait::proto;
 use substrait::proto::expression::literal::user_defined::Val;
@@ -440,8 +442,6 @@ pub(crate) fn from_substrait_literal(
                 return Ok(value);
             }
 
-            // TODO: remove the code below once the producer has been updated
-
             // Helper function to prevent duplicating this code - can be inlined once the non-extension path is removed
             let interval_month_day_nano =
                 |user_defined: &proto::expression::literal::UserDefined| -> datafusion::common::Result<ScalarValue> {
@@ -474,6 +474,36 @@ pub(crate) fn from_substrait_literal(
                 .get(&user_defined.type_reference)
             {
                 match name.as_ref() {
+                    FLOAT_16_TYPE_NAME => {
+                        // Rules for encoding fp16 Substrait literals are defined as part of Arrow here:
+                        //
+                        // https://github.com/apache/arrow/blame/bab558061696ddc1841148d6210424b12923d48e/format/substrait/extension_types.yaml#L112
+
+                        let Some(value) = user_defined.val.as_ref() else {
+                            return substrait_err!("Float16 value is empty");
+                        };
+                        let Val::Value(value_any) = value else {
+                            return substrait_err!(
+                                "Float16 value is not a value type literal"
+                            );
+                        };
+                        if value_any.type_url != "google.protobuf.UInt32Value" {
+                            return substrait_err!(
+                                "Float16 value is not a google.protobuf.UInt32Value"
+                            );
+                        }
+                        let decoded_value =
+                            pbjson_types::UInt32Value::decode(value_any.value.clone())
+                                .map_err(|err| {
+                                    substrait_datafusion_err!(
+                                        "Failed to decode float16 value: {err}"
+                                    )
+                                })?;
+                        let u32_bytes = decoded_value.value.to_le_bytes();
+                        let f16_val =
+                            half::f16::from_le_bytes(u32_bytes[0..2].try_into().unwrap());
+                        return Ok(ScalarValue::Float16(Some(f16_val)));
+                    }
                     // Kept for backwards compatibility - producers should use IntervalCompound instead
                     #[allow(deprecated)]
                     INTERVAL_MONTH_DAY_NANO_TYPE_NAME => {

datafusion/substrait/src/logical_plan/consumer/expr/literal.rs

@@ -17,6 +17,7 @@
 
 use super::utils::{from_substrait_precision, next_struct_field_name, DEFAULT_TIMEZONE};
 use super::SubstraitConsumer;
+use crate::variation_const::FLOAT_16_TYPE_NAME;
 #[allow(deprecated)]
 use crate::variation_const::{
     DATE_32_TYPE_VARIATION_REF, DATE_64_TYPE_VARIATION_REF,
@@ -251,6 +252,7 @@ pub fn from_substrait_type(
                     match name.as_ref() {
                         // Kept for backwards compatibility, producers should use IntervalCompound instead
                         INTERVAL_MONTH_DAY_NANO_TYPE_NAME => Ok(DataType::Interval(IntervalUnit::MonthDayNano)),
+                        FLOAT_16_TYPE_NAME => Ok(DataType::Float16),
                         _ => not_impl_err!(
                                 "Unsupported Substrait user defined type with ref {} and variation {}",
                                 u.type_reference,
@@ -304,15 +306,15 @@ pub fn from_substrait_named_struct(
         })?,
         &base_schema.names,
         &mut name_idx,
-    );
+    )?;
     if name_idx != base_schema.names.len() {
         return substrait_err!(
             "Names list must match exactly to nested schema, but found {} uses for {} names",
             name_idx,
             base_schema.names.len()
         );
     }
-    DFSchema::try_from(Schema::new(fields?))
+    DFSchema::try_from(Schema::new(fields))
 }
 
 fn from_substrait_struct_type(

datafusion/substrait/src/logical_plan/consumer/types.rs

@@ -40,7 +40,7 @@ pub fn from_cast(
                 nullable: true,
                 type_variation_reference: DEFAULT_TYPE_VARIATION_REF,
                 literal_type: Some(LiteralType::Null(to_substrait_type(
-                    data_type, true,
+                    producer, data_type, true,
                 )?)),
             };
             return Ok(Expression {
@@ -51,7 +51,7 @@ pub fn from_cast(
     Ok(Expression {
         rex_type: Some(RexType::Cast(Box::new(
             substrait::proto::expression::Cast {
-                r#type: Some(to_substrait_type(data_type, true)?),
+                r#type: Some(to_substrait_type(producer, data_type, true)?),
                 input: Some(Box::new(producer.handle_expr(expr, schema)?)),
                 failure_behavior: FailureBehavior::ThrowException.into(),
             },
@@ -68,7 +68,7 @@ pub fn from_try_cast(
     Ok(Expression {
         rex_type: Some(RexType::Cast(Box::new(
             substrait::proto::expression::Cast {
-                r#type: Some(to_substrait_type(data_type, true)?),
+                r#type: Some(to_substrait_type(producer, data_type, true)?),
                 input: Some(Box::new(producer.handle_expr(expr, schema)?)),
                 failure_behavior: FailureBehavior::ReturnNull.into(),
             },
@@ -79,7 +79,9 @@ pub fn from_try_cast(
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::logical_plan::producer::to_substrait_extended_expr;
+    use crate::logical_plan::producer::{
+        to_substrait_extended_expr, DefaultSubstraitProducer,
+    };
     use datafusion::arrow::datatypes::{DataType, Field};
     use datafusion::common::DFSchema;
     use datafusion::execution::SessionStateBuilder;
@@ -92,6 +94,8 @@ mod tests {
         let empty_schema = DFSchemaRef::new(DFSchema::empty());
         let field = Field::new("out", DataType::Int32, false);
 
+        let mut producer = DefaultSubstraitProducer::new(&state);
+
         let expr = Expr::Literal(ScalarValue::Null, None)
             .cast_to(&DataType::Int32, &empty_schema)
             .unwrap();
@@ -107,7 +111,7 @@ mod tests {
                 nullable: true,
                 type_variation_reference: DEFAULT_TYPE_VARIATION_REF,
                 literal_type: Some(LiteralType::Null(
-                    to_substrait_type(&DataType::Int32, true).unwrap(),
+                    to_substrait_type(&mut producer, &DataType::Int32, true).unwrap(),
                 )),
             };
             let expected = Expression {
@@ -131,13 +135,16 @@ mod tests {
             typed_null.referred_expr[0].expr_type.as_ref().unwrap()
         {
             let cast_expr = substrait::proto::expression::Cast {
-                r#type: Some(to_substrait_type(&DataType::Int32, true).unwrap()),
+                r#type: Some(
+                    to_substrait_type(&mut producer, &DataType::Int32, true).unwrap(),
+                ),
                 input: Some(Box::new(Expression {
                     rex_type: Some(RexType::Literal(Literal {
                         nullable: true,
                         type_variation_reference: DEFAULT_TYPE_VARIATION_REF,
                         literal_type: Some(LiteralType::Null(
-                            to_substrait_type(&DataType::Int64, true).unwrap(),
+                            to_substrait_type(&mut producer, &DataType::Int64, true)
+                                .unwrap(),
                         )),
                     })),
                 })),

datafusion/substrait/src/logical_plan/producer/expr/cast.rs

@@ -18,7 +18,7 @@
 use crate::logical_plan::producer::{to_substrait_type, SubstraitProducer};
 use crate::variation_const::{
     DATE_32_TYPE_VARIATION_REF, DECIMAL_128_TYPE_VARIATION_REF,
-    DEFAULT_CONTAINER_TYPE_VARIATION_REF, DEFAULT_TYPE_VARIATION_REF,
+    DEFAULT_CONTAINER_TYPE_VARIATION_REF, DEFAULT_TYPE_VARIATION_REF, FLOAT_16_TYPE_NAME,
     LARGE_CONTAINER_TYPE_VARIATION_REF, TIME_32_TYPE_VARIATION_REF,
     TIME_64_TYPE_VARIATION_REF, UNSIGNED_INTEGER_TYPE_VARIATION_REF,
     VIEW_CONTAINER_TYPE_VARIATION_REF,
@@ -61,6 +61,7 @@ pub(crate) fn to_substrait_literal(
             nullable: true,
             type_variation_reference: DEFAULT_TYPE_VARIATION_REF,
             literal_type: Some(LiteralType::Null(to_substrait_type(
+                producer,
                 &value.data_type(),
                 true,
             )?)),
@@ -94,6 +95,41 @@ pub(crate) fn to_substrait_literal(
             LiteralType::I64(*n as i64),
             UNSIGNED_INTEGER_TYPE_VARIATION_REF,
         ),
+        ScalarValue::Float16(Some(f)) => {
+            // Rules for encoding fp16 Substrait literals are defined as part of Arrow here:
+            //
+            // https://github.com/apache/arrow/blame/bab558061696ddc1841148d6210424b12923d48e/format/substrait/extension_types.yaml#L112
+            //
+            // fp16 literals are encoded as user defined literals with
+            // a google.protobuf.UInt32Value message where the lower 16 bits are
+            // the fp16 value.
+            let type_anchor = producer.register_type(FLOAT_16_TYPE_NAME.to_string());
+
+            // The spec says "lower 16 bits" but neglects to mention the endianness.
+            // Let's just use little-endian for now.
+            //
+            // See https://github.com/apache/arrow/issues/47846
+            let f_bytes = f.to_le_bytes();
+            let value = u32::from_le_bytes([f_bytes[0], f_bytes[1], 0, 0]);
+
+            let value = pbjson_types::UInt32Value { value };
+            let encoded_value = prost::Message::encode_to_vec(&value);
+            (
+                LiteralType::UserDefined(
+                    substrait::proto::expression::literal::UserDefined {
+                        type_reference: type_anchor,
+                        type_parameters: vec![],
+                        val: Some(substrait::proto::expression::literal::user_defined::Val::Value(
+                            pbjson_types::Any {
+                                type_url: "google.protobuf.UInt32Value".to_string(),
+                                value: encoded_value.into(),
+                            },
+                        )),
+                    },
+                ),
+                DEFAULT_TYPE_VARIATION_REF,
+            )
+        }
         ScalarValue::Float32(Some(f)) => {
             (LiteralType::Fp32(*f), DEFAULT_TYPE_VARIATION_REF)
         }
@@ -241,7 +277,7 @@ pub(crate) fn to_substrait_literal(
         ),
         ScalarValue::Map(m) => {
             let map = if m.is_empty() || m.value(0).is_empty() {
-                let mt = to_substrait_type(m.data_type(), m.is_nullable())?;
+                let mt = to_substrait_type(producer, m.data_type(), m.is_nullable())?;
                 let mt = match mt {
                     substrait::proto::Type {
                         kind: Some(r#type::Kind::Map(mt)),
@@ -354,12 +390,13 @@ fn convert_array_to_literal_list<T: OffsetSizeTrait>(
         .collect::<datafusion::common::Result<Vec<_>>>()?;
 
     if values.is_empty() {
-        let lt = match to_substrait_type(array.data_type(), array.is_nullable())? {
-            substrait::proto::Type {
-                kind: Some(r#type::Kind::List(lt)),
-            } => lt.as_ref().to_owned(),
-            _ => unreachable!(),
-        };
+        let lt =
+            match to_substrait_type(producer, array.data_type(), array.is_nullable())? {
+                substrait::proto::Type {
+                    kind: Some(r#type::Kind::List(lt)),
+                } => lt.as_ref().to_owned(),
+                _ => unreachable!(),
+            };
         Ok(LiteralType::EmptyList(lt))
     } else {
         Ok(LiteralType::List(List { values }))

datafusion/substrait/src/logical_plan/producer/expr/literal.rs

@@ -78,7 +78,7 @@ pub fn to_substrait_extended_expr(
             })
         })
         .collect::<datafusion::common::Result<Vec<_>>>()?;
-    let substrait_schema = to_substrait_named_struct(schema)?;
+    let substrait_schema = to_substrait_named_struct(&mut producer, schema)?;
 
     let extensions = producer.get_extensions();
     Ok(Box::new(ExtendedExpression {

datafusion/substrait/src/logical_plan/producer/expr/mod.rs

@@ -36,7 +36,7 @@ pub fn to_substrait_plan(
     let plan_rels = vec![PlanRel {
         rel_type: Some(plan_rel::RelType::Root(RelRoot {
             input: Some(*producer.handle_plan(plan)?),
-            names: to_substrait_named_struct(plan.schema())?.names,
+            names: to_substrait_named_struct(&mut producer, plan.schema())?.names,
         })),
     }];
 

datafusion/substrait/src/logical_plan/producer/plan.rs

@@ -48,7 +48,7 @@ pub fn from_table_scan(
     });
 
     let table_schema = scan.source.schema().to_dfschema_ref()?;
-    let base_schema = to_substrait_named_struct(&table_schema)?;
+    let base_schema = to_substrait_named_struct(producer, &table_schema)?;
 
     let filter_option = if scan.filters.is_empty() {
         None
@@ -83,15 +83,18 @@ pub fn from_table_scan(
     }))
 }
 
-pub fn from_empty_relation(e: &EmptyRelation) -> datafusion::common::Result<Box<Rel>> {
+pub fn from_empty_relation(
+    producer: &mut impl SubstraitProducer,
+    e: &EmptyRelation,
+) -> datafusion::common::Result<Box<Rel>> {
     if e.produce_one_row {
         return not_impl_err!("Producing a row from empty relation is unsupported");
     }
     #[allow(deprecated)]
     Ok(Box::new(Rel {
         rel_type: Some(RelType::Read(Box::new(ReadRel {
             common: None,
-            base_schema: Some(to_substrait_named_struct(&e.schema)?),
+            base_schema: Some(to_substrait_named_struct(producer, &e.schema)?),
             filter: None,
             best_effort_filter: None,
             projection: None,
@@ -135,7 +138,7 @@ pub fn from_values(
     Ok(Box::new(Rel {
         rel_type: Some(RelType::Read(Box::new(ReadRel {
             common: None,
-            base_schema: Some(to_substrait_named_struct(&v.schema)?),
+            base_schema: Some(to_substrait_named_struct(producer, &v.schema)?),
             filter: None,
             best_effort_filter: None,
             projection: None,

datafusion/substrait/src/logical_plan/producer/rel/read_rel.rs

@@ -70,6 +70,10 @@ use substrait::proto::{
 ///        self.extensions.register_function(signature)
 ///     }
 ///
+///     fn register_type(&mut self, type_name: String) -> u32 {
+///         self.extensions.register_type(type_name)
+///     }
+///
 ///     fn get_extensions(self) -> Extensions {
 ///         self.extensions
 ///     }
@@ -114,6 +118,15 @@ pub trait SubstraitProducer: Send + Sync + Sized {
     /// there is one. Otherwise, it should generate a new anchor.
     fn register_function(&mut self, signature: String) -> u32;
 
+    /// Within a Substrait plan, user defined types are referenced using type anchors that are stored at
+    /// the top level of the [Plan](substrait::proto::Plan) within
+    /// [ExtensionType](substrait::proto::extensions::simple_extension_declaration::ExtensionType)
+    /// messages.
+    ///
+    /// When given a type name, this method should return the existing anchor for it if
+    /// there is one. Otherwise, it should generate a new anchor.
+    fn register_type(&mut self, name: String) -> u32;
+
     /// Consume the producer to generate the [Extensions] for the Substrait plan based on the
     /// functions that have been registered
     fn get_extensions(self) -> Extensions;
@@ -182,7 +195,7 @@ pub trait SubstraitProducer: Send + Sync + Sized {
         &mut self,
         plan: &EmptyRelation,
     ) -> datafusion::common::Result<Box<Rel>> {
-        from_empty_relation(plan)
+        from_empty_relation(self, plan)
     }
 
     fn handle_subquery_alias(
@@ -367,6 +380,10 @@ impl SubstraitProducer for DefaultSubstraitProducer<'_> {
         self.extensions.register_function(fn_name)
     }
 
+    fn register_type(&mut self, type_name: String) -> u32 {
+        self.extensions.register_type(type_name)
+    }
+
     fn get_extensions(self) -> Extensions {
         self.extensions
     }