Add `union_tag` scalar function #14687

gstvg · 2025-02-16T01:14:35Z

Which issue does this PR close?

Closes Add union_tag function #11080

Rationale for this change

Retrieve the name of the currently selected field on a union, as there's no way to do it today

What changes are included in this PR?

union_tag scalar function implementation

Are these changes tested?

Yes, with sqllogictests when possible, and with unit tests for union scalars, which are not supported in SQL yet

Are there any user-facing changes?

A new scalar function union_tag

gstvg · 2025-02-17T19:13:52Z

datafusion/functions/src/core/union_tag.rs

+                // Union fields type IDs only constraints are being unique and in the 0..128 range:
+                // They may not start at 0, be sequential, or even contiguous.
+                // Therefore, we allocate a values vector with a length equal to the highest type ID plus one,
+                // ensuring that each field's name can be placed at the index corresponding to its type ID.


The union column used on the sqllogictests contains a single field with type id 3, so this is put to the test

datafusion/datafusion/sqllogictest/src/test_context.rs

Lines 411 to 430 in e4b78c7

fn register_union_table(ctx: &SessionContext) {

let union = UnionArray::try_new(

UnionFields::new(vec![3], vec![Field::new("int", DataType::Int32, false)]),

ScalarBuffer::from(vec![3, 3]),

None,

vec![Arc::new(Int32Array::from(vec![1, 2]))],

)

.unwrap();

let schema = Schema::new(vec![Field::new(

"union_column",

union.data_type().clone(),

false,

)]);

let batch =

RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(union)]).unwrap();

ctx.register_batch("union_table", batch).unwrap();

}

datafusion/datafusion/sqllogictest/src/test_context.rs

Lines 117 to 120 in e4b78c7

"union_function.slt" => {

info!("Registering table with union column");

register_union_table(test_ctx.session_ctx())

}

Omega359 · 2025-02-18T19:01:59Z

@alamb - here is another function coming in (xxhash, regexp_extract (both versions of it), array_min/array_max functions) where it is not clear what should be accepted and what shouldn't be. Since we've already accepted union_extract this may be a case of fleshing that series of functions out.

I strongly think we need 'official' documentation as to what will be and won't be accepted and an recommended repository/sub project where additional functions can be located. Hopefully under the apache umbrella such that they hopefully can maintained by the community and work across multiple DF versions.

alamb · 2025-02-19T12:31:29Z

@alamb - here is another function coming in (xxhash, regexp_extract (both versions of it), array_min/array_max functions) where it is not clear what should be accepted and what shouldn't be. Since we've already accepted union_extract this may be a case of fleshing that series of functions out.

I strongly think we need 'official' documentation as to what will be and won't be accepted and an recommended repository/sub project where additional functions can be located. Hopefully under the apache umbrella such that they hopefully can maintained by the community and work across multiple DF versions.

Yeah I agree. I think we should file a "discussion" type ticket to have this discussion. I can file one at some point later (I am low on time this week) or if you can that would be sweet.

we have some generic guidance here: https://datafusion.apache.org/contributor-guide/index.html#what-contributions-are-good-fits

Omega359 · 2025-02-19T14:20:52Z

Yeah I agree. I think we should file a "discussion" type ticket to have this discussion. I can file one at some point later (I am low on time this week) or if you can that would be sweet.

we have some generic guidance here: https://datafusion.apache.org/contributor-guide/index.html#what-contributions-are-good-fits

on it. #14777

github-actions · 2025-04-22T02:07:01Z

Thank you for your contribution. Unfortunately, this pull request is stale because it has been open 60 days with no activity. Please remove the stale label or comment or this will be closed in 7 days.

Omega359 · 2025-04-22T14:45:19Z

@alamb, thoughts on this?

alamb

If you think we should merge this @Omega359 I am fine with doing so too.

I have merged up from main to get a clean CI run, and hopefully we'll be good.

I think this will nicely round out the UnionArray related functions

Omega359 · 2025-04-29T13:43:01Z

I'll review it today @alamb

Omega359 · 2025-04-29T21:04:46Z

LGTM. I think I'd like to see a test with multiple columns but the logic looks solid to me. I believe the use of unsafe is indeed ok given the conditions outlined.

alamb · 2025-04-29T22:38:18Z

datafusion/sqllogictest/test_files/union_function.slt

@@ -23,7 +26,8 @@ query ?I
 select union_column, union_extract(union_column, 'int') from union_table;
 ----
 {int=1} 1
-{int=2} 2
+{string=bar} NULL


I added a new row to the table so we could test union_tag with a field that did not exist, per @Omega359 's suggestion

alamb · 2025-04-29T22:38:52Z

