Skip to content

feat: support literal for ARRAY top level #1978

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.

Sign up for GitHub

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

Draft
wants to merge 10 commits into
base: main
Choose a base branch
from
+523 −18
Draft

feat: support literal for ARRAY top level #1978

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
b37cc38
feat: support literal for ARRAY top level
comphead Jul 2, 2025
09d43c8
feat: support literal for ARRAY top level
comphead Jul 2, 2025
3ce906d
feat: support literal for ARRAY top level
comphead Jul 2, 2025
8aa045b
fixed size list
comphead Jul 2, 2025
3308c72
feat: support literal for ARRAY top level
comphead Jul 2, 2025
6b70695
feat: support literal for ARRAY top level
comphead Jul 10, 2025
eb5f2bf
feat: support literal for ARRAY top level
comphead Jul 10, 2025
d8029cb
feat: support literal for ARRAY top level
comphead Jul 10, 2025
5c5522b
feat: support literal for ARRAY top level
comphead Jul 10, 2025
2148f5d
feat: support literal for ARRAY top level
comphead Jul 12, 2025
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
241 changes: 240 additions & 1 deletion native/core/src/execution/planner.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,12 @@ use datafusion::physical_expr::window::WindowExpr;
use datafusion::physical_expr::LexOrdering;

use crate::parquet::parquet_exec::init_datasource_exec;
use arrow::array::{
BinaryArray, BinaryBuilder, BooleanBuilder, Date32Builder, Decimal128Array, Decimal128Builder,
Float32Builder, Float64Builder, Int16Builder, Int32Builder, Int64Builder, Int8Builder,
NullArray, StringBuilder, TimestampMicrosecondBuilder,
};
use datafusion::common::utils::SingleRowListArrayBuilder;
use datafusion::physical_plan::coalesce_batches::CoalesceBatchesExec;
use datafusion::physical_plan::filter::FilterExec as DataFusionFilterExec;
use datafusion_comet_proto::spark_operator::SparkFilePartition;
Expand Down Expand Up @@ -440,6 +446,240 @@ impl PhysicalPlanner {
)))
}
}
},
Value::ListVal(values) => {
if let DataType::List(f) = data_type {
match f.data_type() {
DataType::Null => {
SingleRowListArrayBuilder::new(Arc::new(NullArray::new(values.clone().null_mask.len())))
.build_list_scalar()
}
DataType::Boolean => {
let vals = values.clone();
Copy link
Contributor Author

Choose a reason for hiding this comment

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

I'll create a macro in followup PR to remove repeated code

let len = vals.boolean_values.len();
let mut arr = BooleanBuilder::with_capacity(len);

for i in 0 .. len {
if !vals.null_mask[i] {
arr.append_value(vals.boolean_values[i]);
} else {
arr.append_null();
}
}

SingleRowListArrayBuilder::new(Arc::new(arr.finish()))
.build_list_scalar()
}
DataType::Int8 => {
let vals = values.clone();
let len = vals.byte_values.len();
let mut arr = Int8Builder::with_capacity(len);

for i in 0 .. len {
if !vals.null_mask[i] {
arr.append_value(vals.byte_values[i] as i8);
} else {
arr.append_null();
}
}

SingleRowListArrayBuilder::new(Arc::new(arr.finish()))
.build_list_scalar()
}
DataType::Int16 => {
let vals = values.clone();
let len = vals.short_values.len();
let mut arr = Int16Builder::with_capacity(len);

for i in 0 .. len {
if !vals.null_mask[i] {
arr.append_value(vals.short_values[i] as i16);
} else {
arr.append_null();
}
}

SingleRowListArrayBuilder::new(Arc::new(arr.finish()))
.build_list_scalar()
}
DataType::Int32 => {
let vals = values.clone();
let len = vals.int_values.len();
let mut arr = Int32Builder::with_capacity(len);

for i in 0 .. len {
if !vals.null_mask[i] {
arr.append_value(vals.int_values[i]);
} else {
arr.append_null();
}
}

SingleRowListArrayBuilder::new(Arc::new(arr.finish()))
.build_list_scalar()
}
DataType::Int64 => {
let vals = values.clone();
let len = vals.long_values.len();
let mut arr = Int64Builder::with_capacity(len);

for i in 0 .. len {
if !vals.null_mask[i] {
arr.append_value(vals.long_values[i]);
} else {
arr.append_null();
}
}

SingleRowListArrayBuilder::new(Arc::new(arr.finish()))
.build_list_scalar()
}
DataType::Float32 => {
let vals = values.clone();
let len = vals.float_values.len();
let mut arr = Float32Builder::with_capacity(len);

for i in 0 .. len {
if !vals.null_mask[i] {
arr.append_value(vals.float_values[i]);
} else {
arr.append_null();
}
}

SingleRowListArrayBuilder::new(Arc::new(arr.finish()))
.build_list_scalar()
}
DataType::Float64 => {
let vals = values.clone();
let len = vals.double_values.len();
let mut arr = Float64Builder::with_capacity(len);

for i in 0 .. len {
if !vals.null_mask[i] {
arr.append_value(vals.double_values[i]);
} else {
arr.append_null();
}
}

SingleRowListArrayBuilder::new(Arc::new(arr.finish()))
.build_list_scalar()
}
DataType::Timestamp(TimeUnit::Microsecond, None) => {
let vals = values.clone();
let len = vals.long_values.len();
let mut arr = TimestampMicrosecondBuilder::with_capacity(len);

for i in 0 .. len {
if !vals.null_mask[i] {
arr.append_value(vals.long_values[i]);
} else {
arr.append_null();
}
}

SingleRowListArrayBuilder::new(Arc::new(arr.finish()))
.build_list_scalar()
}
DataType::Timestamp(TimeUnit::Microsecond, Some(tz)) => {
let vals = values.clone();
let len = vals.long_values.len();
let mut arr = TimestampMicrosecondBuilder::with_capacity(len);

for i in 0 .. len {
if !vals.null_mask[i] {
arr.append_value(vals.long_values[i]);
} else {
arr.append_null();
}
}

SingleRowListArrayBuilder::new(Arc::new(arr.finish().with_timezone(Arc::clone(tz))))
.build_list_scalar()
}
DataType::Date32 => {
let vals = values.clone();
let len = vals.int_values.len();
let mut arr = Date32Builder::with_capacity(len);

for i in 0 .. len {
if !vals.null_mask[i] {
arr.append_value(vals.int_values[i]);
} else {
arr.append_null();
}
}

SingleRowListArrayBuilder::new(Arc::new(arr.finish()))
.build_list_scalar()
}
DataType::Binary => {
let vals = values.clone();
let mut arr = BinaryBuilder::new();

for (i, v) in vals.bytes_values.into_iter().enumerate() {
if !vals.null_mask[i] {
arr.append_value(v);
} else {
arr.append_null();
}
}

let binary_array: BinaryArray = arr.finish();
SingleRowListArrayBuilder::new(Arc::new(binary_array))
.build_list_scalar()
}
DataType::Utf8 => {
let vals = values.clone();
let len = vals.string_values.len();
let mut arr = StringBuilder::with_capacity(len, len);

for (i, v) in vals.string_values.into_iter().enumerate() {
if !vals.null_mask[i] {
arr.append_value(v);
} else {
arr.append_null();
}
}

SingleRowListArrayBuilder::new(Arc::new(arr.finish()))
.build_list_scalar()
}
DataType::Decimal128(p, s) => {
let vals = values.clone();
let mut arr = Decimal128Builder::new().with_precision_and_scale(*p, *s)?;

for (i, v) in vals.decimal_values.into_iter().enumerate() {
if !vals.null_mask[i] {
let big_integer = BigInt::from_signed_bytes_be(&v);
let integer = big_integer.to_i128().ok_or_else(|| {
GeneralError(format!(
"Cannot parse {big_integer:?} as i128 for Decimal literal"
))
})?;
arr.append_value(integer);
} else {
arr.append_null();
}
}

let decimal_array: Decimal128Array = arr.finish();
SingleRowListArrayBuilder::new(Arc::new(decimal_array))
.build_list_scalar()
}
dt => {
return Err(GeneralError(format!(
"DataType::List literal does not support {dt:?} type"
)))
}
}

} else {
return Err(GeneralError(format!(
"Expected DataType::List but got {data_type:?}"
)))
}
}
}
};
Expand Down Expand Up @@ -2273,7 +2513,6 @@ impl PhysicalPlanner {
other => other,
};
let func = self.session_ctx.udf(fun_name)?;

