diff --git a/datafusion/catalog/src/information_schema.rs b/datafusion/catalog/src/information_schema.rs index e9e6f2d9ae750..bd92a72c4f24c 100644 --- a/datafusion/catalog/src/information_schema.rs +++ b/datafusion/catalog/src/information_schema.rs @@ -36,6 +36,7 @@ use datafusion_expr::logical_plan::information_schema::{ pub use datafusion_expr::logical_plan::information_schema::{ InformationSchemataBuilder, schemata_schema, }; +use datafusion_expr::registry::ExtensionTypeRegistryRef; use datafusion_expr::type_coercion::functions::fields_with_udf; use datafusion_expr::{ AggregateUDF, Documentation, ReturnFieldArgs, ScalarUDF, Signature, TypeSignature, @@ -93,6 +94,7 @@ impl InformationSchemaProvider { catalog_list, catalog_filter: None, table_functions: HashMap::new(), + extension_types: None, }, } } @@ -114,6 +116,18 @@ impl InformationSchemaProvider { self } + /// Attach the session's extension type registry so that `DESCRIBE`, + /// `SHOW COLUMNS` and `information_schema.columns` show the + /// [`display_name`](datafusion_common::types::DFExtensionType::display_name) + /// of registered extension types. + pub fn with_extension_type_registry( + mut self, + extension_types: ExtensionTypeRegistryRef, + ) -> Self { + self.config.extension_types = Some(extension_types); + self + } + /// Restricts the reported objects to the given catalog. When `None`, all /// catalogs are reported. pub fn with_catalog_filter(mut self, catalog_name: Option) -> Self { @@ -136,7 +150,8 @@ impl InformationSchemaProvider { /// Creates an execution plan for a `DESCRIBE ` statement, where /// `schema` is the output schema of the query. pub fn describe_schema(&self, schema: &Schema) -> Result> { - let mut builder = InformationSchemaColumnsBuilder::new(ColumnsDetail::Describe); + let mut builder = InformationSchemaColumnsBuilder::new(ColumnsDetail::Describe) + .with_extension_type_registry(self.config.extension_types.clone()); for (position, field) in schema.fields().iter().enumerate() { builder.add_column("", "", "", position, field); } @@ -152,7 +167,8 @@ impl InformationSchemaProvider { ) -> Result> { let table = self.resolve_table(table_ref).await?; - let mut builder = InformationSchemaColumnsBuilder::new(detail); + let mut builder = InformationSchemaColumnsBuilder::new(detail) + .with_extension_type_registry(self.config.extension_types.clone()); for (position, field) in table.schema().fields().iter().enumerate() { builder.add_column( &table_ref.catalog, @@ -407,6 +423,7 @@ struct InformationSchemaConfig { catalog_list: Arc, catalog_filter: Option, table_functions: HashMap>, + extension_types: Option, } impl InformationSchemaConfig { @@ -1041,7 +1058,8 @@ impl PartitionStream for InformationSchemaColumns { } fn execute(&self, _ctx: Arc) -> SendableRecordBatchStream { - let mut builder = InformationSchemaColumnsBuilder::new(ColumnsDetail::Full); + let mut builder = InformationSchemaColumnsBuilder::new(ColumnsDetail::Full) + .with_extension_type_registry(self.config.extension_types.clone()); let config = self.config.clone(); Box::pin(RecordBatchStreamAdapter::new( Arc::clone(&self.schema), @@ -1223,6 +1241,7 @@ mod tests { catalog_list: Arc::new(Fixture), catalog_filter: None, table_functions: HashMap::new(), + extension_types: None, }; let mut builder = InformationSchemaTablesBuilder::new(); diff --git a/datafusion/common/src/types/canonical_extensions/json.rs b/datafusion/common/src/types/canonical_extensions/json.rs index 6d7dddfef32f8..9012313dc11d9 100644 --- a/datafusion/common/src/types/canonical_extensions/json.rs +++ b/datafusion/common/src/types/canonical_extensions/json.rs @@ -6,6 +6,7 @@ use crate::Result; use crate::types::extension::DFExtensionType; use arrow::datatypes::DataType; use arrow_schema::extension::{ExtensionType, Json}; +use std::borrow::Cow; /// Defines the extension type logic for the canonical `arrow.json` extension type. This extension /// type defines that a particular string field stores JSON values. @@ -42,4 +43,8 @@ impl DFExtensionType for DFJson { fn serialize_metadata(&self) -> Option { self.inner.serialize_metadata() } + + fn display_name(&self) -> Option> { + Some(Cow::Borrowed("JSON")) + } } diff --git a/datafusion/common/src/types/canonical_extensions/uuid.rs b/datafusion/common/src/types/canonical_extensions/uuid.rs index 38ae91190d2ae..e2cbe9f29a8ac 100644 --- a/datafusion/common/src/types/canonical_extensions/uuid.rs +++ b/datafusion/common/src/types/canonical_extensions/uuid.rs @@ -9,6 +9,7 @@ use arrow::array::{Array, FixedSizeBinaryArray}; use arrow::datatypes::DataType; use arrow::util::display::{ArrayFormatter, DisplayIndex, FormatOptions, FormatResult}; use arrow_schema::extension::{ExtensionType, Uuid}; +use std::borrow::Cow; use std::fmt::Write; use uuid::Bytes; @@ -43,6 +44,10 @@ impl DFExtensionType for DFUuid { self.0.serialize_metadata() } + fn display_name(&self) -> Option> { + Some(Cow::Borrowed("UUID")) + } + fn create_array_formatter<'fmt>( &self, array: &'fmt dyn Array, diff --git a/datafusion/common/src/types/extension.rs b/datafusion/common/src/types/extension.rs index d371ec4ed9793..160dc834bd3de 100644 --- a/datafusion/common/src/types/extension.rs +++ b/datafusion/common/src/types/extension.rs @@ -6,6 +6,7 @@ use crate::error::Result; use arrow::array::Array; use arrow::util::display::{ArrayFormatter, FormatOptions}; use arrow_schema::DataType; +use std::borrow::Cow; use std::fmt::Debug; use std::sync::Arc; @@ -30,6 +31,7 @@ pub type DFExtensionTypeRef = Arc; /// Note that DataFusion's extension type support is still young and therefore might not cover all /// relevant use cases. Currently, the following operations can be customized: /// - Pretty-printing values in record batches +/// - The type name shown by schema introspection (`DESCRIBE`, `information_schema.columns`) /// /// # Relation to Arrow's [`ExtensionType`](arrow_schema::extension::ExtensionType) /// @@ -63,6 +65,16 @@ pub trait DFExtensionType: Debug + Send + Sync { /// Returns the serialized metadata. fn serialize_metadata(&self) -> Option; + /// Returns the type's presentation label for schema introspection, such as + /// `DESCRIBE`, `SHOW COLUMNS` and `information_schema.columns`. + /// + /// Names may borrow from this instance or be computed from its metadata. + /// The label is not a unique type identifier or SQL syntax guaranteed to round-trip. + /// If `None` is returned, the storage type is shown. + fn display_name(&self) -> Option> { + None + } + /// Returns an [`ArrayFormatter`] that can format values of this type. /// /// If `Ok(None)` is returned, the default implementation will be used. diff --git a/datafusion/core/src/execution/session_state.rs b/datafusion/core/src/execution/session_state.rs index 0149b38a1f32a..aeaf80984f8e3 100644 --- a/datafusion/core/src/execution/session_state.rs +++ b/datafusion/core/src/execution/session_state.rs @@ -419,7 +419,10 @@ impl SessionState { let information_schema_provider = InformationSchemaProvider::new(Arc::clone(self.catalog_list())) .with_information_schema_enabled(true) - .with_table_functions(self.inner.table_functions.clone()); + .with_table_functions(self.inner.table_functions.clone()) + .with_extension_type_registry(Arc::clone( + &self.inner.extension_types, + )); // If partitioned information schema is enabled, restrict the information schema to // the containing catalog diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 198e5bb7412b2..dd49e4ed6a673 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -2354,7 +2354,9 @@ fn create_information_schema_provider( // Do not set a catalog filter. `SHOW` statements describe all catalogs, // regardless of `SessionConfig::partition_information_schema` let provider = InformationSchemaProvider::new(session.catalog_list()); - provider.with_information_schema_enabled(session.config().information_schema()) + provider + .with_information_schema_enabled(session.config().information_schema()) + .with_extension_type_registry(Arc::clone(session.extension_type_registry())) } /// Expand and align a GROUPING SET expression. diff --git a/datafusion/expr/src/logical_plan/information_schema.rs b/datafusion/expr/src/logical_plan/information_schema.rs index fa201f618ed94..9aa1f95233696 100644 --- a/datafusion/expr/src/logical_plan/information_schema.rs +++ b/datafusion/expr/src/logical_plan/information_schema.rs @@ -6,12 +6,14 @@ //! tables and `SHOW` statements. use crate::TableType; +use crate::registry::{ExtensionTypeRegistry, ExtensionTypeRegistryRef}; use arrow::array::ArrayRef; use arrow::array::builder::{BooleanBuilder, StringBuilder, UInt8Builder, UInt64Builder}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use datafusion_common::Result; use datafusion_common::config::ConfigEntry; +use std::fmt::Write; use std::sync::Arc; /// Returns the Arrow schema of `information_schema.tables` rows. @@ -230,6 +232,7 @@ pub struct InformationSchemaColumnsBuilder { numeric_scales: Option, datetime_precisions: Option, interval_types: Option, + extension_types: Option, } impl InformationSchemaColumnsBuilder { @@ -269,9 +272,21 @@ impl InformationSchemaColumnsBuilder { datetime_precisions: full .then(|| UInt64Builder::with_capacity(default_capacity)), interval_types: full.then(StringBuilder::new), + extension_types: None, } } + /// Use the extension type registry to show the + /// [`display_name`](datafusion_common::types::DFExtensionType::display_name) + /// of registered extension types in `data_type`. + pub fn with_extension_type_registry( + mut self, + extension_types: Option, + ) -> Self { + self.extension_types = extension_types; + self + } + /// Append one row describing `field`. Values not needed at this builder's /// detail level are discarded. pub fn add_column( @@ -310,8 +325,13 @@ impl InformationSchemaColumnsBuilder { let nullable_str = if field.is_nullable() { "YES" } else { "NO" }; self.is_nullables.append_value(nullable_str); - // "System supplied type" --> Use display format of the datatype - self.data_types.append_value(field.data_type().to_string()); + // "System supplied type" --> Use the display name of the extension type or the + // display format of the datatype + append_field_type_display_name( + self.extension_types.as_deref(), + field, + &mut self.data_types, + ); // "If data_type identifies a character or bit string type, the // declared maximum length; null for all other data types or @@ -435,6 +455,35 @@ impl InformationSchemaColumnsBuilder { } } +/// Appends the name of the type of `field` shown in `data_type` to `builder`. +/// +/// If the field has an extension type registered in `registry`, appends +/// [`display_name`] if possible. appended. Otherwise, the storage type (raw +/// Arrow physical type name) is appended. +/// +/// [`display_name`]: datafusion_common::types::DFExtensionType::display_name +fn append_field_type_display_name( + registry: Option<&dyn ExtensionTypeRegistry>, + field: &Field, + builder: &mut StringBuilder, +) { + let extension_type = registry + .and_then(|registry| registry.create_extension_type_for_field(field).ok()) + .flatten(); + match extension_type + .as_ref() + .and_then(|extension_type| extension_type.display_name()) + { + Some(name) => builder.append_value(name), + None => { + // Write the storage type straight into the builder, then finish the row. + // Writing to a `StringBuilder` does not fail. + let _ = write!(builder, "{}", field.data_type()); + builder.append_value(""); + } + } +} + /// Returns the Arrow schema of `information_schema.schemata` rows. /// /// Columns and nullability match @@ -827,3 +876,56 @@ pub fn show_functions_schema() -> Schema { Field::new("syntax_example", DataType::Utf8, true), ]) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::registry::MemoryExtensionTypeRegistry; + use arrow::array::Array; + use arrow_schema::extension::{EXTENSION_TYPE_NAME_KEY, Uuid}; + use std::collections::HashMap; + + fn display_names( + registry: Option<&dyn ExtensionTypeRegistry>, + fields: &[Field], + ) -> Vec { + let mut builder = StringBuilder::new(); + for field in fields { + append_field_type_display_name(registry, field, &mut builder); + } + let array = builder.finish(); + (0..array.len()) + .map(|i| array.value(i).to_owned()) + .collect() + } + + /// Fields whose extension type cannot be resolved show their storage type. The + /// resolvable cases are covered by `sql_extension_types.slt`. + #[test] + fn test_append_field_type_display_name_fallbacks() { + let registry = MemoryExtensionTypeRegistry::new_with_canonical_extension_types(); + + let unregistered = Field::new("x", DataType::Int32, true).with_metadata( + HashMap::from([(EXTENSION_TYPE_NAME_KEY.to_owned(), "example.x".to_owned())]), + ); + // `arrow.uuid` with a storage type that the extension type does not support + let invalid = + Field::new("i", DataType::Int32, true).with_metadata(HashMap::from([( + EXTENSION_TYPE_NAME_KEY.to_owned(), + "arrow.uuid".to_owned(), + )])); + assert_eq!( + display_names(Some(®istry), &[unregistered, invalid]), + ["Int32", "Int32"] + ); + + let uuid = Field::new("u", DataType::FixedSizeBinary(16), true) + .with_extension_type(Uuid); + let empty_registry = MemoryExtensionTypeRegistry::new_empty(); + assert_eq!( + display_names(Some(&empty_registry), std::slice::from_ref(&uuid)), + ["FixedSizeBinary(16)"] + ); + assert_eq!(display_names(None, &[uuid]), ["FixedSizeBinary(16)"]); + } +} diff --git a/datafusion/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs index a9a3d7629171c..f4df452abdce2 100644 --- a/datafusion/sqllogictest/src/test_context.rs +++ b/datafusion/sqllogictest/src/test_context.rs @@ -27,6 +27,7 @@ use datafusion::common::stats::Precision; use datafusion::common::{DataFusionError, Result, not_impl_err}; use datafusion::functions::math::abs; use datafusion::logical_expr::async_udf::{AsyncScalarUDF, AsyncScalarUDFImpl}; +use datafusion::logical_expr::registry::MemoryExtensionTypeRegistry; use datafusion::logical_expr::{ ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, create_udf, @@ -116,6 +117,12 @@ impl TestContext { } } + if file_name == "sql_extension_types.slt" { + state_builder = state_builder.with_extension_type_registry(Arc::new( + MemoryExtensionTypeRegistry::new_with_canonical_extension_types(), + )); + } + if matches!( relative_path.file_name().and_then(|name| name.to_str()), Some("statistics_registry.slt") diff --git a/datafusion/sqllogictest/test_files/sql_extension_types.slt b/datafusion/sqllogictest/test_files/sql_extension_types.slt index 1d81fee43e58b..957f3dedfc09e 100644 --- a/datafusion/sqllogictest/test_files/sql_extension_types.slt +++ b/datafusion/sqllogictest/test_files/sql_extension_types.slt @@ -4,6 +4,9 @@ # Tests for built-in mapping of SQL types to Arrow canonical extension types. # See: https://arrow.apache.org/docs/format/CanonicalExtensions.html +# +# The test context registers the canonical extension types for this file, so schema +# introspection shows their display names, such as UUID and JSON. ####### # UUID Extension Type @@ -23,7 +26,7 @@ CREATE TABLE test_uuid (value UUID); query TTT DESCRIBE test_uuid; ---- -value FixedSizeBinary(16) YES +value UUID YES # 1. NULL insertion statement ok @@ -76,7 +79,7 @@ CREATE TABLE test_json (data JSON); query TTT DESCRIBE test_json; ---- -data Utf8 YES +data JSON YES # 1. NULL insertion statement ok @@ -108,6 +111,61 @@ NULL statement ok DROP TABLE test_json; +####### +# Schema introspection +####### + +statement ok +set datafusion.catalog.information_schema = true; + +statement ok +CREATE TABLE test_ext (u UUID, j JSON, i INT NOT NULL); + +query TTTTTT +SHOW COLUMNS FROM test_ext; +---- +datafusion public test_ext u UUID YES +datafusion public test_ext j JSON YES +datafusion public test_ext i Int32 NO + +query TT +SELECT column_name, data_type +FROM information_schema.columns +WHERE table_name = 'test_ext' +ORDER BY ordinal_position; +---- +u UUID +j JSON +i Int32 + +# `DESCRIBE` of a query shows the display name of its output columns. Columns computed +# from an extension type column keep its display name only if they keep the extension type +query TTT +DESCRIBE SELECT u, arrow_cast(u, 'FixedSizeBinary(16)') AS raw, j, i FROM test_ext; +---- +u UUID YES +raw FixedSizeBinary(16) YES +j JSON YES +i Int32 NO + +statement ok +CREATE VIEW test_ext_view AS SELECT u, arrow_cast(u, 'FixedSizeBinary(16)') AS raw FROM test_ext; + +query TTT +DESCRIBE test_ext_view; +---- +u UUID YES +raw FixedSizeBinary(16) YES + +statement ok +DROP VIEW test_ext_view; + +statement ok +DROP TABLE test_ext; + +statement ok +RESET datafusion.catalog.information_schema; + ####### # Error Handling #######