Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
247 changes: 241 additions & 6 deletions native/spark-expr/src/struct_funcs/get_struct_field.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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:?}"
Expand All @@ -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)?,

Copy link
Copy Markdown
Member

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 ListView and LargeListView too. I evaluated GetStructField on a one-row Literal(Struct) whose list-view element had ARROW:extension:name metadata. For both offset widths, expr.data_type() retained the metadata but expr.evaluate() returned a scalar without it. The base expression preserved it on the same input. DataFusion 55.2 rebuilds these two views in try_from_array just 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 native PhysicalExpr API. Spark’s standard ArrayType maps to List, and I found no Spark SQL producer for this case.

Copy link
Copy Markdown
Contributor Author

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. ListView and LargeListView now wrap the one-row array directly, like the other list types. The metadata test now covers List, ListView and LargeListView. With the two new arms removed, it fails on ListView because the evaluated type has no metadata.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Preserve the list element field when extracting a dictionary scalar. I used CreateNamedStruct with a literal Dictionary<Int32, List<Int32>> whose list element carries ARROW:extension:name, then evaluated GetStructField on a one-row batch. The dictionary reaches this fallback, where DataFusion recursively converts its list value and drops the element metadata. data_type() still declares that metadata, so RecordBatch::try_new rejects the resulting column. The same batch validates on base. Could the dictionary conversion extract the referenced value through the same metadata-preserving path and retain the key type? A regression test can assert that the result constructs a batch with the expression’s declared schema. I reproduced this through the native API and found no Spark SQL producer for this encoding.

Copy link
Copy Markdown
Contributor Author

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: before the change, RecordBatch::try_new rejects the column with "column types must match schema types". A dictionary now converts the value its key points at through single_row_scalar, so a list value is wrapped directly, and it keeps its key type. A null key gives a null of the value type, which DataFusion builds from the list's field, so that keeps the metadata too. The new test builds a batch with the declared schema for a present key, a null key and a null struct, and checks that the key type is still Int32 and the value is the list the key points at.

})
}

impl Display for GetStructField {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(
Expand All @@ -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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This loop only compares value.data_type() for a present list, so it would still pass if the wrapped list held the wrong row or lost the parent's nulls. Could it also assert the value (the [1, 2] list) and add a null-struct case for each list width, since that is where child_with_parent_nulls feeds the directly wrapped array?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done. The test now covers every list type that single_row_scalar wraps directly: List, LargeList, FixedSizeList, ListView and LargeListView. For each one it checks that the value is the [1, 2] list, and that a null struct whose child slot still holds the list gives a null of the declared type. It fails on the null struct when the scalar arm skips child_with_parent_nulls, and on the value when the List arm wraps a different list.

}

// 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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This only reads an integer field, so the list arms of single_row_scalar and the struct and map arms that go through try_from_array are not reached from SQL. Could the fixture also read an array field, a nested struct field and a map field of the same kind of cast named_struct over the scalar subquery, for example CAST(named_struct('l', array((SELECT max(a) FROM test_struct_scalar))) AS STRUCT<l: ARRAY<BIGINT>>).l, named_struct('s', named_struct('y', <subquery>)) and named_struct('m', map('k', <subquery>))? I expect CreateNamedStruct to return a scalar struct for these too since every field value is a scalar, but I have not run them. A NULL struct would also be worth a row if IF(<scalar subquery> > 100, named_struct(...), NULL) stays a scalar natively. Today the null-struct case is only covered by the Rust test.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The 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 get_struct_field.sql, one query each. For the map I used map_from_arrays(array('k'), array(<subquery>)): Comet runs map() through codegen dispatch, which does not take a scalar subquery, so with map() the query ran in Spark.

The IF(<subquery> > 100, named_struct(...), NULL) form did not reach native GetStructField at first. PushFoldableIntoBranches pushes .x into both branches because the NULL branch is foldable. With that rule excluded, native gets a null scalar struct and main fails the same way, so the file now excludes it and has that query, with a non-null branch next to it. The earlier queries in the file still pass, with a native GetStructField, when the rule is excluded.

These queries now use expect_native(getstructfield), so a rewrite that removes the field access fails the test instead of passing it. Without the ExcludeRules line, the null struct query fails that check.


-- 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