1 // Licensed to the Apache Software Foundation (ASF) under one
2 // or more contributor license agreements. See the NOTICE file
3 // distributed with this work for additional information
4 // regarding copyright ownership. The ASF licenses this file
5 // to you under the Apache License, Version 2.0 (the
6 // "License"); you may not use this file except in compliance
7 // with the License. You may obtain a copy of the License at
8 //
9 // http://www.apache.org/licenses/LICENSE-2.0
10 //
11 // Unless required by applicable law or agreed to in writing,
12 // software distributed under the License is distributed on an
13 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14 // KIND, either express or implied. See the License for the
15 // specific language governing permissions and limitations
16 // under the License.
17
18 //! Column expression
19
20 use std::sync::Arc;
21
22 use arrow::{
23 datatypes::{DataType, Schema},
24 record_batch::RecordBatch,
25 };
26
27 use crate::error::Result;
28 use crate::physical_plan::{ColumnarValue, PhysicalExpr};
29
30 /// Represents the column at a given index in a RecordBatch
31 #[derive(Debug)]
32 pub struct Column {
33 name: String,
34 }
35
36 impl Column {
37 /// Create a new column expression
new(name: &str) -> Self38 pub fn new(name: &str) -> Self {
39 Self {
40 name: name.to_owned(),
41 }
42 }
43
44 /// Get the column name
name(&self) -> &str45 pub fn name(&self) -> &str {
46 &self.name
47 }
48 }
49
50 impl std::fmt::Display for Column {
fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result51 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
52 write!(f, "{}", self.name)
53 }
54 }
55
56 impl PhysicalExpr for Column {
57 /// Return a reference to Any that can be used for downcasting
as_any(&self) -> &dyn std::any::Any58 fn as_any(&self) -> &dyn std::any::Any {
59 self
60 }
61
62 /// Get the data type of this expression, given the schema of the input
data_type(&self, input_schema: &Schema) -> Result<DataType>63 fn data_type(&self, input_schema: &Schema) -> Result<DataType> {
64 Ok(input_schema
65 .field_with_name(&self.name)?
66 .data_type()
67 .clone())
68 }
69
70 /// Decide whehter this expression is nullable, given the schema of the input
nullable(&self, input_schema: &Schema) -> Result<bool>71 fn nullable(&self, input_schema: &Schema) -> Result<bool> {
72 Ok(input_schema.field_with_name(&self.name)?.is_nullable())
73 }
74
75 /// Evaluate the expression
evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue>76 fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
77 Ok(ColumnarValue::Array(
78 batch.column(batch.schema().index_of(&self.name)?).clone(),
79 ))
80 }
81 }
82
83 /// Create a column expression
col(name: &str) -> Arc<dyn PhysicalExpr>84 pub fn col(name: &str) -> Arc<dyn PhysicalExpr> {
85 Arc::new(Column::new(name))
86 }
87