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
21 changes: 20 additions & 1 deletion native/spark-expr/src/array_funcs/array_slice.rs
Original file line number Diff line number Diff line change
Expand Up @@ -188,6 +188,14 @@ mod tests {
use arrow::datatypes::{Field, Int32Type};

fn build_list(rows: Vec<Option<Vec<Option<i32>>>>) -> Arc<ListArray> {
let field = Arc::new(Field::new("item", DataType::Int32, true));
build_list_with_field(field, rows)
}

fn build_list_with_field(
field: FieldRef,
rows: Vec<Option<Vec<Option<i32>>>>,
) -> Arc<ListArray> {
let mut offsets = vec![0i32];
let mut values: Vec<Option<i32>> = Vec::new();
let mut nulls = NullBufferBuilder::new(rows.len());
Expand All @@ -202,7 +210,6 @@ mod tests {
offsets.push(values.len() as i32);
}
let values = Arc::new(Int32Array::from(values)) as ArrayRef;
let field = Arc::new(Field::new("item", DataType::Int32, true));
Arc::new(ListArray::new(
field,
OffsetBuffer::new(offsets.into()),
Expand Down Expand Up @@ -371,4 +378,16 @@ mod tests {
let length = Int64Array::from(vec![Some(-1)]);
assert!(slice_list::<i32>(list.as_ref(), &start, &length).is_err());
}

#[test]
fn preserves_non_nullable_element_field() {
let mut metadata = std::collections::HashMap::new();
metadata.insert("spark.element".to_string(), "non-null".to_string());
let field = Arc::new(Field::new("element", DataType::Int32, false).with_metadata(metadata));
let list = build_list_with_field(field, vec![Some(vec![Some(1), Some(2), Some(3)])]);
let start = Int64Array::from(vec![Some(1)]);
let length = Int64Array::from(vec![Some(2)]);
let result = slice_list::<i32>(list.as_ref(), &start, &length).unwrap();
assert_eq!(result.data_type(), list.data_type());
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -1489,6 +1489,23 @@ class CometArrayExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelp
}
}

// Spark declares split and sequence as ArrayType(..., containsNull=false). Native
// Parquet already normalizes stored children to nullable, so the non-null element
// field is produced after the scan. CometSlice must infer its return type from that
// input field; a planned nullable element disagrees with the kernel and crashes.
// https://github.com/apache/datafusion-comet/issues/5743
test("slice over expression-produced non-null element arrays (#5743)") {
val input = Seq((1, "axb", 2), (2, "", 3), (3, "cxd", 2))
withParquetDataFrame(input) { parquet =>
withParquetTable(parquet.toDF("id", "s", "n"), "t") {
checkSparkAnswerAndOperator(sql("SELECT id, slice(split(s, 'x'), 1, n) AS a FROM t"))
checkSparkAnswerAndOperator(sql("SELECT id, slice(sequence(1, n), 1, 2) AS a FROM t"))
checkSparkAnswerAndOperator(
sql("SELECT id, slice(concat(split(s, 'x'), array('z')), 1, n) AS a FROM t"))
}
}
}

// https://issues.apache.org/jira/browse/SPARK-55747
test("(ansi) GetArrayItem on null array from split()") {
withSQLConf(
Expand Down
Loading