datafusion/sqllogictest/test_files/union_function.slt

+select union_column, union_tag(union_column) from union_table;
+----
+{int=1} int
+{string=bar} string


I added a new test as suggested

alamb · 2025-04-29T22:40:09Z

LGTM. I think I'd like to see a test with multiple columns but the logic looks solid to me. I believe the use of unsafe is indeed ok given the conditions outlined.

Thanks again for the review. I added the requested test. I think the CI is going to fail because of #15149 (comment)

Once that is fixed I'll refresh the PR

alamb · 2025-05-01T02:03:30Z

Thank you @Omega359 and @gstvg for your patience

* feat: add union_tag scalar function * update for new api * Add test for second field type --------- Co-authored-by: Andrew Lamb <[email protected]>

feat: add union_tag scalar function

8d07d08

github-actions bot added documentation Improvements or additions to documentation sqllogictest SQL Logic Tests (.slt) functions Changes to functions implementation labels Feb 16, 2025

alamb mentioned this pull request Feb 17, 2025

alamb's review queue #14698

Closed

gstvg commented Feb 17, 2025

View reviewed changes

This was referenced Feb 24, 2025

Weekly Plan (Andrew Lamb) Feb 24, 2025 #14850

Closed

Weekly Plan (Andrew Lamb) March 3, 2025 #14978

Closed

alamb mentioned this pull request Mar 10, 2025

Weekly Plan (Andrew Lamb) March 10, 2025 #15121

Closed

11 tasks

github-actions bot added the Stale PR has not had any activity for some time label Apr 22, 2025

github-actions bot removed the Stale PR has not had any activity for some time label Apr 24, 2025

Merge remote-tracking branch 'apache/main' into union_tag

1e3675b

alamb changed the title ~~Add union_tag scalar function~~ Add union_tag scalar function Apr 29, 2025

alamb approved these changes Apr 29, 2025

View reviewed changes

alamb added 3 commits April 29, 2025 18:26

Merge remote-tracking branch 'apache/main' into union_tag

288a3c1

update for new api

47e785a

Add test for second field type

680e7fa

alamb reviewed Apr 29, 2025

View reviewed changes

Merge remote-tracking branch 'apache/main' into union_tag

8d9dda0

alamb merged commit a726e5a into apache:main May 1, 2025
29 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add `union_tag` scalar function #14687

Add `union_tag` scalar function #14687

gstvg commented Feb 16, 2025

gstvg Feb 17, 2025 •

edited

Loading

Omega359 commented Feb 18, 2025

alamb commented Feb 19, 2025

Omega359 commented Feb 19, 2025 •

edited

Loading

github-actions bot commented Apr 22, 2025

Omega359 commented Apr 22, 2025

alamb left a comment

Omega359 commented Apr 29, 2025

Omega359 commented Apr 29, 2025

alamb Apr 29, 2025

alamb Apr 29, 2025

alamb commented Apr 29, 2025

alamb commented May 1, 2025

	fn register_union_table(ctx: &SessionContext) {
	let union = UnionArray::try_new(
	UnionFields::new(vec![3], vec![Field::new("int", DataType::Int32, false)]),
	ScalarBuffer::from(vec![3, 3]),
	None,
	vec![Arc::new(Int32Array::from(vec![1, 2]))],
	)
	.unwrap();

	let schema = Schema::new(vec![Field::new(
	"union_column",
	union.data_type().clone(),
	false,
	)]);

	let batch =
	RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(union)]).unwrap();

	ctx.register_batch("union_table", batch).unwrap();
	}

	"union_function.slt" => {
	info!("Registering table with union column");
	register_union_table(test_ctx.session_ctx())
	}

Add union_tag scalar function #14687

Add union_tag scalar function #14687

Conversation

gstvg commented Feb 16, 2025

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

gstvg Feb 17, 2025 • edited Loading

Choose a reason for hiding this comment

Omega359 commented Feb 18, 2025

alamb commented Feb 19, 2025

Omega359 commented Feb 19, 2025 • edited Loading

github-actions bot commented Apr 22, 2025

Omega359 commented Apr 22, 2025

alamb left a comment

Choose a reason for hiding this comment

Omega359 commented Apr 29, 2025

Omega359 commented Apr 29, 2025

alamb Apr 29, 2025

Choose a reason for hiding this comment

alamb Apr 29, 2025

Choose a reason for hiding this comment

alamb commented Apr 29, 2025

alamb commented May 1, 2025

Add `union_tag` scalar function #14687

Add `union_tag` scalar function #14687

gstvg Feb 17, 2025 •

edited

Loading

Omega359 commented Feb 19, 2025 •

edited

Loading