Repository navigation
fix: return a scalar field from a scalar struct in GetStructField #7381
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -15,7 +15,7 @@ | |
| // specific language governing permissions and limitations | ||
| // under the License. | ||
|
|
||
| use arrow::array::{Array, StructArray}; | ||
| use arrow::array::{Array, ArrayRef, AsArray, StructArray}; | ||
| use arrow::datatypes::{DataType, Field, Schema}; | ||
| use arrow::record_batch::RecordBatch; | ||
| use datafusion::common::{DataFusionError, Result as DataFusionResult, ScalarValue}; | ||
|
|
@@ -99,8 +99,10 @@ impl PhysicalExpr for GetStructField { | |
| self.ordinal, | ||
| )?)) | ||
| } | ||
| ColumnarValue::Scalar(ScalarValue::Struct(struct_array)) => Ok(ColumnarValue::Array( | ||
| child_with_parent_nulls(&struct_array, self.ordinal)?, | ||
| // A scalar struct holds one row, so its field is a scalar too. A one-row array would | ||
| // not match the batch's row count. | ||
| ColumnarValue::Scalar(ScalarValue::Struct(struct_array)) => Ok(ColumnarValue::Scalar( | ||
| single_row_scalar(child_with_parent_nulls(&struct_array, self.ordinal)?)?, | ||
| )), | ||
| value => Err(DataFusionError::Execution(format!( | ||
| "Expected a struct array, got {value:?}" | ||
|
|
@@ -123,6 +125,37 @@ impl PhysicalExpr for GetStructField { | |
| } | ||
| } | ||
|
|
||
| /// The scalar that a one-row array holds. A list wraps the array as is: `try_from_array` rebuilds | ||
| /// a list's element field from its name and nullability, which drops metadata such as a Variant | ||
| /// element's extension type, so the scalar's type would no longer match `data_type`. A dictionary | ||
| /// converts the value its key points at here for the same reason, since `try_from_array` would | ||
| /// convert it with itself. | ||
| fn single_row_scalar(array: ArrayRef) -> DataFusionResult<ScalarValue> { | ||
| Ok(match array.data_type() { | ||
| DataType::List(_) => ScalarValue::List(Arc::new(array.as_list::<i32>().clone())), | ||
| DataType::LargeList(_) => ScalarValue::LargeList(Arc::new(array.as_list::<i64>().clone())), | ||
| DataType::FixedSizeList(_, _) => { | ||
| ScalarValue::FixedSizeList(Arc::new(array.as_fixed_size_list().clone())) | ||
| } | ||
| DataType::ListView(_) => { | ||
| ScalarValue::ListView(Arc::new(array.as_list_view::<i32>().clone())) | ||
| } | ||
| DataType::LargeListView(_) => { | ||
| ScalarValue::LargeListView(Arc::new(array.as_list_view::<i64>().clone())) | ||
| } | ||
| DataType::Dictionary(key_type, _) => { | ||
| let dict = array.as_any_dictionary(); | ||
| let value = if dict.keys().is_valid(0) { | ||
| single_row_scalar(dict.values().slice(dict.normalized_keys()[0], 1))? | ||
| } else { | ||
| ScalarValue::try_from(dict.values().data_type())? | ||
| }; | ||
| ScalarValue::Dictionary(key_type.clone(), Box::new(value)) | ||
| } | ||
| _ => ScalarValue::try_from_array(&array, 0)?, | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Preserve the list element field when extracting a dictionary scalar. I used
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Thanks, I could reproduce that: before the change, |
||
| }) | ||
| } | ||
|
|
||
| impl Display for GetStructField { | ||
| fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { | ||
| write!( | ||
|
|
@@ -136,10 +169,212 @@ impl Display for GetStructField { | |
| #[cfg(test)] | ||
| mod tests { | ||
| use super::*; | ||
| use arrow::array::{ArrayRef, Int64Array}; | ||
| use arrow::buffer::NullBuffer; | ||
| use arrow::array::{ | ||
| DictionaryArray, FixedSizeListArray, Int32Array, Int64Array, LargeListArray, | ||
| LargeListViewArray, ListArray, ListViewArray, NullArray, | ||
| }; | ||
| use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; | ||
| use arrow::datatypes::Fields; | ||
| use datafusion::physical_expr::expressions::Column; | ||
| use datafusion::physical_expr::expressions::{Column, Literal}; | ||
|
|
||
| // A field of a scalar struct is a scalar, so a consumer broadcasts it over the batch instead of | ||
| // reading a one-row array beside full columns. A null struct gives a null field, for a typed | ||
| // field and a Null-typed one. | ||
| #[test] | ||
| fn field_of_scalar_struct_is_scalar() { | ||
| let fields: Fields = Fields::from(vec![ | ||
| Field::new("a", DataType::Int64, true), | ||
| Field::new("n", DataType::Null, true), | ||
| ]); | ||
| let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); | ||
| let batch = RecordBatch::try_new( | ||
| Arc::clone(&schema), | ||
| vec![Arc::new(Int64Array::from(vec![0_i64, 1, 2, 3]))], | ||
| ) | ||
| .unwrap(); | ||
|
|
||
| let present = ScalarValue::Struct(Arc::new(StructArray::new( | ||
| fields.clone(), | ||
| vec![ | ||
| Arc::new(Int64Array::from(vec![7_i64])), | ||
| Arc::new(NullArray::new(1)), | ||
| ], | ||
| None, | ||
| ))); | ||
| let absent = ScalarValue::Struct(Arc::new(StructArray::new_null(fields.clone(), 1))); | ||
| // A null struct whose child slot still holds a value: the field is null all the same. | ||
| let hidden = ScalarValue::Struct(Arc::new(StructArray::new( | ||
| fields, | ||
| vec![ | ||
| Arc::new(Int64Array::from(vec![7_i64])), | ||
| Arc::new(NullArray::new(1)), | ||
| ], | ||
| Some(NullBuffer::new_null(1)), | ||
| ))); | ||
|
|
||
| for (scalar, ordinal, expected) in [ | ||
| (present.clone(), 0, ScalarValue::Int64(Some(7))), | ||
| (present, 1, ScalarValue::Null), | ||
| (absent.clone(), 0, ScalarValue::Int64(None)), | ||
| (absent, 1, ScalarValue::Null), | ||
| (hidden, 0, ScalarValue::Int64(None)), | ||
| ] { | ||
| let expr = GetStructField::new(Arc::new(Literal::new(scalar)), ordinal); | ||
| match expr.evaluate(&batch).unwrap() { | ||
| ColumnarValue::Scalar(value) => assert_eq!(value, expected), | ||
| ColumnarValue::Array(array) => { | ||
| panic!("expected a scalar for ordinal {ordinal}, got array {array:?}") | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
||
| // A list field of a scalar struct is the struct's list, and a null struct gives a null list. | ||
| // Either way it keeps its element field, metadata included, so the scalar's type is the one | ||
| // `data_type` declares. | ||
| #[test] | ||
| fn list_field_of_scalar_struct_keeps_its_type() { | ||
| let item = | ||
| Arc::new(Field::new("item", DataType::Int64, true).with_metadata( | ||
| [("ARROW:extension:name".to_string(), "test.ext".to_string())].into(), | ||
| )); | ||
| let values: ArrayRef = Arc::new(Int64Array::from(vec![1_i64, 2])); | ||
| let lists: Vec<ArrayRef> = vec![ | ||
| Arc::new(ListArray::new( | ||
| Arc::clone(&item), | ||
| OffsetBuffer::from_lengths([2]), | ||
| Arc::clone(&values), | ||
| None, | ||
| )), | ||
| Arc::new(LargeListArray::new( | ||
| Arc::clone(&item), | ||
| OffsetBuffer::from_lengths([2]), | ||
| Arc::clone(&values), | ||
| None, | ||
| )), | ||
| Arc::new(FixedSizeListArray::new( | ||
| Arc::clone(&item), | ||
| 2, | ||
| Arc::clone(&values), | ||
| None, | ||
| )), | ||
| Arc::new(ListViewArray::new( | ||
| Arc::clone(&item), | ||
| ScalarBuffer::from(vec![0_i32]), | ||
| ScalarBuffer::from(vec![2_i32]), | ||
| Arc::clone(&values), | ||
| None, | ||
| )), | ||
| Arc::new(LargeListViewArray::new( | ||
| Arc::clone(&item), | ||
| ScalarBuffer::from(vec![0_i64]), | ||
| ScalarBuffer::from(vec![2_i64]), | ||
| Arc::clone(&values), | ||
| None, | ||
| )), | ||
| ]; | ||
| let schema = Schema::new(vec![Field::new("id", DataType::Int64, false)]); | ||
| let batch = RecordBatch::try_new( | ||
| Arc::new(schema.clone()), | ||
| vec![Arc::new(Int64Array::from(vec![0_i64, 1, 2]))], | ||
| ) | ||
| .unwrap(); | ||
|
|
||
| for list in lists { | ||
| let fields = Fields::from(vec![Field::new("l", list.data_type().clone(), true)]); | ||
| // The second struct is null although its child slot holds the list. | ||
| for struct_nulls in [None, Some(NullBuffer::new_null(1))] { | ||
| let is_null = struct_nulls.is_some(); | ||
| let scalar = ScalarValue::Struct(Arc::new(StructArray::new( | ||
| fields.clone(), | ||
| vec![Arc::clone(&list)], | ||
| struct_nulls, | ||
| ))); | ||
| let expr = GetStructField::new(Arc::new(Literal::new(scalar)), 0); | ||
| let value = match expr.evaluate(&batch).unwrap() { | ||
| ColumnarValue::Scalar(value) => value, | ||
| ColumnarValue::Array(array) => panic!("expected a scalar, got array {array:?}"), | ||
| }; | ||
| let kind = list.data_type(); | ||
| assert_eq!( | ||
| value.data_type(), | ||
| expr.data_type(&schema).unwrap(), | ||
| "{kind}" | ||
| ); | ||
| assert_eq!(value.is_null(), is_null, "{kind}, null struct: {is_null}"); | ||
| if !is_null { | ||
| assert_eq!(value.to_array().unwrap().as_ref(), list.as_ref(), "{kind}"); | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
||
| // A dictionary field of a scalar struct is the value its key points at, still dictionary | ||
| // encoded with the same key type. A list value keeps its element field, so the field broadcasts | ||
| // into a column of the type `data_type` declares. | ||
| #[test] | ||
| fn dictionary_field_of_scalar_struct_keeps_its_type() { | ||
| let item = | ||
| Arc::new(Field::new("item", DataType::Int64, true).with_metadata( | ||
| [("ARROW:extension:name".to_string(), "test.ext".to_string())].into(), | ||
| )); | ||
| let lists = ListArray::new( | ||
| Arc::clone(&item), | ||
| OffsetBuffer::from_lengths([1, 2]), | ||
| Arc::new(Int64Array::from(vec![0_i64, 1, 2])), | ||
| None, | ||
| ); | ||
| let dict = |key: Option<i32>| -> ArrayRef { | ||
| Arc::new(DictionaryArray::new( | ||
| Int32Array::from(vec![key]), | ||
| Arc::new(lists.clone()), | ||
| )) | ||
| }; | ||
| let schema = Schema::new(vec![Field::new("id", DataType::Int64, false)]); | ||
| let batch = RecordBatch::try_new( | ||
| Arc::new(schema.clone()), | ||
| vec![Arc::new(Int64Array::from(vec![0_i64, 1, 2]))], | ||
| ) | ||
| .unwrap(); | ||
|
|
||
| // The key points at the second list, a null key gives a null, and so does a null struct | ||
| // whose child slot still holds a key. | ||
| for (field, struct_nulls, expected) in [ | ||
| (dict(Some(1)), None, Some(lists.slice(1, 1))), | ||
| (dict(None), None, None), | ||
| (dict(Some(1)), Some(NullBuffer::new_null(1)), None), | ||
| ] { | ||
| let fields = Fields::from(vec![Field::new("d", field.data_type().clone(), true)]); | ||
| let scalar = ScalarValue::Struct(Arc::new(StructArray::new( | ||
| fields, | ||
| vec![field], | ||
| struct_nulls, | ||
| ))); | ||
| let expr = GetStructField::new(Arc::new(Literal::new(scalar)), 0); | ||
| let value = match expr.evaluate(&batch).unwrap() { | ||
| ColumnarValue::Scalar(value) => value, | ||
| ColumnarValue::Array(array) => panic!("expected a scalar, got array {array:?}"), | ||
| }; | ||
| let declared = Arc::new(Schema::new(vec![Field::new( | ||
| "d", | ||
| expr.data_type(&schema).unwrap(), | ||
| true, | ||
| )])); | ||
| RecordBatch::try_new( | ||
| declared, | ||
| vec![value.to_array_of_size(batch.num_rows()).unwrap()], | ||
| ) | ||
| .unwrap(); | ||
| let ScalarValue::Dictionary(key_type, inner) = value else { | ||
| panic!("expected a dictionary scalar, got {value:?}") | ||
| }; | ||
| assert_eq!(*key_type, DataType::Int32); | ||
| match expected { | ||
| Some(list) => assert_eq!(*inner, ScalarValue::List(Arc::new(list))), | ||
| None => assert!(inner.is_null(), "expected a null, got {inner:?}"), | ||
| } | ||
| } | ||
|
Comment on lines
+283
to
+376
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This loop only compares
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Done. The test now covers every list type that |
||
| } | ||
|
|
||
| // A field of a NULL struct must be NULL (Spark semantics) even when the child buffer holds a | ||
| // non-null value at that row -- Arrow stores child validity independently of the parent | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -15,6 +15,11 @@ | |
| -- specific language governing permissions and limitations | ||
| -- under the License. | ||
|
|
||
| -- PushFoldableIntoBranches is excluded for the null struct query at the end, which it would | ||
| -- rewrite so that native no longer reads a field of the struct. The other queries keep a native | ||
| -- GetStructField either way. | ||
| -- ExcludeRules: org.apache.spark.sql.catalyst.optimizer.PushFoldableIntoBranches | ||
|
|
||
| statement | ||
| CREATE TABLE test_struct(s struct<name: string, age: int, score: double>) USING parquet | ||
|
|
||
|
|
@@ -38,3 +43,33 @@ INSERT INTO test_struct_null_field VALUES (1), (-1), (NULL) | |
| query | ||
| SELECT s.x, s.n FROM ( | ||
| SELECT IF(a > 0, named_struct('x', a, 'n', null), null) AS s FROM test_struct_null_field) | ||
|
|
||
| -- A field of a struct that native evaluates as a scalar (a cast of named_struct over a scalar | ||
| -- subquery) is one value for every row. The table is one file, so at the default batch size its | ||
| -- four rows share a batch. expect_native checks that the field access itself reaches native, so | ||
| -- that an optimizer rewrite cannot turn these into tests of something else. | ||
| statement | ||
| CREATE TABLE test_struct_scalar(a int) USING parquet | ||
|
|
||
| statement | ||
| INSERT INTO test_struct_scalar SELECT /*+ REPARTITION(1) */ * FROM VALUES (1), (-1), (NULL), (4) AS v(a) | ||
|
|
||
| query expect_native(getstructfield) | ||
| SELECT a, CAST(named_struct('x', (SELECT max(a) FROM test_struct_scalar)) AS STRUCT<x: BIGINT>).x, CAST(named_struct('x', (SELECT max(a) FROM test_struct_scalar WHERE a > 100)) AS STRUCT<x: BIGINT>).x FROM test_struct_scalar | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This only reads an integer field, so the list arms of
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Thanks. All three run natively on a scalar struct and fail on main with the same length error, so I added them to The These queries now use |
||
|
|
||
| -- The same holds for an array, a struct and a map field. | ||
| query expect_native(getstructfield) | ||
| SELECT a, CAST(named_struct('l', array((SELECT max(a) FROM test_struct_scalar))) AS STRUCT<l: ARRAY<BIGINT>>).l FROM test_struct_scalar | ||
|
|
||
| query expect_native(getstructfield) | ||
| SELECT a, CAST(named_struct('s', named_struct('y', (SELECT max(a) FROM test_struct_scalar))) AS STRUCT<s: STRUCT<y: BIGINT>>).s FROM test_struct_scalar | ||
|
|
||
| -- The map comes from map_from_arrays, which runs natively; map() over a scalar subquery would run | ||
| -- in Spark. | ||
| query expect_native(getstructfield, map_from_arrays) | ||
| SELECT a, CAST(named_struct('m', map_from_arrays(array('k'), array((SELECT max(a) FROM test_struct_scalar)))) AS STRUCT<m: MAP<STRING, BIGINT>>).m FROM test_struct_scalar | ||
|
|
||
| -- A field of a null struct that native evaluates as a scalar (an IF over a scalar subquery whose | ||
| -- other branch is NULL) is NULL for every row, next to the same IF selecting the struct. | ||
| query expect_native(getstructfield) | ||
| SELECT a, IF((SELECT max(a) FROM test_struct_scalar) > 100, named_struct('x', (SELECT max(a) FROM test_struct_scalar)), NULL).x, IF((SELECT max(a) FROM test_struct_scalar) > 0, named_struct('x', (SELECT max(a) FROM test_struct_scalar)), NULL).x FROM test_struct_scalar | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Preserve the element field for
ListViewandLargeListViewtoo. I evaluatedGetStructFieldon a one-rowLiteral(Struct)whose list-view element hadARROW:extension:namemetadata. For both offset widths,expr.data_type()retained the metadata butexpr.evaluate()returned a scalar without it. The base expression preserved it on the same input. DataFusion 55.2 rebuilds these two views intry_from_arrayjust as it does the other list types. Wrapping their one-row arrays directly and extending the existing metadata test to both widths would keep the declared and evaluated types consistent. I reproduced this through the nativePhysicalExprAPI. Spark’s standardArrayTypemaps toList, and I found no Spark SQL producer for this case.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks, I could reproduce that.
ListViewandLargeListViewnow wrap the one-row array directly, like the other list types. The metadata test now coversList,ListViewandLargeListView. With the two new arms removed, it fails onListViewbecause the evaluated type has no metadata.