let coerced_types = func
.coerce_types(&input_expr_types)
.unwrap_or_else(|_| input_expr_types.clone());
Expand Down
1 change: 1 addition & 0 deletions native/proto/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@

// Include generated modules from .proto files.
#[allow(missing_docs)]
#[allow(clippy::large_enum_variant)]
Copy link
Contributor Author

Choose a reason for hiding this comment

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

clippy complains on proto enums

pub mod spark_expression {
include!(concat!("generated", "/spark.spark_expression.rs"));
}
Expand Down
20 changes: 11 additions & 9 deletions native/proto/src/proto/expr.proto
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,8 @@ syntax = "proto3";

package spark.spark_expression;

import "types.proto";

option java_package = "org.apache.comet.serde";

// The basic message representing a Spark expression.
Expand Down Expand Up @@ -110,13 +112,13 @@ enum StatisticsType {
}

message Count {
repeated Expr children = 1;
repeated Expr children = 1;
}

message Sum {
Expr child = 1;
DataType datatype = 2;
bool fail_on_error = 3;
Expr child = 1;
DataType datatype = 2;
bool fail_on_error = 3;
}

message Min {
Expand Down Expand Up @@ -213,10 +215,11 @@ message Literal {
string string_val = 8;
bytes bytes_val = 9;
bytes decimal_val = 10;
}
ListLiteral list_val = 11;
}

