Skip to content

Commit

Permalink
unify cast_to function of ScalarValue (#13122)
Browse files Browse the repository at this point in the history
  • Loading branch information
JasonLi-cn authored Oct 29, 2024
1 parent d00a089 commit 4e38abd
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 35 deletions.
33 changes: 21 additions & 12 deletions datafusion/common/src/scalar/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@ use arrow::{
use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano, ScalarBuffer};
use arrow_schema::{UnionFields, UnionMode};

use crate::format::DEFAULT_CAST_OPTIONS;
use half::f16;
pub use struct_builder::ScalarStructBuilder;

Expand Down Expand Up @@ -2809,22 +2810,30 @@ impl ScalarValue {

/// Try to parse `value` into a ScalarValue of type `target_type`
pub fn try_from_string(value: String, target_type: &DataType) -> Result<Self> {
let value = ScalarValue::from(value);
let cast_options = CastOptions {
safe: false,
format_options: Default::default(),
};
let cast_arr = cast_with_options(&value.to_array()?, target_type, &cast_options)?;
ScalarValue::try_from_array(&cast_arr, 0)
ScalarValue::from(value).cast_to(target_type)
}

/// Try to cast this value to a ScalarValue of type `data_type`
pub fn cast_to(&self, data_type: &DataType) -> Result<Self> {
let cast_options = CastOptions {
safe: false,
format_options: Default::default(),
pub fn cast_to(&self, target_type: &DataType) -> Result<Self> {
self.cast_to_with_options(target_type, &DEFAULT_CAST_OPTIONS)
}

/// Try to cast this value to a ScalarValue of type `data_type` with [`CastOptions`]
pub fn cast_to_with_options(
&self,
target_type: &DataType,
cast_options: &CastOptions<'static>,
) -> Result<Self> {
let scalar_array = match (self, target_type) {
(
ScalarValue::Float64(Some(float_ts)),
DataType::Timestamp(TimeUnit::Nanosecond, None),
) => ScalarValue::Int64(Some((float_ts * 1_000_000_000_f64).trunc() as i64))
.to_array()?,
_ => self.to_array()?,
};
let cast_arr = cast_with_options(&self.to_array()?, data_type, &cast_options)?;

let cast_arr = cast_with_options(&scalar_array, target_type, cast_options)?;
ScalarValue::try_from_array(&cast_arr, 0)
}

Expand Down
27 changes: 4 additions & 23 deletions datafusion/expr-common/src/columnar_value.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
use arrow::array::{Array, ArrayRef, NullArray};
use arrow::compute::{kernels, CastOptions};
use arrow::datatypes::{DataType, TimeUnit};
use arrow::datatypes::DataType;
use datafusion_common::format::DEFAULT_CAST_OPTIONS;
use datafusion_common::{internal_err, Result, ScalarValue};
use std::sync::Arc;
Expand Down Expand Up @@ -193,28 +193,9 @@ impl ColumnarValue {
ColumnarValue::Array(array) => Ok(ColumnarValue::Array(
kernels::cast::cast_with_options(array, cast_type, &cast_options)?,
)),
ColumnarValue::Scalar(scalar) => {
let scalar_array =
if cast_type == &DataType::Timestamp(TimeUnit::Nanosecond, None) {
if let ScalarValue::Float64(Some(float_ts)) = scalar {
ScalarValue::Int64(Some(
(float_ts * 1_000_000_000_f64).trunc() as i64,
))
.to_array()?
} else {
scalar.to_array()?
}
} else {
scalar.to_array()?
};
let cast_array = kernels::cast::cast_with_options(
&scalar_array,
cast_type,
&cast_options,
)?;
let cast_scalar = ScalarValue::try_from_array(&cast_array, 0)?;
Ok(ColumnarValue::Scalar(cast_scalar))
}
ColumnarValue::Scalar(scalar) => Ok(ColumnarValue::Scalar(
scalar.cast_to_with_options(cast_type, &cast_options)?,
)),
}
}
}
Expand Down

0 comments on commit 4e38abd

Please sign in to comment.