| 218 | }; |
| 219 | |
| 220 | arrow::Result<Expression> CheckSupportedArithmeticExpression( |
| 221 | const Expression &expr, const arrow::Schema &schema) { |
| 222 | // Case 0: Literal, must be numeric type |
| 223 | if (auto literal = expr.literal()) { |
| 224 | auto type = literal->type(); |
| 225 | if (IsNumericType(type)) { |
| 226 | return expr; |
| 227 | } else { |
| 228 | return arrow::Status::Invalid("Only numeric literals are allowed, got: ", |
| 229 | literal->ToString()); |
| 230 | } |
| 231 | } |
| 232 | |
| 233 | // Case 1: Single column reference (e.g., col) |
| 234 | if (auto field_ref = expr.field_ref()) { |
| 235 | auto field = schema.GetFieldByName(*field_ref->name()); |
| 236 | if (!field) { |
| 237 | return arrow::Status::Invalid("Field not found: ", *field_ref->name()); |
| 238 | } |
| 239 | if (!IsNumericType(field->type())) { |
| 240 | return arrow::Status::Invalid( |
| 241 | "Only numeric columns are allowed, but got: ", field->ToString()); |
| 242 | } |
| 243 | return expr; // Valid, return directly |
| 244 | } |
| 245 | |
| 246 | // Step 2: Handle function calls (unary, binary, etc.) |
| 247 | if (auto call = expr.call()) { |
| 248 | const auto &func_name = call->function_name; |
| 249 | |
| 250 | // Case 2: Binary arithmetic operations (e.g., col + 1) |
| 251 | if (func_name == "add" || func_name == "subtract" || |
| 252 | func_name == "multiply" || func_name == "divide") { |
| 253 | if (call->arguments.size() != 2) { |
| 254 | return arrow::Status::Invalid("Expected two arguments for '", func_name, |
| 255 | "'"); |
| 256 | } |
| 257 | |
| 258 | const auto &left = call->arguments[0]; |
| 259 | const auto &right = call->arguments[1]; |
| 260 | |
| 261 | // One must be field_ref, the other must be literal |
| 262 | bool left_is_field = left.field_ref() != nullptr; |
| 263 | bool right_is_literal = right.literal() != nullptr; |
| 264 | |
| 265 | if (left_is_field && right_is_literal) { |
| 266 | auto field = schema.GetFieldByName(*left.field_ref()->name()); |
| 267 | if (!field) { |
| 268 | return arrow::Status::Invalid("Field not found: ", |
| 269 | *left.field_ref()->name()); |
| 270 | } |
| 271 | if (!IsNumericType(field->type())) { |
| 272 | return arrow::Status::Invalid("Column is not numeric: ", |
| 273 | field->ToString()); |
| 274 | } |
| 275 | return expr; |
| 276 | } |
| 277 | |