Add union_extract scalar function #12116

gstvg · 2024-08-22T18:02:24Z

Which issue does this PR close?

What changes are included in this PR?

union_extract implementation and docs
Add tables with union column on the session context of sqllogictests

Are these changes tested?

Yes. There's sqllogictests for one sparse and one dense union, and for each error case
Since it's not possible yet to create scalar unions in SQL, those cases are tested as unit tests

Are there any user-facing changes?

A new union_extract scalar function, and it's docs

…n_extract

gstvg · 2024-08-22T18:11:32Z

datafusion/functions/src/core/union_extract.rs

+}
+
+const EQ_SCALAR_CHUNK_SIZE: usize = 512;
+


Those pub items with #[doc(hidden)] are there to be benchmarked, in case any reviewer wants to, but maybe it can be removed before merge. Just let me know

Great! I think we could remove them.

Removed on arrow-rs PR

gstvg · 2024-08-22T18:15:35Z

datafusion/functions/benches/union_extract.rs

+
+fn criterion_benchmark(c: &mut Criterion) {
+    let union_extract = UnionExtractFun::new();
+


Although most part of these instantiations are very repetitive, the differences between them are so subtle that I personally prefer to be as explicit as possible. Just ask for simplification if that's preferable.

I'm not sure do we even need this 🤔

If you aren't sure, I'm more than happy to not even need to include them on the arrow-rs PR 😀

gstvg · 2024-08-22T18:18:20Z

datafusion/sqllogictest/src/test_context.rs

+    ctx.register_batch(table_name, batch).unwrap();
+}
+
+fn sparse_1_1_single_field(ctx: &SessionContext) {


see other comment

…n_extract

samuelcolvin

I haven't gone through the implementation line by line, but overall this looks pretty good to me.

samuelcolvin · 2024-08-23T15:59:45Z

datafusion/functions/src/core/union_extract.rs

+impl UnionExtractFun {
+    pub fn new() -> Self {
+        Self {
+            signature: Signature::any(2, Volatility::Immutable),


Can you be more specific than this?

Maybe, I will take a look at Signature::UserDefnied and ScalarUDFImpl::coerce_types

I pushed using ScalarUDFImpl::coerce_types in case you wanna take a look
It changed the error messages, from this to this, but I'm not sure if it's better

samuelcolvin · 2024-08-23T16:00:54Z

datafusion/functions/src/core/union_extract.rs

+            );
+        }
+
+        let fields = if let DataType::Union(fields, _) = &arg_types[0] {


nicer to write this like

let DataType::Union(fields, _) = &arg.types[0] else { return exec_err!(... };

Cool, this syntax wasn't stable when I started rust, didn't know about it. Thanks!
Done

samuelcolvin · 2024-08-23T16:01:33Z

datafusion/functions/src/core/union_extract.rs

+            );
+        };
+
+        let field_name = if let Expr::Literal(ScalarValue::Utf8(Some(field_name))) =


jayzhan211 · 2024-08-27T00:26:14Z

datafusion/functions/src/core/union_extract.rs

+
+// This is like MutableBuffer::collect_bool(type_ids.len(), |i| type_ids[i] == target) with fast paths for all true or all false values.
+#[doc(hidden)]
+pub fn eq_scalar_generic<const N: usize>(type_ids: &[i8], target: i8) -> BoolValue {


Couldn't we make them private?

I made them private on arrow-rs

jayzhan211 · 2024-08-27T00:35:47Z

datafusion/sqllogictest/test_files/union_datatype.slt

@@ -0,0 +1,255 @@
+# Licensed to the Apache Software Foundation (ASF) under one


Maybe name with union_function.slt?

Looks better
Done

jayzhan211 · 2024-08-27T00:37:37Z

datafusion/sqllogictest/src/test_context.rs

@@ -361,3 +369,893 @@ fn create_example_udf() -> ScalarUDF {
        adder,
    )
 }
+
+fn register_union_tables(ctx: &SessionContext) {


Can't we create table in slt file?

Until #10206 is fixed, I don't think so. Anyway, with the core of this PR moved to arrow-rs, all these tests are unit tests on arrow-rs now 🙏

jayzhan211 · 2024-08-27T00:44:03Z

datafusion/functions/src/core/union_extract.rs

+                            .build_unchecked()
+                    };
+
+                    Ok(ColumnarValue::Array(make_array(data)))


The code here seems equivalent to

new_null_array(target.data_type(), target.len());

Semantically yes, but new_null_array also allocates data buffers

gstvg · 2024-09-11T23:55:58Z

Thanks @samuelcolvin and @jayzhan211 for your reviews. I moved part of this PR to apache/arrow-rs#6387. I made all your requests that were applicable there, and then I will make the requests that still apply here.
Sorry for any review work hours that became unused after this 😅

gstvg added 2 commits August 22, 2024 17:49

feat: add union_extract scalar function

3faf768

Merge branch 'main' of https://github.com/apache/datafusion into unio…

b65309c

…n_extract

github-actions bot added documentation Improvements or additions to documentation sqllogictest SQL Logic Tests (.slt) functions labels Aug 22, 2024

gstvg commented Aug 22, 2024

View reviewed changes

gstvg mentioned this pull request Aug 22, 2024

Casting to and from unions apache/arrow-rs#6247

Open

gstvg added 2 commits August 22, 2024 22:46

Merge branch 'main' of https://github.com/apache/datafusion into unio…

1740dfe

…n_extract

fix: docs fmt, add clippy atr, sql error msg

4f34083

samuelcolvin reviewed Aug 23, 2024

View reviewed changes

gstvg marked this pull request as draft August 23, 2024 18:20

jayzhan211 reviewed Aug 27, 2024

View reviewed changes

gstvg added 2 commits September 12, 2024 13:05

use arrow-rs implementation

7250a2a

merge upstream

dc31a3a

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add union_extract scalar function #12116

Add union_extract scalar function #12116

gstvg commented Aug 22, 2024 •

edited

Loading

gstvg Aug 22, 2024

jayzhan211 Aug 27, 2024

gstvg Sep 12, 2024

gstvg Aug 22, 2024 •

edited

Loading

jayzhan211 Aug 27, 2024

gstvg Sep 12, 2024

gstvg Aug 22, 2024

samuelcolvin left a comment

samuelcolvin Aug 23, 2024

gstvg Aug 23, 2024

gstvg Sep 12, 2024 •

edited

Loading

samuelcolvin Aug 23, 2024

gstvg Sep 12, 2024

samuelcolvin Aug 23, 2024

gstvg Sep 12, 2024

jayzhan211 Aug 27, 2024

gstvg Sep 11, 2024

jayzhan211 Aug 27, 2024 •

edited

Loading

gstvg Sep 12, 2024

jayzhan211 Aug 27, 2024

gstvg Sep 11, 2024

jayzhan211 Aug 27, 2024

gstvg Sep 11, 2024

gstvg commented Sep 11, 2024


		fn criterion_benchmark(c: &mut Criterion) {
		let union_extract = UnionExtractFun::new();

		@@ -0,0 +1,255 @@
		# Licensed to the Apache Software Foundation (ASF) under one

Add union_extract scalar function #12116

Are you sure you want to change the base?

Add union_extract scalar function #12116

Conversation

gstvg commented Aug 22, 2024 • edited Loading

Which issue does this PR close?

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

gstvg Aug 22, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

samuelcolvin left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

gstvg Sep 12, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jayzhan211 Aug 27, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

gstvg commented Sep 11, 2024

gstvg commented Aug 22, 2024 •

edited

Loading

gstvg Aug 22, 2024 •

edited

Loading

gstvg Sep 12, 2024 •

edited

Loading

jayzhan211 Aug 27, 2024 •

edited

Loading