Skip to content

Commit 5f4d912

Browse files
tonyalaribetonyalaribe
committed
remove array_eelemnt function since its not needed
1 parent 84d2501 commit 5f4d912

File tree

2 files changed

+4
-216
lines changed

2 files changed

+4
-216
lines changed

‎src/functions.rs‎

Lines changed: 0 additions & 212 deletions
Original file line numberDiff line numberDiff line change
@@ -45,9 +45,6 @@ pub fn register_custom_functions(ctx: &mut datafusion::execution::context::Sessi
4545
// Register approx_percentile scalar function
4646
ctx.register_udf(create_approx_percentile_udf());
4747

48-
// Register array_element function
49-
ctx.register_udf(create_array_element_udf());
50-
5148
Ok(())
5249
}
5350

@@ -942,215 +939,6 @@ impl ScalarUDFImpl for ApproxPercentileUDF {
942939
}
943940
}
944941

945-
/// Create array_element UDF for PostgreSQL-compatible array access
946-
fn create_array_element_udf() -> ScalarUDF {
947-
let udf = ArrayElementUDF::new();
948-
ScalarUDF::new_from_impl(udf)
949-
}
950-
951-
/// UDF implementation for array_element
952-
#[derive(Debug)]
953-
struct ArrayElementUDF {
954-
signature: Signature,
955-
}
956-
957-
impl ArrayElementUDF {
958-
fn new() -> Self {
959-
Self {
960-
signature: Signature::any(2, Volatility::Immutable),
961-
}
962-
}
963-
}
964-
965-
impl ScalarUDFImpl for ArrayElementUDF {
966-
fn as_any(&self) -> &dyn Any {
967-
self
968-
}
969-
970-
fn name(&self) -> &str {
971-
"array_element"
972-
}
973-
974-
fn signature(&self) -> &Signature {
975-
&self.signature
976-
}
977-
978-
fn return_type(&self, arg_types: &[DataType]) -> datafusion::error::Result<DataType> {
979-
if arg_types.len() != 2 {
980-
return Err(DataFusionError::Execution(
981-
"array_element requires exactly 2 arguments".to_string(),
982-
));
983-
}
984-
985-
match &arg_types[0] {
986-
DataType::List(field) => Ok(field.data_type().clone()),
987-
_ => Err(DataFusionError::Execution(
988-
"First argument must be an array".to_string(),
989-
)),
990-
}
991-
}
992-
993-
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> datafusion::error::Result<ColumnarValue> {
994-
if args.args.len() != 2 {
995-
return Err(DataFusionError::Execution(
996-
"array_element requires exactly 2 arguments: array and index".to_string(),
997-
));
998-
}
999-
1000-
let array_arg = &args.args[0];
1001-
let index_arg = &args.args[1];
1002-
1003-
// Convert to arrays
1004-
let list_array = match array_arg {
1005-
ColumnarValue::Array(array) => array.clone(),
1006-
ColumnarValue::Scalar(scalar) => scalar.to_array_of_size(1)?,
1007-
};
1008-
1009-
let index_array = match index_arg {
1010-
ColumnarValue::Array(array) => array.clone(),
1011-
ColumnarValue::Scalar(scalar) => scalar.to_array_of_size(list_array.len())?,
1012-
};
1013-
1014-
// Get the list array
1015-
let list_array = list_array
1016-
.as_any()
1017-
.downcast_ref::<datafusion::arrow::array::ListArray>()
1018-
.ok_or_else(|| DataFusionError::Execution("First argument must be a list array".to_string()))?;
1019-
1020-
let index_array = index_array
1021-
.as_any()
1022-
.downcast_ref::<Int64Array>()
1023-
.ok_or_else(|| DataFusionError::Execution("Second argument must be an integer".to_string()))?;
1024-
1025-
// Get the data type of list elements
1026-
let element_type = match list_array.data_type() {
1027-
DataType::List(field) => field.data_type(),
1028-
_ => return Err(DataFusionError::Execution("Expected list data type".to_string())),
1029-
};
1030-
1031-
// Create a builder for the result based on element type
1032-
let result = match element_type {
1033-
DataType::Float32 => {
1034-
let mut builder = datafusion::arrow::array::Float32Array::builder(list_array.len());
1035-
for i in 0..list_array.len() {
1036-
if list_array.is_null(i) || index_array.is_null(i) {
1037-
builder.append_null();
1038-
} else {
1039-
let idx = index_array.value(i) as usize;
1040-
// PostgreSQL uses 1-based indexing
1041-
if idx == 0 || idx > list_array.value(i).len() {
1042-
builder.append_null();
1043-
} else {
1044-
let values = list_array.value(i);
1045-
let float_values = values
1046-
.as_any()
1047-
.downcast_ref::<datafusion::arrow::array::Float32Array>()
1048-
.ok_or_else(|| DataFusionError::Execution("Expected Float32 array elements".to_string()))?;
1049-
builder.append_value(float_values.value((idx - 1) as usize));
1050-
}
1051-
}
1052-
}
1053-
Arc::new(builder.finish()) as ArrayRef
1054-
}
1055-
DataType::Float64 => {
1056-
let mut builder = Float64Array::builder(list_array.len());
1057-
for i in 0..list_array.len() {
1058-
if list_array.is_null(i) || index_array.is_null(i) {
1059-
builder.append_null();
1060-
} else {
1061-
let idx = index_array.value(i) as usize;
1062-
// PostgreSQL uses 1-based indexing
1063-
if idx == 0 || idx > list_array.value(i).len() {
1064-
builder.append_null();
1065-
} else {
1066-
let values = list_array.value(i);
1067-
let float_values = values
1068-
.as_any()
1069-
.downcast_ref::<Float64Array>()
1070-
.ok_or_else(|| DataFusionError::Execution("Expected Float64 array elements".to_string()))?;
1071-
builder.append_value(float_values.value((idx - 1) as usize));
1072-
}
1073-
}
1074-
}
1075-
Arc::new(builder.finish()) as ArrayRef
1076-
}
1077-
DataType::Utf8 => {
1078-
let mut builder = StringBuilder::new();
1079-
for i in 0..list_array.len() {
1080-
if list_array.is_null(i) || index_array.is_null(i) {
1081-
builder.append_null();
1082-
} else {
1083-
let idx = index_array.value(i) as usize;
1084-
// PostgreSQL uses 1-based indexing
1085-
if idx == 0 || idx > list_array.value(i).len() {
1086-
builder.append_null();
1087-
} else {
1088-
let values = list_array.value(i);
1089-
let string_values = values
1090-
.as_any()
1091-
.downcast_ref::<StringArray>()
1092-
.ok_or_else(|| DataFusionError::Execution("Expected String array elements".to_string()))?;
1093-
builder.append_value(string_values.value((idx - 1) as usize));
1094-
}
1095-
}
1096-
}
1097-
Arc::new(builder.finish()) as ArrayRef
1098-
}
1099-
DataType::Int32 => {
1100-
let mut builder = datafusion::arrow::array::Int32Array::builder(list_array.len());
1101-
for i in 0..list_array.len() {
1102-
if list_array.is_null(i) || index_array.is_null(i) {
1103-
builder.append_null();
1104-
} else {
1105-
let idx = index_array.value(i) as usize;
1106-
// PostgreSQL uses 1-based indexing
1107-
if idx == 0 || idx > list_array.value(i).len() {
1108-
builder.append_null();
1109-
} else {
1110-
let values = list_array.value(i);
1111-
let int_values = values
1112-
.as_any()
1113-
.downcast_ref::<datafusion::arrow::array::Int32Array>()
1114-
.ok_or_else(|| DataFusionError::Execution("Expected Int32 array elements".to_string()))?;
1115-
builder.append_value(int_values.value((idx - 1) as usize));
1116-
}
1117-
}
1118-
}
1119-
Arc::new(builder.finish()) as ArrayRef
1120-
}
1121-
DataType::Int64 => {
1122-
let mut builder = Int64Array::builder(list_array.len());
1123-
for i in 0..list_array.len() {
1124-
if list_array.is_null(i) || index_array.is_null(i) {
1125-
builder.append_null();
1126-
} else {
1127-
let idx = index_array.value(i) as usize;
1128-
// PostgreSQL uses 1-based indexing
1129-
if idx == 0 || idx > list_array.value(i).len() {
1130-
builder.append_null();
1131-
} else {
1132-
let values = list_array.value(i);
1133-
let int_values = values
1134-
.as_any()
1135-
.downcast_ref::<Int64Array>()
1136-
.ok_or_else(|| DataFusionError::Execution("Expected Int64 array elements".to_string()))?;
1137-
builder.append_value(int_values.value((idx - 1) as usize));
1138-
}
1139-
}
1140-
}
1141-
Arc::new(builder.finish()) as ArrayRef
1142-
}
1143-
_ => {
1144-
return Err(DataFusionError::Execution(
1145-
format!("Unsupported array element type: {:?}", element_type),
1146-
));
1147-
}
1148-
};
1149-
1150-
Ok(ColumnarValue::Array(result))
1151-
}
1152-
}
1153-
1154942
#[cfg(test)]
1155943
mod tests {
1156944
use super::*;

‎tests/slt/percentile_functions.slt‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -211,18 +211,18 @@ LIMIT 1
211211
----
212212
true
213213

214-
# Test 4: array_element function with literal array
214+
# Test 4: array indexing with literal array (0-based indexing)
215215
query R
216-
SELECT array_element(ARRAY[10.5, 20.5, 30.5], 2) as second_element
216+
SELECT ARRAY[10.5, 20.5, 30.5][1] as second_element
217217
FROM test_spans
218218
WHERE project_id = 'test-project'
219219
LIMIT 1
220220
----
221221
20.5
222222

223-
# Test 5: array_element with string array
223+
# Test 5: array indexing with string array (0-based indexing)
224224
query T
225-
SELECT array_element(ARRAY['p50', 'p75', 'p90', 'p95'], 3) as third_quantile
225+
SELECT ARRAY['p50', 'p75', 'p90', 'p95'][2] as third_quantile
226226
FROM test_spans
227227
WHERE project_id = 'test-project'
228228
LIMIT 1

0 commit comments

Comments
 (0)