DataType datatype = 11;
bool is_null = 12;
DataType datatype = 12;
bool is_null = 13;
}

message MathExpr {
Expand Down Expand Up @@ -469,5 +472,4 @@ message DataType {
}

DataTypeInfo type_info = 2;
Copy link
Contributor Author

Choose a reason for hiding this comment

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

I'm preferring to move all DataType proto message into types.proto

}

}
41 changes: 41 additions & 0 deletions native/proto/src/proto/types.proto
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.



syntax = "proto3";

package spark.spark_expression;

option java_package = "org.apache.comet.serde";

message ListLiteral {
// Only one of these fields should be populated based on the array type
repeated bool boolean_values = 1;
repeated int32 byte_values = 2;
repeated int32 short_values = 3;
repeated int32 int_values = 4;
repeated int64 long_values = 5;
repeated float float_values = 6;
repeated double double_values = 7;
repeated string string_values = 8;
repeated bytes bytes_values = 9;
repeated bytes decimal_values = 10;
repeated ListLiteral list_values = 11;

repeated bool null_mask = 12;
}
4 changes: 4 additions & 0 deletions native/spark-expr/src/conversion_funcs/cast.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ use crate::utils::array_with_timezone;
use crate::{EvalMode, SparkError, SparkResult};
use arrow::array::builder::StringBuilder;
use arrow::array::{DictionaryArray, StringArray, StructArray};
use arrow::compute::can_cast_types;
use arrow::datatypes::{DataType, Schema};
use arrow::{
array::{
Expand Down Expand Up @@ -967,6 +968,9 @@ fn cast_array(
to_type,
cast_options,
)?),
(List(_), List(_)) if can_cast_types(from_type, to_type) => {
Ok(cast_with_options(&array, to_type, &CAST_OPTIONS)?)
}
(UInt8 | UInt16 | UInt32 | UInt64, Int8 | Int16 | Int32 | Int64)
if cast_options.allow_cast_unsigned_ints =>
{
Expand Down
5 changes: 4 additions & 1 deletion spark/src/main/scala/org/apache/comet/expressions/CometCast.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@

package org.apache.comet.expressions

import org.apache.spark.sql.types.{DataType, DataTypes, DecimalType, StructType}
import org.apache.spark.sql.types.{ArrayType, DataType, DataTypes, DecimalType, NullType, StructType}

sealed trait SupportLevel

Expand Down Expand Up @@ -62,6 +62,9 @@ object CometCast {
}

(fromType, toType) match {
case (dt: ArrayType, _: ArrayType) if dt.elementType == NullType => Compatible()
case (dt: ArrayType, dt1: ArrayType) =>
isSupported(dt.elementType, dt1.elementType, timeZoneId, evalMode)
case (dt: DataType, _) if dt.typeName == "timestamp_ntz" =>
// https://github.com/apache/datafusion-comet/issues/378
toType match {
Expand Down
Loading
Loading