WIP: Extension type casts using extension registry#21071
Draft
paleolimbot wants to merge 39 commits intoapache:mainfrom
Draft
WIP: Extension type casts using extension registry#21071paleolimbot wants to merge 39 commits intoapache:mainfrom
paleolimbot wants to merge 39 commits intoapache:mainfrom
Conversation
github-actions
Bot
added
logical-expr
paleolimbot
commented
Mar 20, 2026
Comment on lines
+94
to
+115
| let state = SessionStateBuilder::default() | ||
| .with_canonical_extension_types()? | ||
| .with_type_planner(Arc::new(CustomTypePlanner {})) | ||
| .build(); | ||
| let ctx = SessionContext::new_with_state(state); | ||
|
|
||
| ctx.register_batch("test", batch)?; | ||
|
|
||
| let df = ctx.sql("SELECT my_uuids::VARCHAR FROM test").await?; | ||
| let batches = df.collect().await?; | ||
|
|
||
| assert_batches_eq!( | ||
| vec![ | ||
| "+--------------------------------------+", | ||
| "| test.my_uuids |", | ||
| "+--------------------------------------+", | ||
| "| 00000000-0000-0000-0000-000000000000 |", | ||
| "| 00010203-0405-0607-0809-000102030506 |", | ||
| "+--------------------------------------+", | ||
| ], | ||
| &batches | ||
| ); |
Member
Author
There was a problem hiding this comment.
Casting from a UUID to something else works!
Comment on lines
+121
to
+141
| async fn create_cast_char_to_uuid() -> Result<()> { | ||
| let state = SessionStateBuilder::default() | ||
| .with_canonical_extension_types()? | ||
| .with_type_planner(Arc::new(CustomTypePlanner {})) | ||
| .build(); | ||
| let ctx = SessionContext::new_with_state(state); | ||
|
|
||
| let df = ctx | ||
| .sql("SELECT '00010203-0405-0607-0809-000102030506'::UUID AS uuid") | ||
| .await?; | ||
| let batches = df.collect().await?; | ||
| assert_batches_eq!( | ||
| vec![ | ||
| "+----------------------------------+", | ||
| "| uuid |", | ||
| "+----------------------------------+", | ||
| "| 00010203040506070809000102030506 |", | ||
| "+----------------------------------+", | ||
| ], | ||
| &batches | ||
| ); |
Member
Author
There was a problem hiding this comment.
We can also go the other direction!
Comment on lines
+74
to
+94
| fn cast_from(&self) -> Result<Arc<dyn CastExtension>> { | ||
| Ok(Arc::new(DefaultExtensionCast {})) | ||
| } | ||
|
|
||
| fn cast_to(&self) -> Result<Arc<dyn CastExtension>> { | ||
| Ok(Arc::new(DefaultExtensionCast {})) | ||
| } | ||
| } | ||
|
|
||
| pub trait CastExtension: Debug + Send + Sync { | ||
| fn can_cast(&self, from: &Field, to: &Field, options: &CastOptions) -> Result<bool>; | ||
|
|
||
| // None for fallback | ||
| fn cast( | ||
| &self, | ||
| value: ArrayRef, | ||
| from: &Field, | ||
| to: &Field, | ||
| options: &CastOptions, | ||
| ) -> Result<ArrayRef>; | ||
| } |
Member
Author
There was a problem hiding this comment.
This is the interface an extension type can implement to define its interactions with other types
Comment on lines
+301
to
+305
| if let Some(registry) = &execution_props.extension_types | ||
| && let Some(extension_type) = | ||
| registry.create_extension_type_for_field(&field)? | ||
| { | ||
| let cast_extension = extension_type.cast_from()?; |
Member
Author
There was a problem hiding this comment.
This is where the cast (to an extension type from something else) gets planned into a physical expression
Comment on lines
+332
to
+334
| let cast_extension = extension_type.cast_to()?; | ||
| if cast_extension.can_cast(&src_field, &field, &DEFAULT_CAST_OPTIONS)? { | ||
| return expressions::cast_with_extension( |
Member
Author
There was a problem hiding this comment.
This is where the cast from something else to an extension type is resolved to a physical expression.
(I didn't handle the case where an extension type is getting cast to another extension type, but that should get handled if this is ever going to merge because we have to make sure the interface can handle either the right OR the left side defining this cast without erroring if the other side doesn't handle it)
Comment on lines
+275
to
+286
| if let Some(cast_extension) = &self.cast_extension { | ||
| let from_field = self.expr.return_field(&batch.schema())?; | ||
| let to_field = self.return_field(&batch.schema())?; | ||
| match value { | ||
| ColumnarValue::Array(array) => { | ||
| Ok(ColumnarValue::Array(cast_extension.cast( | ||
| array, | ||
| &from_field, | ||
| &to_field, | ||
| &self.cast_options, | ||
| )?)) | ||
| } |
Member
Author
There was a problem hiding this comment.
This is where the cast is executed. I chose to tack on the CastExtension to the CastExpr but it could also be its own PhysicalExpr (maybe safer in case I haven't considered some of the things that could happen during physical optimizer passes)
paleolimbot
force-pushed
the
extension-type-registry-cast
branch
from
6326d62 to
4f82129
Compare
Co-authored-by: Copilot <copilot@github.com>
|
Thank you for opening this pull request! Reviewer note: cargo-semver-checks reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch). Details |
github-actions
Bot
removed
datasource
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
This PR is a proof of concept stacked on top of #20312 to demonstrate how casting to and from extension types might be supported with the registry design in that PR. All details up for grabs...most of the work here is just piping the registry so that we can resolve a cast to or from an extension type when creating a physical expression from a logical one (the work to pipe SQL and Logical plan casts to extension types was already done before this PR).
Rationale for this change
What changes are included in this PR?
Are these changes tested?
Are there any user-facing changes?