grorge123 commented on code in PR #7381:
URL: https://github.com/apache/datafusion-comet/pull/7381#discussion_r4500749441
##########
spark/src/test/resources/sql-tests/expressions/struct/get_struct_field.sql:
##########
@@ -38,3 +38,15 @@ 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.
+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
+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
Review Comment:
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.
##########
native/spark-expr/src/struct_funcs/get_struct_field.rs:
##########
@@ -136,10 +158,113 @@ impl Display for GetStructField {
#[cfg(test)]
mod tests {
use super::*;
- use arrow::array::{ArrayRef, Int64Array};
- use arrow::buffer::NullBuffer;
+ use arrow::array::{Int64Array, 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 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(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)]);
+ let scalar = ScalarValue::Struct(Arc::new(StructArray::new(fields,
vec![list], None)));
+ let expr = GetStructField::new(Arc::new(Literal::new(scalar)), 0);
+ match expr.evaluate(&batch).unwrap() {
+ ColumnarValue::Scalar(value) => {
+ assert_eq!(value.data_type(),
expr.data_type(&schema).unwrap())
+ }
+ ColumnarValue::Array(array) => panic!("expected a scalar, got
array {array:?}"),
+ }
+ }
Review Comment:
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.
##########
native/spark-expr/src/struct_funcs/get_struct_field.rs:
##########
@@ -123,6 +125,26 @@ 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`.
+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()))
+ }
+ _ => ScalarValue::try_from_array(&array, 0)?,
Review Comment:
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
