From 44c2aa4c23bfd0d85ad52560ac26de1399107c00 Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Tue, 9 Sep 2025 17:30:31 +0800 Subject: [PATCH] refactor: use DataFusion's return_type in our function trait directly (#6935) Signed-off-by: luofucong --- src/common/function/src/function.rs | 4 +- .../function/src/scalars/date/date_add.rs | 16 ++- .../function/src/scalars/date/date_format.rs | 19 ++- .../function/src/scalars/date/date_sub.rs | 16 ++- .../src/scalars/expression/is_null.rs | 13 +- .../function/src/scalars/geo/geohash.rs | 16 ++- src/common/function/src/scalars/geo/h3.rs | 118 ++++++++++-------- .../function/src/scalars/geo/measure.rs | 14 +-- .../function/src/scalars/geo/relation.rs | 14 +-- src/common/function/src/scalars/geo/s2.rs | 17 ++- src/common/function/src/scalars/geo/wkt.rs | 5 +- src/common/function/src/scalars/hll_count.rs | 11 +- src/common/function/src/scalars/ip/cidr.rs | 10 +- src/common/function/src/scalars/ip/ipv4.rs | 9 +- src/common/function/src/scalars/ip/ipv6.rs | 11 +- src/common/function/src/scalars/ip/range.rs | 11 +- .../function/src/scalars/json/json_get.rs | 36 ++---- .../function/src/scalars/json/json_is.rs | 9 +- .../src/scalars/json/json_path_exists.rs | 10 +- .../src/scalars/json/json_path_match.rs | 10 +- .../src/scalars/json/json_to_string.rs | 10 +- .../function/src/scalars/json/parse_json.rs | 11 +- src/common/function/src/scalars/matches.rs | 4 +- .../function/src/scalars/matches_term.rs | 6 +- src/common/function/src/scalars/math.rs | 4 +- src/common/function/src/scalars/math/clamp.rs | 9 +- .../function/src/scalars/math/modulo.rs | 49 +++----- src/common/function/src/scalars/math/rate.rs | 10 +- src/common/function/src/scalars/test.rs | 5 +- .../src/scalars/timestamp/to_unixtime.rs | 10 +- .../function/src/scalars/uddsketch_calc.rs | 11 +- src/common/function/src/scalars/udf.rs | 18 +-- .../scalars/vector/convert/parse_vector.rs | 6 +- .../vector/convert/vector_to_string.rs | 6 +- .../function/src/scalars/vector/distance.rs | 5 +- .../src/scalars/vector/elem_product.rs | 11 +- .../function/src/scalars/vector/elem_sum.rs | 11 +- .../function/src/scalars/vector/scalar_add.rs | 5 +- .../function/src/scalars/vector/scalar_mul.rs | 5 +- .../function/src/scalars/vector/vector_add.rs | 10 +- .../function/src/scalars/vector/vector_dim.rs | 11 +- .../function/src/scalars/vector/vector_div.rs | 5 +- .../src/scalars/vector/vector_kth_elem.rs | 8 +- .../function/src/scalars/vector/vector_mul.rs | 5 +- .../src/scalars/vector/vector_norm.rs | 6 +- .../function/src/scalars/vector/vector_sub.rs | 10 +- .../src/scalars/vector/vector_subvector.rs | 5 +- src/common/function/src/system/build.rs | 11 +- src/common/function/src/system/database.rs | 32 +++-- .../src/system/pg_catalog/pg_get_userbyid.rs | 7 +- .../src/system/pg_catalog/table_is_visible.rs | 7 +- .../function/src/system/pg_catalog/version.rs | 6 +- src/common/function/src/system/timezone.rs | 12 +- src/common/function/src/system/version.rs | 6 +- src/flow/src/transform.rs | 7 +- src/flow/src/transform/aggr.rs | 2 +- src/flow/src/transform/plan.rs | 2 +- src/servers/src/mysql/writer.rs | 16 ++- src/servers/tests/mysql/mysql_writer_test.rs | 11 +- 59 files changed, 326 insertions(+), 418 deletions(-) diff --git a/src/common/function/src/function.rs b/src/common/function/src/function.rs index a15b4c0d9b..538d3fd227 100644 --- a/src/common/function/src/function.rs +++ b/src/common/function/src/function.rs @@ -16,8 +16,8 @@ use std::fmt; use std::sync::Arc; use common_query::error::Result; +use datafusion::arrow::datatypes::DataType; use datafusion_expr::Signature; -use datatypes::data_type::ConcreteDataType; use datatypes::vectors::VectorRef; use session::context::{QueryContextBuilder, QueryContextRef}; @@ -63,7 +63,7 @@ pub trait Function: fmt::Display + Sync + Send { fn name(&self) -> &str; /// The returned data type of function execution. - fn return_type(&self, input_types: &[ConcreteDataType]) -> Result; + fn return_type(&self, input_types: &[DataType]) -> Result; /// The signature of function. fn signature(&self) -> Signature; diff --git a/src/common/function/src/scalars/date/date_add.rs b/src/common/function/src/scalars/date/date_add.rs index 682bbab7ce..9785b92a91 100644 --- a/src/common/function/src/scalars/date/date_add.rs +++ b/src/common/function/src/scalars/date/date_add.rs @@ -18,7 +18,6 @@ use common_query::error::{ArrowComputeSnafu, IntoVectorSnafu, InvalidFuncArgsSna use datafusion_expr::Signature; use datatypes::arrow::compute::kernels::numeric; use datatypes::arrow::datatypes::{DataType, IntervalUnit, TimeUnit}; -use datatypes::prelude::ConcreteDataType; use datatypes::vectors::{Helper, VectorRef}; use snafu::{ResultExt, ensure}; @@ -38,7 +37,7 @@ impl Function for DateAddFunction { NAME } - fn return_type(&self, input_types: &[ConcreteDataType]) -> Result { + fn return_type(&self, input_types: &[DataType]) -> Result { Ok(input_types[0].clone()) } @@ -93,7 +92,6 @@ mod tests { use datafusion_expr::{TypeSignature, Volatility}; use datatypes::arrow::datatypes::IntervalDayTime; - use datatypes::prelude::ConcreteDataType; use datatypes::value::Value; use datatypes::vectors::{ DateVector, IntervalDayTimeVector, IntervalYearMonthVector, TimestampSecondVector, @@ -106,18 +104,18 @@ mod tests { let f = DateAddFunction; assert_eq!("date_add", f.name()); assert_eq!( - ConcreteDataType::timestamp_microsecond_datatype(), - f.return_type(&[ConcreteDataType::timestamp_microsecond_datatype()]) + DataType::Timestamp(TimeUnit::Microsecond, None), + f.return_type(&[DataType::Timestamp(TimeUnit::Microsecond, None)]) .unwrap() ); assert_eq!( - ConcreteDataType::timestamp_second_datatype(), - f.return_type(&[ConcreteDataType::timestamp_second_datatype()]) + DataType::Timestamp(TimeUnit::Second, None), + f.return_type(&[DataType::Timestamp(TimeUnit::Second, None)]) .unwrap() ); assert_eq!( - ConcreteDataType::date_datatype(), - f.return_type(&[ConcreteDataType::date_datatype()]).unwrap() + DataType::Date32, + f.return_type(&[DataType::Date32]).unwrap() ); assert!( matches!(f.signature(), diff --git a/src/common/function/src/scalars/date/date_format.rs b/src/common/function/src/scalars/date/date_format.rs index 95104c473d..278f8ccc0b 100644 --- a/src/common/function/src/scalars/date/date_format.rs +++ b/src/common/function/src/scalars/date/date_format.rs @@ -36,8 +36,8 @@ impl Function for DateFormatFunction { NAME } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::string_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Utf8) } fn signature(&self) -> Signature { @@ -129,7 +129,7 @@ mod tests { use std::sync::Arc; use datafusion_expr::{TypeSignature, Volatility}; - use datatypes::prelude::{ConcreteDataType, ScalarVector}; + use datatypes::prelude::ScalarVector; use datatypes::value::Value; use datatypes::vectors::{DateVector, StringVector, TimestampSecondVector}; @@ -140,19 +140,16 @@ mod tests { let f = DateFormatFunction; assert_eq!("date_format", f.name()); assert_eq!( - ConcreteDataType::string_datatype(), - f.return_type(&[ConcreteDataType::timestamp_microsecond_datatype()]) + DataType::Utf8, + f.return_type(&[DataType::Timestamp(TimeUnit::Microsecond, None)]) .unwrap() ); assert_eq!( - ConcreteDataType::string_datatype(), - f.return_type(&[ConcreteDataType::timestamp_second_datatype()]) + DataType::Utf8, + f.return_type(&[DataType::Timestamp(TimeUnit::Second, None)]) .unwrap() ); - assert_eq!( - ConcreteDataType::string_datatype(), - f.return_type(&[ConcreteDataType::date_datatype()]).unwrap() - ); + assert_eq!(DataType::Utf8, f.return_type(&[DataType::Date32]).unwrap()); assert!(matches!(f.signature(), Signature { type_signature: TypeSignature::OneOf(sigs), diff --git a/src/common/function/src/scalars/date/date_sub.rs b/src/common/function/src/scalars/date/date_sub.rs index 939ce93892..e451ee3c6e 100644 --- a/src/common/function/src/scalars/date/date_sub.rs +++ b/src/common/function/src/scalars/date/date_sub.rs @@ -18,7 +18,6 @@ use common_query::error::{ArrowComputeSnafu, IntoVectorSnafu, InvalidFuncArgsSna use datafusion_expr::Signature; use datatypes::arrow::compute::kernels::numeric; use datatypes::arrow::datatypes::{DataType, IntervalUnit, TimeUnit}; -use datatypes::prelude::ConcreteDataType; use datatypes::vectors::{Helper, VectorRef}; use snafu::{ResultExt, ensure}; @@ -38,7 +37,7 @@ impl Function for DateSubFunction { NAME } - fn return_type(&self, input_types: &[ConcreteDataType]) -> Result { + fn return_type(&self, input_types: &[DataType]) -> Result { Ok(input_types[0].clone()) } @@ -93,7 +92,6 @@ mod tests { use datafusion_expr::{TypeSignature, Volatility}; use datatypes::arrow::datatypes::IntervalDayTime; - use datatypes::prelude::ConcreteDataType; use datatypes::value::Value; use datatypes::vectors::{ DateVector, IntervalDayTimeVector, IntervalYearMonthVector, TimestampSecondVector, @@ -106,18 +104,18 @@ mod tests { let f = DateSubFunction; assert_eq!("date_sub", f.name()); assert_eq!( - ConcreteDataType::timestamp_microsecond_datatype(), - f.return_type(&[ConcreteDataType::timestamp_microsecond_datatype()]) + DataType::Timestamp(TimeUnit::Microsecond, None), + f.return_type(&[DataType::Timestamp(TimeUnit::Microsecond, None)]) .unwrap() ); assert_eq!( - ConcreteDataType::timestamp_second_datatype(), - f.return_type(&[ConcreteDataType::timestamp_second_datatype()]) + DataType::Timestamp(TimeUnit::Second, None), + f.return_type(&[DataType::Timestamp(TimeUnit::Second, None)]) .unwrap() ); assert_eq!( - ConcreteDataType::date_datatype(), - f.return_type(&[ConcreteDataType::date_datatype()]).unwrap() + DataType::Date32, + f.return_type(&[DataType::Date32]).unwrap() ); assert!( matches!(f.signature(), diff --git a/src/common/function/src/scalars/expression/is_null.rs b/src/common/function/src/scalars/expression/is_null.rs index bf8abe3c5a..fd28f8682a 100644 --- a/src/common/function/src/scalars/expression/is_null.rs +++ b/src/common/function/src/scalars/expression/is_null.rs @@ -17,11 +17,11 @@ use std::fmt::Display; use std::sync::Arc; use common_query::error; -use common_query::error::{ArrowComputeSnafu, InvalidFuncArgsSnafu}; +use common_query::error::{ArrowComputeSnafu, InvalidFuncArgsSnafu, Result}; use datafusion::arrow::array::ArrayRef; use datafusion::arrow::compute::is_null; +use datafusion::arrow::datatypes::DataType; use datafusion_expr::{Signature, Volatility}; -use datatypes::data_type::ConcreteDataType; use datatypes::prelude::VectorRef; use datatypes::vectors::Helper; use snafu::{ResultExt, ensure}; @@ -45,8 +45,8 @@ impl Function for IsNullFunction { NAME } - fn return_type(&self, _: &[ConcreteDataType]) -> common_query::error::Result { - Ok(ConcreteDataType::boolean_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Boolean) } fn signature(&self) -> Signature { @@ -88,10 +88,7 @@ mod tests { fn test_is_null_function() { let is_null = IsNullFunction; assert_eq!("isnull", is_null.name()); - assert_eq!( - ConcreteDataType::boolean_datatype(), - is_null.return_type(&[]).unwrap() - ); + assert_eq!(DataType::Boolean, is_null.return_type(&[]).unwrap()); assert_eq!( is_null.signature(), Signature { diff --git a/src/common/function/src/scalars/geo/geohash.rs b/src/common/function/src/scalars/geo/geohash.rs index 8b8b00957d..59e3444ffb 100644 --- a/src/common/function/src/scalars/geo/geohash.rs +++ b/src/common/function/src/scalars/geo/geohash.rs @@ -13,10 +13,12 @@ // limitations under the License. use std::fmt; +use std::sync::Arc; use common_error::ext::{BoxedError, PlainError}; use common_error::status_code::StatusCode; use common_query::error::{self, InvalidFuncArgsSnafu, Result}; +use datafusion::arrow::datatypes::Field; use datafusion_expr::type_coercion::aggregates::INTEGERS; use datafusion_expr::{Signature, TypeSignature, Volatility}; use datatypes::arrow::datatypes::DataType; @@ -86,8 +88,8 @@ impl Function for GeohashFunction { Self::NAME } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::string_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Utf8) } fn signature(&self) -> Signature { @@ -172,10 +174,12 @@ impl Function for GeohashNeighboursFunction { GeohashNeighboursFunction::NAME } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::list_datatype( - ConcreteDataType::string_datatype(), - )) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::List(Arc::new(Field::new( + "x", + DataType::Utf8, + false, + )))) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/scalars/geo/h3.rs b/src/common/function/src/scalars/geo/h3.rs index 43cb0058cb..35ee31c68b 100644 --- a/src/common/function/src/scalars/geo/h3.rs +++ b/src/common/function/src/scalars/geo/h3.rs @@ -13,14 +13,14 @@ // limitations under the License. use std::str::FromStr; -use std::sync::LazyLock; +use std::sync::{Arc, LazyLock}; use common_error::ext::{BoxedError, PlainError}; use common_error::status_code::StatusCode; use common_query::error::{self, Result}; use datafusion_expr::type_coercion::aggregates::INTEGERS; use datafusion_expr::{Signature, TypeSignature, Volatility}; -use datatypes::arrow::datatypes::DataType; +use datatypes::arrow::datatypes::{DataType, Field}; use datatypes::prelude::ConcreteDataType; use datatypes::scalars::{Scalar, ScalarVectorBuilder}; use datatypes::value::{ListValue, Value}; @@ -59,8 +59,8 @@ impl Function for H3LatLngToCell { "h3_latlng_to_cell" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::uint64_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::UInt64) } fn signature(&self) -> Signature { @@ -131,8 +131,8 @@ impl Function for H3LatLngToCellString { "h3_latlng_to_cell_string" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::string_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Utf8) } fn signature(&self) -> Signature { @@ -200,8 +200,8 @@ impl Function for H3CellToString { "h3_cell_to_string" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::string_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Utf8) } fn signature(&self) -> Signature { @@ -235,8 +235,8 @@ impl Function for H3StringToCell { "h3_string_to_cell" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::uint64_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::UInt64) } fn signature(&self) -> Signature { @@ -285,10 +285,12 @@ impl Function for H3CellCenterLatLng { "h3_cell_center_latlng" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::list_datatype( - ConcreteDataType::float64_datatype(), - )) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::List(Arc::new(Field::new( + "x", + DataType::Float64, + false, + )))) } fn signature(&self) -> Signature { @@ -332,8 +334,8 @@ impl Function for H3CellResolution { "h3_cell_resolution" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::uint8_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::UInt8) } fn signature(&self) -> Signature { @@ -368,8 +370,8 @@ impl Function for H3CellBase { "h3_cell_base" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::uint8_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::UInt8) } fn signature(&self) -> Signature { @@ -404,8 +406,8 @@ impl Function for H3CellIsPentagon { "h3_cell_is_pentagon" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::boolean_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Boolean) } fn signature(&self) -> Signature { @@ -440,8 +442,8 @@ impl Function for H3CellCenterChild { "h3_cell_center_child" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::uint64_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::UInt64) } fn signature(&self) -> Signature { @@ -480,8 +482,8 @@ impl Function for H3CellParent { "h3_cell_parent" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::uint64_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::UInt64) } fn signature(&self) -> Signature { @@ -518,10 +520,12 @@ impl Function for H3CellToChildren { "h3_cell_to_children" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::list_datatype( - ConcreteDataType::uint64_datatype(), - )) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::List(Arc::new(Field::new( + "item", + DataType::UInt64, + true, + )))) } fn signature(&self) -> Signature { @@ -569,8 +573,8 @@ impl Function for H3CellToChildrenSize { "h3_cell_to_children_size" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::uint64_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::UInt64) } fn signature(&self) -> Signature { @@ -606,8 +610,8 @@ impl Function for H3CellToChildPos { "h3_cell_to_child_pos" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::uint64_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::UInt64) } fn signature(&self) -> Signature { @@ -643,8 +647,8 @@ impl Function for H3ChildPosToCell { "h3_child_pos_to_cell" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::uint64_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::UInt64) } fn signature(&self) -> Signature { @@ -695,10 +699,12 @@ impl Function for H3GridDisk { "h3_grid_disk" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::list_datatype( - ConcreteDataType::uint64_datatype(), - )) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::List(Arc::new(Field::new( + "item", + DataType::UInt64, + true, + )))) } fn signature(&self) -> Signature { @@ -748,10 +754,12 @@ impl Function for H3GridDiskDistances { "h3_grid_disk_distances" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::list_datatype( - ConcreteDataType::uint64_datatype(), - )) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::List(Arc::new(Field::new( + "item", + DataType::UInt64, + true, + )))) } fn signature(&self) -> Signature { @@ -800,8 +808,8 @@ impl Function for H3GridDistance { fn name(&self) -> &str { "h3_grid_distance" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::int32_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Int32) } fn signature(&self) -> Signature { @@ -854,10 +862,12 @@ impl Function for H3GridPathCells { "h3_grid_path_cells" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::list_datatype( - ConcreteDataType::uint64_datatype(), - )) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::List(Arc::new(Field::new( + "item", + DataType::UInt64, + true, + )))) } fn signature(&self) -> Signature { @@ -921,8 +931,8 @@ impl Function for H3CellContains { "h3_cells_contains" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::boolean_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Boolean) } fn signature(&self) -> Signature { @@ -992,8 +1002,8 @@ impl Function for H3CellDistanceSphereKm { fn name(&self) -> &str { "h3_distance_sphere_km" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::float64_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Float64) } fn signature(&self) -> Signature { @@ -1047,8 +1057,8 @@ impl Function for H3CellDistanceEuclideanDegree { fn name(&self) -> &str { "h3_distance_degree" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::float64_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Float64) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/scalars/geo/measure.rs b/src/common/function/src/scalars/geo/measure.rs index 37a5efe7ee..cab77fae6e 100644 --- a/src/common/function/src/scalars/geo/measure.rs +++ b/src/common/function/src/scalars/geo/measure.rs @@ -15,8 +15,8 @@ use common_error::ext::{BoxedError, PlainError}; use common_error::status_code::StatusCode; use common_query::error::{self, Result}; +use datafusion::arrow::datatypes::DataType; use datafusion_expr::{Signature, Volatility}; -use datatypes::prelude::ConcreteDataType; use datatypes::scalars::ScalarVectorBuilder; use datatypes::vectors::{Float64VectorBuilder, MutableVector, VectorRef}; use derive_more::Display; @@ -39,8 +39,8 @@ impl Function for STDistance { "st_distance" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::float64_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Float64) } fn signature(&self) -> Signature { @@ -87,8 +87,8 @@ impl Function for STDistanceSphere { "st_distance_sphere_m" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::float64_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Float64) } fn signature(&self) -> Signature { @@ -145,8 +145,8 @@ impl Function for STArea { "st_area" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::float64_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Float64) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/scalars/geo/relation.rs b/src/common/function/src/scalars/geo/relation.rs index f195655796..6dc116c4b3 100644 --- a/src/common/function/src/scalars/geo/relation.rs +++ b/src/common/function/src/scalars/geo/relation.rs @@ -13,8 +13,8 @@ // limitations under the License. use common_query::error::Result; +use datafusion::arrow::datatypes::DataType; use datafusion_expr::{Signature, Volatility}; -use datatypes::prelude::ConcreteDataType; use datatypes::scalars::ScalarVectorBuilder; use datatypes::vectors::{BooleanVectorBuilder, MutableVector, VectorRef}; use derive_more::Display; @@ -36,8 +36,8 @@ impl Function for STContains { "st_contains" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::boolean_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Boolean) } fn signature(&self) -> Signature { @@ -84,8 +84,8 @@ impl Function for STWithin { "st_within" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::boolean_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Boolean) } fn signature(&self) -> Signature { @@ -132,8 +132,8 @@ impl Function for STIntersects { "st_intersects" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::boolean_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Boolean) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/scalars/geo/s2.rs b/src/common/function/src/scalars/geo/s2.rs index 0ef5e94e4b..4e356520c2 100644 --- a/src/common/function/src/scalars/geo/s2.rs +++ b/src/common/function/src/scalars/geo/s2.rs @@ -17,7 +17,6 @@ use std::sync::LazyLock; use common_query::error::{InvalidFuncArgsSnafu, Result}; use datafusion_expr::{Signature, TypeSignature, Volatility}; use datatypes::arrow::datatypes::DataType; -use datatypes::prelude::ConcreteDataType; use datatypes::scalars::ScalarVectorBuilder; use datatypes::value::Value; use datatypes::vectors::{MutableVector, StringVectorBuilder, UInt64VectorBuilder, VectorRef}; @@ -49,8 +48,8 @@ impl Function for S2LatLngToCell { "s2_latlng_to_cell" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::uint64_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::UInt64) } fn signature(&self) -> Signature { @@ -112,8 +111,8 @@ impl Function for S2CellLevel { "s2_cell_level" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::uint64_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::UInt64) } fn signature(&self) -> Signature { @@ -148,8 +147,8 @@ impl Function for S2CellToToken { "s2_cell_to_token" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::string_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Utf8) } fn signature(&self) -> Signature { @@ -184,8 +183,8 @@ impl Function for S2CellParent { "s2_cell_parent" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::uint64_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::UInt64) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/scalars/geo/wkt.rs b/src/common/function/src/scalars/geo/wkt.rs index 3a98cacbeb..b6938bf7ee 100644 --- a/src/common/function/src/scalars/geo/wkt.rs +++ b/src/common/function/src/scalars/geo/wkt.rs @@ -19,7 +19,6 @@ use common_error::status_code::StatusCode; use common_query::error::{self, Result}; use datafusion_expr::{Signature, TypeSignature, Volatility}; use datatypes::arrow::datatypes::DataType; -use datatypes::prelude::ConcreteDataType; use datatypes::scalars::ScalarVectorBuilder; use datatypes::vectors::{MutableVector, StringVectorBuilder, VectorRef}; use derive_more::Display; @@ -43,8 +42,8 @@ impl Function for LatLngToPointWkt { "wkt_point_from_latlng" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::string_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Utf8) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/scalars/hll_count.rs b/src/common/function/src/scalars/hll_count.rs index 0ec0d8eec3..fa6e787db0 100644 --- a/src/common/function/src/scalars/hll_count.rs +++ b/src/common/function/src/scalars/hll_count.rs @@ -20,7 +20,6 @@ use std::fmt::Display; use common_query::error::{DowncastVectorSnafu, InvalidFuncArgsSnafu, Result}; use datafusion_expr::{Signature, Volatility}; use datatypes::arrow::datatypes::DataType; -use datatypes::data_type::ConcreteDataType; use datatypes::prelude::Vector; use datatypes::scalars::{ScalarVector, ScalarVectorBuilder}; use datatypes::vectors::{BinaryVector, MutableVector, UInt64VectorBuilder, VectorRef}; @@ -59,8 +58,8 @@ impl Function for HllCalcFunction { NAME } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::uint64_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::UInt64) } fn signature(&self) -> Signature { @@ -126,10 +125,8 @@ mod tests { let function = HllCalcFunction; assert_eq!("hll_count", function.name()); assert_eq!( - ConcreteDataType::uint64_datatype(), - function - .return_type(&[ConcreteDataType::uint64_datatype()]) - .unwrap() + DataType::UInt64, + function.return_type(&[DataType::UInt64]).unwrap() ); // Create a test HLL diff --git a/src/common/function/src/scalars/ip/cidr.rs b/src/common/function/src/scalars/ip/cidr.rs index bd74dd24b1..cbd5fbd922 100644 --- a/src/common/function/src/scalars/ip/cidr.rs +++ b/src/common/function/src/scalars/ip/cidr.rs @@ -19,7 +19,7 @@ use common_query::error::{InvalidFuncArgsSnafu, Result}; use datafusion_common::types; use datafusion_expr::{Coercion, Signature, TypeSignature, TypeSignatureClass, Volatility}; use datatypes::arrow::datatypes::DataType; -use datatypes::prelude::{ConcreteDataType, Value}; +use datatypes::prelude::Value; use datatypes::scalars::ScalarVectorBuilder; use datatypes::vectors::{MutableVector, StringVectorBuilder, VectorRef}; use derive_more::Display; @@ -45,8 +45,8 @@ impl Function for Ipv4ToCidr { "ipv4_to_cidr" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::string_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Utf8) } fn signature(&self) -> Signature { @@ -174,8 +174,8 @@ impl Function for Ipv6ToCidr { "ipv6_to_cidr" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::string_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Utf8) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/scalars/ip/ipv4.rs b/src/common/function/src/scalars/ip/ipv4.rs index d3adb44c54..56a1409456 100644 --- a/src/common/function/src/scalars/ip/ipv4.rs +++ b/src/common/function/src/scalars/ip/ipv4.rs @@ -18,7 +18,6 @@ use std::str::FromStr; use common_query::error::{InvalidFuncArgsSnafu, Result}; use datafusion_expr::{Signature, TypeSignature, Volatility}; use datatypes::arrow::datatypes::DataType; -use datatypes::prelude::ConcreteDataType; use datatypes::scalars::ScalarVectorBuilder; use datatypes::vectors::{MutableVector, StringVectorBuilder, UInt32VectorBuilder, VectorRef}; use derive_more::Display; @@ -53,8 +52,8 @@ impl Function for Ipv4NumToString { "ipv4_num_to_string" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::string_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Utf8) } fn signature(&self) -> Signature { @@ -116,8 +115,8 @@ impl Function for Ipv4StringToNum { "ipv4_string_to_num" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::uint32_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::UInt32) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/scalars/ip/ipv6.rs b/src/common/function/src/scalars/ip/ipv6.rs index d625884808..57bf2c2082 100644 --- a/src/common/function/src/scalars/ip/ipv6.rs +++ b/src/common/function/src/scalars/ip/ipv6.rs @@ -16,8 +16,9 @@ use std::net::{Ipv4Addr, Ipv6Addr}; use std::str::FromStr; use common_query::error::{InvalidFuncArgsSnafu, Result}; +use datafusion::arrow::datatypes::DataType; use datafusion_expr::{Signature, Volatility}; -use datatypes::prelude::{ConcreteDataType, Value}; +use datatypes::prelude::Value; use datatypes::scalars::ScalarVectorBuilder; use datatypes::vectors::{BinaryVectorBuilder, MutableVector, StringVectorBuilder, VectorRef}; use derive_more::Display; @@ -39,8 +40,8 @@ impl Function for Ipv6NumToString { "ipv6_num_to_string" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::string_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Utf8) } fn signature(&self) -> Signature { @@ -128,8 +129,8 @@ impl Function for Ipv6StringToNum { "ipv6_string_to_num" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::binary_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Binary) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/scalars/ip/range.rs b/src/common/function/src/scalars/ip/range.rs index c3e28d0d8b..ebda747e9b 100644 --- a/src/common/function/src/scalars/ip/range.rs +++ b/src/common/function/src/scalars/ip/range.rs @@ -16,8 +16,9 @@ use std::net::{Ipv4Addr, Ipv6Addr}; use std::str::FromStr; use common_query::error::{InvalidFuncArgsSnafu, Result}; +use datafusion::arrow::datatypes::DataType; use datafusion_expr::{Signature, Volatility}; -use datatypes::prelude::{ConcreteDataType, Value}; +use datatypes::prelude::Value; use datatypes::scalars::ScalarVectorBuilder; use datatypes::vectors::{BooleanVectorBuilder, MutableVector, VectorRef}; use derive_more::Display; @@ -43,8 +44,8 @@ impl Function for Ipv4InRange { "ipv4_in_range" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::boolean_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Boolean) } fn signature(&self) -> Signature { @@ -132,8 +133,8 @@ impl Function for Ipv6InRange { "ipv6_in_range" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::boolean_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Boolean) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/scalars/json/json_get.rs b/src/common/function/src/scalars/json/json_get.rs index 97c6061490..66858d5470 100644 --- a/src/common/function/src/scalars/json/json_get.rs +++ b/src/common/function/src/scalars/json/json_get.rs @@ -58,8 +58,8 @@ macro_rules! json_get { stringify!([<$name:snake>]) } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::[<$type:snake _datatype>]()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::[<$type>]) } fn signature(&self) -> Signature { @@ -159,8 +159,8 @@ impl Function for JsonGetString { "json_get_string" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::string_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Utf8) } fn signature(&self) -> Signature { @@ -242,12 +242,9 @@ mod tests { assert_eq!("json_get_int", json_get_int.name()); assert_eq!( - ConcreteDataType::int64_datatype(), + DataType::Int64, json_get_int - .return_type(&[ - ConcreteDataType::json_datatype(), - ConcreteDataType::string_datatype() - ]) + .return_type(&[DataType::Binary, DataType::Utf8]) .unwrap() ); @@ -295,12 +292,9 @@ mod tests { assert_eq!("json_get_float", json_get_float.name()); assert_eq!( - ConcreteDataType::float64_datatype(), + DataType::Float64, json_get_float - .return_type(&[ - ConcreteDataType::json_datatype(), - ConcreteDataType::string_datatype() - ]) + .return_type(&[DataType::Binary, DataType::Utf8]) .unwrap() ); @@ -348,12 +342,9 @@ mod tests { assert_eq!("json_get_bool", json_get_bool.name()); assert_eq!( - ConcreteDataType::boolean_datatype(), + DataType::Boolean, json_get_bool - .return_type(&[ - ConcreteDataType::json_datatype(), - ConcreteDataType::string_datatype() - ]) + .return_type(&[DataType::Binary, DataType::Utf8]) .unwrap() ); @@ -401,12 +392,9 @@ mod tests { assert_eq!("json_get_string", json_get_string.name()); assert_eq!( - ConcreteDataType::string_datatype(), + DataType::Utf8, json_get_string - .return_type(&[ - ConcreteDataType::json_datatype(), - ConcreteDataType::string_datatype() - ]) + .return_type(&[DataType::Binary, DataType::Utf8]) .unwrap() ); diff --git a/src/common/function/src/scalars/json/json_is.rs b/src/common/function/src/scalars/json/json_is.rs index d9ab30b58f..de2f55b53b 100644 --- a/src/common/function/src/scalars/json/json_is.rs +++ b/src/common/function/src/scalars/json/json_is.rs @@ -37,8 +37,8 @@ macro_rules! json_is { stringify!([<$name:snake>]) } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::boolean_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Boolean) } fn signature(&self) -> Signature { @@ -161,9 +161,8 @@ mod tests { for (func, expected_name) in json_is_functions.iter().zip(expected_names.iter()) { assert_eq!(func.name(), *expected_name); assert_eq!( - func.return_type(&[ConcreteDataType::json_datatype()]) - .unwrap(), - ConcreteDataType::boolean_datatype() + func.return_type(&[DataType::Binary]).unwrap(), + DataType::Boolean ); assert_eq!( func.signature(), diff --git a/src/common/function/src/scalars/json/json_path_exists.rs b/src/common/function/src/scalars/json/json_path_exists.rs index 0a11a88367..a5446ec614 100644 --- a/src/common/function/src/scalars/json/json_path_exists.rs +++ b/src/common/function/src/scalars/json/json_path_exists.rs @@ -36,8 +36,8 @@ impl Function for JsonPathExistsFunction { NAME } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::boolean_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Boolean) } fn signature(&self) -> Signature { @@ -134,10 +134,8 @@ mod tests { assert_eq!("json_path_exists", json_path_exists.name()); assert_eq!( - ConcreteDataType::boolean_datatype(), - json_path_exists - .return_type(&[ConcreteDataType::json_datatype()]) - .unwrap() + DataType::Boolean, + json_path_exists.return_type(&[DataType::Binary]).unwrap() ); assert!(matches!(json_path_exists.signature(), diff --git a/src/common/function/src/scalars/json/json_path_match.rs b/src/common/function/src/scalars/json/json_path_match.rs index e7b60b8ef7..3adbe7eab9 100644 --- a/src/common/function/src/scalars/json/json_path_match.rs +++ b/src/common/function/src/scalars/json/json_path_match.rs @@ -36,8 +36,8 @@ impl Function for JsonPathMatchFunction { NAME } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::boolean_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Boolean) } fn signature(&self) -> Signature { @@ -126,10 +126,8 @@ mod tests { assert_eq!("json_path_match", json_path_match.name()); assert_eq!( - ConcreteDataType::boolean_datatype(), - json_path_match - .return_type(&[ConcreteDataType::json_datatype()]) - .unwrap() + DataType::Boolean, + json_path_match.return_type(&[DataType::Binary]).unwrap() ); assert!(matches!(json_path_match.signature(), diff --git a/src/common/function/src/scalars/json/json_to_string.rs b/src/common/function/src/scalars/json/json_to_string.rs index 4160f73b12..a5197a249c 100644 --- a/src/common/function/src/scalars/json/json_to_string.rs +++ b/src/common/function/src/scalars/json/json_to_string.rs @@ -36,8 +36,8 @@ impl Function for JsonToStringFunction { NAME } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::string_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Utf8) } fn signature(&self) -> Signature { @@ -122,10 +122,8 @@ mod tests { assert_eq!("json_to_string", json_to_string.name()); assert_eq!( - ConcreteDataType::string_datatype(), - json_to_string - .return_type(&[ConcreteDataType::json_datatype()]) - .unwrap() + DataType::Utf8, + json_to_string.return_type(&[DataType::Binary]).unwrap() ); assert!(matches!(json_to_string.signature(), diff --git a/src/common/function/src/scalars/json/parse_json.rs b/src/common/function/src/scalars/json/parse_json.rs index 71c31c93ac..40914ae791 100644 --- a/src/common/function/src/scalars/json/parse_json.rs +++ b/src/common/function/src/scalars/json/parse_json.rs @@ -15,6 +15,7 @@ use std::fmt::{self, Display}; use common_query::error::{InvalidFuncArgsSnafu, Result, UnsupportedInputDataTypeSnafu}; +use datafusion::arrow::datatypes::DataType; use datafusion_expr::{Signature, Volatility}; use datatypes::data_type::ConcreteDataType; use datatypes::prelude::VectorRef; @@ -35,8 +36,8 @@ impl Function for ParseJsonFunction { NAME } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::json_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Binary) } fn signature(&self) -> Signature { @@ -118,10 +119,8 @@ mod tests { assert_eq!("parse_json", parse_json.name()); assert_eq!( - ConcreteDataType::json_datatype(), - parse_json - .return_type(&[ConcreteDataType::json_datatype()]) - .unwrap() + DataType::Binary, + parse_json.return_type(&[DataType::Binary]).unwrap() ); let json_strings = [ diff --git a/src/common/function/src/scalars/matches.rs b/src/common/function/src/scalars/matches.rs index d513427356..c450ac144d 100644 --- a/src/common/function/src/scalars/matches.rs +++ b/src/common/function/src/scalars/matches.rs @@ -58,8 +58,8 @@ impl Function for MatchesFunction { "matches" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::boolean_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Boolean) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/scalars/matches_term.rs b/src/common/function/src/scalars/matches_term.rs index 13001e8919..d4261c2881 100644 --- a/src/common/function/src/scalars/matches_term.rs +++ b/src/common/function/src/scalars/matches_term.rs @@ -17,8 +17,8 @@ use std::iter::repeat_n; use std::sync::Arc; use common_query::error::{InvalidFuncArgsSnafu, Result}; +use datafusion::arrow::datatypes::DataType; use datafusion_expr::{Signature, Volatility}; -use datatypes::prelude::ConcreteDataType; use datatypes::scalars::ScalarVectorBuilder; use datatypes::vectors::{BooleanVector, BooleanVectorBuilder, MutableVector, VectorRef}; use memchr::memmem; @@ -92,8 +92,8 @@ impl Function for MatchesTermFunction { "matches_term" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::boolean_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Boolean) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/scalars/math.rs b/src/common/function/src/scalars/math.rs index 76e289da4c..cde96c2f81 100644 --- a/src/common/function/src/scalars/math.rs +++ b/src/common/function/src/scalars/math.rs @@ -20,9 +20,9 @@ use std::fmt; pub use clamp::{ClampFunction, ClampMaxFunction, ClampMinFunction}; use common_query::error::{GeneralDataFusionSnafu, Result}; +use datafusion::arrow::datatypes::DataType; use datafusion::error::DataFusionError; use datafusion_expr::{Signature, Volatility}; -use datatypes::prelude::ConcreteDataType; use datatypes::vectors::VectorRef; pub use rate::RateFunction; use snafu::ResultExt; @@ -61,7 +61,7 @@ impl Function for RangeFunction { } // The first argument to range_fn is the expression to be evaluated - fn return_type(&self, input_types: &[ConcreteDataType]) -> Result { + fn return_type(&self, input_types: &[DataType]) -> Result { input_types .first() .cloned() diff --git a/src/common/function/src/scalars/math/clamp.rs b/src/common/function/src/scalars/math/clamp.rs index 81ad16bf10..37b17d231f 100644 --- a/src/common/function/src/scalars/math/clamp.rs +++ b/src/common/function/src/scalars/math/clamp.rs @@ -17,10 +17,11 @@ use std::sync::Arc; use common_query::error::{InvalidFuncArgsSnafu, Result}; use datafusion::arrow::array::{ArrayIter, PrimitiveArray}; +use datafusion::arrow::datatypes::DataType as ArrowDataType; use datafusion::logical_expr::Volatility; use datafusion_expr::Signature; use datafusion_expr::type_coercion::aggregates::NUMERICS; -use datatypes::data_type::{ConcreteDataType, DataType}; +use datatypes::data_type::DataType; use datatypes::prelude::VectorRef; use datatypes::types::LogicalPrimitiveType; use datatypes::value::TryAsPrimitive; @@ -67,7 +68,7 @@ impl Function for ClampFunction { CLAMP_NAME } - fn return_type(&self, input_types: &[ConcreteDataType]) -> Result { + fn return_type(&self, input_types: &[ArrowDataType]) -> Result { // Type check is done by `signature` Ok(input_types[0].clone()) } @@ -187,7 +188,7 @@ impl Function for ClampMinFunction { CLAMP_MIN_NAME } - fn return_type(&self, input_types: &[ConcreteDataType]) -> Result { + fn return_type(&self, input_types: &[ArrowDataType]) -> Result { Ok(input_types[0].clone()) } @@ -268,7 +269,7 @@ impl Function for ClampMaxFunction { CLAMP_MAX_NAME } - fn return_type(&self, input_types: &[ConcreteDataType]) -> Result { + fn return_type(&self, input_types: &[ArrowDataType]) -> Result { Ok(input_types[0].clone()) } diff --git a/src/common/function/src/scalars/math/modulo.rs b/src/common/function/src/scalars/math/modulo.rs index a317edcb21..89df6e99bb 100644 --- a/src/common/function/src/scalars/math/modulo.rs +++ b/src/common/function/src/scalars/math/modulo.rs @@ -17,6 +17,7 @@ use std::fmt::Display; use common_query::error; use common_query::error::{ArrowComputeSnafu, InvalidFuncArgsSnafu, Result}; +use datafusion::arrow::datatypes::DataType; use datafusion_expr::type_coercion::aggregates::NUMERICS; use datafusion_expr::{Signature, Volatility}; use datatypes::arrow::compute; @@ -45,13 +46,13 @@ impl Function for ModuloFunction { NAME } - fn return_type(&self, input_types: &[ConcreteDataType]) -> Result { - if input_types.iter().all(ConcreteDataType::is_signed) { - Ok(ConcreteDataType::int64_datatype()) - } else if input_types.iter().all(ConcreteDataType::is_unsigned) { - Ok(ConcreteDataType::uint64_datatype()) + fn return_type(&self, input_types: &[DataType]) -> Result { + if input_types.iter().all(DataType::is_signed_integer) { + Ok(DataType::Int64) + } else if input_types.iter().all(DataType::is_unsigned_integer) { + Ok(DataType::UInt64) } else { - Ok(ConcreteDataType::float64_datatype()) + Ok(DataType::Float64) } } @@ -108,16 +109,12 @@ mod tests { let function = ModuloFunction; assert_eq!("mod", function.name()); assert_eq!( - ConcreteDataType::int64_datatype(), - function - .return_type(&[ConcreteDataType::int64_datatype()]) - .unwrap() + DataType::Int64, + function.return_type(&[DataType::Int64]).unwrap() ); assert_eq!( - ConcreteDataType::int64_datatype(), - function - .return_type(&[ConcreteDataType::int32_datatype()]) - .unwrap() + DataType::Int64, + function.return_type(&[DataType::Int32]).unwrap() ); let nums = vec![18, -17, 5, -6]; @@ -140,16 +137,12 @@ mod tests { let function = ModuloFunction; assert_eq!("mod", function.name()); assert_eq!( - ConcreteDataType::uint64_datatype(), - function - .return_type(&[ConcreteDataType::uint64_datatype()]) - .unwrap() + DataType::UInt64, + function.return_type(&[DataType::UInt64]).unwrap() ); assert_eq!( - ConcreteDataType::uint64_datatype(), - function - .return_type(&[ConcreteDataType::uint32_datatype()]) - .unwrap() + DataType::UInt64, + function.return_type(&[DataType::UInt32]).unwrap() ); let nums: Vec = vec![18, 17, 5, 6]; @@ -172,16 +165,12 @@ mod tests { let function = ModuloFunction; assert_eq!("mod", function.name()); assert_eq!( - ConcreteDataType::float64_datatype(), - function - .return_type(&[ConcreteDataType::float64_datatype()]) - .unwrap() + DataType::Float64, + function.return_type(&[DataType::Float64]).unwrap() ); assert_eq!( - ConcreteDataType::float64_datatype(), - function - .return_type(&[ConcreteDataType::float32_datatype()]) - .unwrap() + DataType::Float64, + function.return_type(&[DataType::Float32]).unwrap() ); let nums = vec![18.0, 17.0, 5.0, 6.0]; diff --git a/src/common/function/src/scalars/math/rate.rs b/src/common/function/src/scalars/math/rate.rs index a528f1d52a..c899cfaf17 100644 --- a/src/common/function/src/scalars/math/rate.rs +++ b/src/common/function/src/scalars/math/rate.rs @@ -20,7 +20,6 @@ use datafusion_expr::type_coercion::aggregates::NUMERICS; use datafusion_expr::{Signature, Volatility}; use datatypes::arrow::compute::kernels::cast; use datatypes::arrow::datatypes::DataType; -use datatypes::prelude::*; use datatypes::vectors::{Helper, VectorRef}; use snafu::ResultExt; @@ -41,8 +40,8 @@ impl Function for RateFunction { "rate" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::float64_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Float64) } fn signature(&self) -> Signature { @@ -84,10 +83,7 @@ mod tests { fn test_rate_function() { let rate = RateFunction; assert_eq!("rate", rate.name()); - assert_eq!( - ConcreteDataType::float64_datatype(), - rate.return_type(&[]).unwrap() - ); + assert_eq!(DataType::Float64, rate.return_type(&[]).unwrap()); assert!(matches!(rate.signature(), Signature { type_signature: TypeSignature::Uniform(2, valid_types), diff --git a/src/common/function/src/scalars/test.rs b/src/common/function/src/scalars/test.rs index 8c1911405e..0623060350 100644 --- a/src/common/function/src/scalars/test.rs +++ b/src/common/function/src/scalars/test.rs @@ -18,7 +18,6 @@ use std::sync::Arc; use common_query::error::Result; use datafusion_expr::{Signature, Volatility}; use datatypes::arrow::datatypes::DataType; -use datatypes::data_type::ConcreteDataType; use datatypes::prelude::VectorRef; use crate::function::{Function, FunctionContext}; @@ -32,8 +31,8 @@ impl Function for TestAndFunction { "test_and" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::boolean_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Boolean) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/scalars/timestamp/to_unixtime.rs b/src/common/function/src/scalars/timestamp/to_unixtime.rs index 5647772ed6..2d4ad0e6e5 100644 --- a/src/common/function/src/scalars/timestamp/to_unixtime.rs +++ b/src/common/function/src/scalars/timestamp/to_unixtime.rs @@ -61,8 +61,8 @@ impl Function for ToUnixtimeFunction { NAME } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::int64_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Int64) } fn signature(&self) -> Signature { @@ -131,7 +131,6 @@ impl fmt::Display for ToUnixtimeFunction { #[cfg(test)] mod tests { use datafusion_expr::TypeSignature; - use datatypes::prelude::ConcreteDataType; use datatypes::value::Value; use datatypes::vectors::{ DateVector, StringVector, TimestampMillisecondVector, TimestampSecondVector, @@ -143,10 +142,7 @@ mod tests { fn test_string_to_unixtime() { let f = ToUnixtimeFunction; assert_eq!("to_unixtime", f.name()); - assert_eq!( - ConcreteDataType::int64_datatype(), - f.return_type(&[]).unwrap() - ); + assert_eq!(DataType::Int64, f.return_type(&[]).unwrap()); assert!(matches!(f.signature(), Signature { diff --git a/src/common/function/src/scalars/uddsketch_calc.rs b/src/common/function/src/scalars/uddsketch_calc.rs index 0e36839f21..2713d8ceaa 100644 --- a/src/common/function/src/scalars/uddsketch_calc.rs +++ b/src/common/function/src/scalars/uddsketch_calc.rs @@ -20,7 +20,6 @@ use std::fmt::Display; use common_query::error::{DowncastVectorSnafu, InvalidFuncArgsSnafu, Result}; use datafusion_expr::{Signature, Volatility}; use datatypes::arrow::datatypes::DataType; -use datatypes::data_type::ConcreteDataType; use datatypes::prelude::Vector; use datatypes::scalars::{ScalarVector, ScalarVectorBuilder}; use datatypes::vectors::{BinaryVector, Float64VectorBuilder, MutableVector, VectorRef}; @@ -59,8 +58,8 @@ impl Function for UddSketchCalcFunction { NAME } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::float64_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Float64) } fn signature(&self) -> Signature { @@ -141,10 +140,8 @@ mod tests { let function = UddSketchCalcFunction; assert_eq!("uddsketch_calc", function.name()); assert_eq!( - ConcreteDataType::float64_datatype(), - function - .return_type(&[ConcreteDataType::float64_datatype()]) - .unwrap() + DataType::Float64, + function.return_type(&[DataType::Float64]).unwrap() ); // Create a test sketch diff --git a/src/common/function/src/scalars/udf.rs b/src/common/function/src/scalars/udf.rs index 867024eb5c..674bc4240e 100644 --- a/src/common/function/src/scalars/udf.rs +++ b/src/common/function/src/scalars/udf.rs @@ -17,10 +17,9 @@ use std::fmt::{Debug, Formatter}; use std::sync::Arc; use common_query::prelude::ColumnarValue; +use datafusion::arrow::datatypes::DataType; use datafusion::logical_expr::{ScalarFunctionArgs, ScalarUDFImpl}; use datafusion_expr::ScalarUDF; -use datatypes::data_type::DataType; -use datatypes::prelude::*; use session::context::QueryContextRef; use crate::function::{FunctionContext, FunctionRef}; @@ -58,16 +57,8 @@ impl ScalarUDFImpl for ScalarUdf { &self.signature } - fn return_type( - &self, - arg_types: &[datatypes::arrow::datatypes::DataType], - ) -> datafusion_common::Result { - let arg_types = arg_types - .iter() - .map(ConcreteDataType::from_arrow_type) - .collect::>(); - let t = self.function.return_type(&arg_types)?; - Ok(t.as_arrow_type()) + fn return_type(&self, arg_types: &[DataType]) -> datafusion_common::Result { + self.function.return_type(arg_types).map_err(Into::into) } fn invoke_with_args( @@ -110,8 +101,9 @@ mod tests { use datafusion::arrow::array::BooleanArray; use datafusion_common::config::ConfigOptions; use datatypes::arrow::datatypes::Field; - use datatypes::data_type::ConcreteDataType; + use datatypes::data_type::{ConcreteDataType, DataType}; use datatypes::prelude::VectorRef; + use datatypes::value::Value; use datatypes::vectors::{BooleanVector, ConstantVector}; use session::context::QueryContextBuilder; diff --git a/src/common/function/src/scalars/vector/convert/parse_vector.rs b/src/common/function/src/scalars/vector/convert/parse_vector.rs index 8cdc55324e..84ec51d959 100644 --- a/src/common/function/src/scalars/vector/convert/parse_vector.rs +++ b/src/common/function/src/scalars/vector/convert/parse_vector.rs @@ -15,8 +15,8 @@ use std::fmt::Display; use common_query::error::{InvalidFuncArgsSnafu, InvalidVectorStringSnafu, Result}; +use datafusion::arrow::datatypes::DataType; use datafusion_expr::{Signature, Volatility}; -use datatypes::prelude::ConcreteDataType; use datatypes::scalars::ScalarVectorBuilder; use datatypes::types::parse_string_to_vector_type_value; use datatypes::vectors::{BinaryVectorBuilder, MutableVector, VectorRef}; @@ -34,8 +34,8 @@ impl Function for ParseVectorFunction { NAME } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::binary_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Binary) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/scalars/vector/convert/vector_to_string.rs b/src/common/function/src/scalars/vector/convert/vector_to_string.rs index b2cbacd59b..58fb22b61a 100644 --- a/src/common/function/src/scalars/vector/convert/vector_to_string.rs +++ b/src/common/function/src/scalars/vector/convert/vector_to_string.rs @@ -15,9 +15,9 @@ use std::fmt::Display; use common_query::error::{InvalidFuncArgsSnafu, Result}; +use datafusion::arrow::datatypes::DataType; use datafusion_expr::type_coercion::aggregates::BINARYS; use datafusion_expr::{Signature, Volatility}; -use datatypes::prelude::ConcreteDataType; use datatypes::scalars::ScalarVectorBuilder; use datatypes::types::vector_type_value_to_string; use datatypes::value::Value; @@ -36,8 +36,8 @@ impl Function for VectorToStringFunction { NAME } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::string_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Utf8) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/scalars/vector/distance.rs b/src/common/function/src/scalars/vector/distance.rs index 1dc2f8748c..864e2405b0 100644 --- a/src/common/function/src/scalars/vector/distance.rs +++ b/src/common/function/src/scalars/vector/distance.rs @@ -22,7 +22,6 @@ use std::fmt::Display; use common_query::error::{InvalidFuncArgsSnafu, Result}; use datafusion_expr::Signature; use datatypes::arrow::datatypes::DataType; -use datatypes::prelude::ConcreteDataType; use datatypes::scalars::ScalarVectorBuilder; use datatypes::vectors::{Float32VectorBuilder, MutableVector, VectorRef}; use snafu::ensure; @@ -44,8 +43,8 @@ macro_rules! define_distance_function { $display_name } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::float32_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Float32) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/scalars/vector/elem_product.rs b/src/common/function/src/scalars/vector/elem_product.rs index eee0363de5..954f0f73a6 100644 --- a/src/common/function/src/scalars/vector/elem_product.rs +++ b/src/common/function/src/scalars/vector/elem_product.rs @@ -15,10 +15,10 @@ use std::borrow::Cow; use std::fmt::Display; -use common_query::error::InvalidFuncArgsSnafu; +use common_query::error::{InvalidFuncArgsSnafu, Result}; +use datafusion::arrow::datatypes::DataType; use datafusion::logical_expr_common::type_coercion::aggregates::{BINARYS, STRINGS}; use datafusion_expr::{Signature, TypeSignature, Volatility}; -use datatypes::prelude::ConcreteDataType; use datatypes::scalars::ScalarVectorBuilder; use datatypes::vectors::{Float32VectorBuilder, MutableVector, VectorRef}; use nalgebra::DVectorView; @@ -50,11 +50,8 @@ impl Function for ElemProductFunction { NAME } - fn return_type( - &self, - _input_types: &[ConcreteDataType], - ) -> common_query::error::Result { - Ok(ConcreteDataType::float32_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Float32) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/scalars/vector/elem_sum.rs b/src/common/function/src/scalars/vector/elem_sum.rs index 06e39698c3..8bea6d6c32 100644 --- a/src/common/function/src/scalars/vector/elem_sum.rs +++ b/src/common/function/src/scalars/vector/elem_sum.rs @@ -15,10 +15,10 @@ use std::borrow::Cow; use std::fmt::Display; -use common_query::error::InvalidFuncArgsSnafu; +use common_query::error::{InvalidFuncArgsSnafu, Result}; +use datafusion::arrow::datatypes::DataType; use datafusion_expr::type_coercion::aggregates::{BINARYS, STRINGS}; use datafusion_expr::{Signature, TypeSignature, Volatility}; -use datatypes::prelude::ConcreteDataType; use datatypes::scalars::ScalarVectorBuilder; use datatypes::vectors::{Float32VectorBuilder, MutableVector, VectorRef}; use nalgebra::DVectorView; @@ -37,11 +37,8 @@ impl Function for ElemSumFunction { NAME } - fn return_type( - &self, - _input_types: &[ConcreteDataType], - ) -> common_query::error::Result { - Ok(ConcreteDataType::float32_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Float32) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/scalars/vector/scalar_add.rs b/src/common/function/src/scalars/vector/scalar_add.rs index 6eaf9ade22..81a532132e 100644 --- a/src/common/function/src/scalars/vector/scalar_add.rs +++ b/src/common/function/src/scalars/vector/scalar_add.rs @@ -18,7 +18,6 @@ use std::fmt::Display; use common_query::error::{InvalidFuncArgsSnafu, Result}; use datafusion_expr::Signature; use datatypes::arrow::datatypes::DataType; -use datatypes::prelude::ConcreteDataType; use datatypes::scalars::ScalarVectorBuilder; use datatypes::vectors::{BinaryVectorBuilder, MutableVector, VectorRef}; use nalgebra::DVectorView; @@ -60,8 +59,8 @@ impl Function for ScalarAddFunction { NAME } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::binary_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Binary) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/scalars/vector/scalar_mul.rs b/src/common/function/src/scalars/vector/scalar_mul.rs index 32f1f5b64a..8985987331 100644 --- a/src/common/function/src/scalars/vector/scalar_mul.rs +++ b/src/common/function/src/scalars/vector/scalar_mul.rs @@ -18,7 +18,6 @@ use std::fmt::Display; use common_query::error::{InvalidFuncArgsSnafu, Result}; use datafusion_expr::Signature; use datatypes::arrow::datatypes::DataType; -use datatypes::prelude::ConcreteDataType; use datatypes::scalars::ScalarVectorBuilder; use datatypes::vectors::{BinaryVectorBuilder, MutableVector, VectorRef}; use nalgebra::DVectorView; @@ -60,8 +59,8 @@ impl Function for ScalarMulFunction { NAME } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::binary_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Binary) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/scalars/vector/vector_add.rs b/src/common/function/src/scalars/vector/vector_add.rs index 0ec72fe0a7..3ecab37938 100644 --- a/src/common/function/src/scalars/vector/vector_add.rs +++ b/src/common/function/src/scalars/vector/vector_add.rs @@ -15,10 +15,9 @@ use std::borrow::Cow; use std::fmt::Display; -use common_query::error::InvalidFuncArgsSnafu; +use common_query::error::{InvalidFuncArgsSnafu, Result}; use datafusion_expr::Signature; use datatypes::arrow::datatypes::DataType; -use datatypes::prelude::ConcreteDataType; use datatypes::scalars::ScalarVectorBuilder; use datatypes::vectors::{BinaryVectorBuilder, MutableVector, VectorRef}; use nalgebra::DVectorView; @@ -51,11 +50,8 @@ impl Function for VectorAddFunction { NAME } - fn return_type( - &self, - _input_types: &[ConcreteDataType], - ) -> common_query::error::Result { - Ok(ConcreteDataType::binary_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Binary) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/scalars/vector/vector_dim.rs b/src/common/function/src/scalars/vector/vector_dim.rs index 4afec35be5..2f877ed9ea 100644 --- a/src/common/function/src/scalars/vector/vector_dim.rs +++ b/src/common/function/src/scalars/vector/vector_dim.rs @@ -15,10 +15,10 @@ use std::borrow::Cow; use std::fmt::Display; -use common_query::error::InvalidFuncArgsSnafu; +use common_query::error::{InvalidFuncArgsSnafu, Result}; +use datafusion::arrow::datatypes::DataType; use datafusion::logical_expr_common::type_coercion::aggregates::{BINARYS, STRINGS}; use datafusion_expr::{Signature, TypeSignature, Volatility}; -use datatypes::prelude::ConcreteDataType; use datatypes::scalars::ScalarVectorBuilder; use datatypes::vectors::{MutableVector, UInt64VectorBuilder, VectorRef}; use snafu::ensure; @@ -49,11 +49,8 @@ impl Function for VectorDimFunction { NAME } - fn return_type( - &self, - _input_types: &[ConcreteDataType], - ) -> common_query::error::Result { - Ok(ConcreteDataType::uint64_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::UInt64) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/scalars/vector/vector_div.rs b/src/common/function/src/scalars/vector/vector_div.rs index 49d152ca83..2e2d6898a1 100644 --- a/src/common/function/src/scalars/vector/vector_div.rs +++ b/src/common/function/src/scalars/vector/vector_div.rs @@ -18,7 +18,6 @@ use std::fmt::Display; use common_query::error::{InvalidFuncArgsSnafu, Result}; use datafusion_expr::Signature; use datatypes::arrow::datatypes::DataType; -use datatypes::prelude::ConcreteDataType; use datatypes::scalars::ScalarVectorBuilder; use datatypes::vectors::{BinaryVectorBuilder, MutableVector, VectorRef}; use nalgebra::DVectorView; @@ -52,8 +51,8 @@ impl Function for VectorDivFunction { NAME } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::binary_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Binary) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/scalars/vector/vector_kth_elem.rs b/src/common/function/src/scalars/vector/vector_kth_elem.rs index 2906146fc8..d1d110c47c 100644 --- a/src/common/function/src/scalars/vector/vector_kth_elem.rs +++ b/src/common/function/src/scalars/vector/vector_kth_elem.rs @@ -18,7 +18,6 @@ use std::fmt::Display; use common_query::error::{InvalidFuncArgsSnafu, Result}; use datafusion_expr::Signature; use datatypes::arrow::datatypes::DataType; -use datatypes::prelude::ConcreteDataType; use datatypes::scalars::ScalarVectorBuilder; use datatypes::vectors::{Float32VectorBuilder, MutableVector, VectorRef}; use snafu::ensure; @@ -53,11 +52,8 @@ impl Function for VectorKthElemFunction { NAME } - fn return_type( - &self, - _input_types: &[ConcreteDataType], - ) -> common_query::error::Result { - Ok(ConcreteDataType::float32_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Float32) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/scalars/vector/vector_mul.rs b/src/common/function/src/scalars/vector/vector_mul.rs index 69847ad098..99bf3efe5d 100644 --- a/src/common/function/src/scalars/vector/vector_mul.rs +++ b/src/common/function/src/scalars/vector/vector_mul.rs @@ -18,7 +18,6 @@ use std::fmt::Display; use common_query::error::{InvalidFuncArgsSnafu, Result}; use datafusion_expr::Signature; use datatypes::arrow::datatypes::DataType; -use datatypes::prelude::ConcreteDataType; use datatypes::scalars::ScalarVectorBuilder; use datatypes::vectors::{BinaryVectorBuilder, MutableVector, VectorRef}; use nalgebra::DVectorView; @@ -52,8 +51,8 @@ impl Function for VectorMulFunction { NAME } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::binary_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Binary) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/scalars/vector/vector_norm.rs b/src/common/function/src/scalars/vector/vector_norm.rs index 1d99d27022..bd82efbdbc 100644 --- a/src/common/function/src/scalars/vector/vector_norm.rs +++ b/src/common/function/src/scalars/vector/vector_norm.rs @@ -16,9 +16,9 @@ use std::borrow::Cow; use std::fmt::Display; use common_query::error::{InvalidFuncArgsSnafu, Result}; +use datafusion::arrow::datatypes::DataType; use datafusion::logical_expr_common::type_coercion::aggregates::{BINARYS, STRINGS}; use datafusion_expr::{Signature, TypeSignature, Volatility}; -use datatypes::prelude::ConcreteDataType; use datatypes::scalars::ScalarVectorBuilder; use datatypes::vectors::{BinaryVectorBuilder, MutableVector, VectorRef}; use nalgebra::DVectorView; @@ -52,8 +52,8 @@ impl Function for VectorNormFunction { NAME } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::binary_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Binary) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/scalars/vector/vector_sub.rs b/src/common/function/src/scalars/vector/vector_sub.rs index f5e948f42c..a849df2a3d 100644 --- a/src/common/function/src/scalars/vector/vector_sub.rs +++ b/src/common/function/src/scalars/vector/vector_sub.rs @@ -15,10 +15,9 @@ use std::borrow::Cow; use std::fmt::Display; -use common_query::error::InvalidFuncArgsSnafu; +use common_query::error::{InvalidFuncArgsSnafu, Result}; use datafusion_expr::Signature; use datatypes::arrow::datatypes::DataType; -use datatypes::prelude::ConcreteDataType; use datatypes::scalars::ScalarVectorBuilder; use datatypes::vectors::{BinaryVectorBuilder, MutableVector, VectorRef}; use nalgebra::DVectorView; @@ -51,11 +50,8 @@ impl Function for VectorSubFunction { NAME } - fn return_type( - &self, - _input_types: &[ConcreteDataType], - ) -> common_query::error::Result { - Ok(ConcreteDataType::binary_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Binary) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/scalars/vector/vector_subvector.rs b/src/common/function/src/scalars/vector/vector_subvector.rs index 3d4378f8e5..cda358d0f9 100644 --- a/src/common/function/src/scalars/vector/vector_subvector.rs +++ b/src/common/function/src/scalars/vector/vector_subvector.rs @@ -18,7 +18,6 @@ use std::fmt::Display; use common_query::error::{InvalidFuncArgsSnafu, Result}; use datafusion_expr::{Signature, TypeSignature, Volatility}; use datatypes::arrow::datatypes::DataType; -use datatypes::prelude::ConcreteDataType; use datatypes::scalars::ScalarVectorBuilder; use datatypes::vectors::{BinaryVectorBuilder, MutableVector, VectorRef}; use snafu::ensure; @@ -52,8 +51,8 @@ impl Function for VectorSubvectorFunction { NAME } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::binary_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Binary) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/system/build.rs b/src/common/function/src/system/build.rs index ef51e63cc0..5517994a52 100644 --- a/src/common/function/src/system/build.rs +++ b/src/common/function/src/system/build.rs @@ -16,8 +16,8 @@ use std::fmt; use std::sync::Arc; use common_query::error::Result; +use datafusion::arrow::datatypes::DataType; use datafusion_expr::{Signature, Volatility}; -use datatypes::prelude::*; use datatypes::vectors::{StringVector, VectorRef}; use crate::function::{Function, FunctionContext}; @@ -37,8 +37,8 @@ impl Function for BuildFunction { "build" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::string_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Utf8) } fn signature(&self) -> Signature { @@ -61,10 +61,7 @@ mod tests { fn test_build_function() { let build = BuildFunction; assert_eq!("build", build.name()); - assert_eq!( - ConcreteDataType::string_datatype(), - build.return_type(&[]).unwrap() - ); + assert_eq!(DataType::Utf8, build.return_type(&[]).unwrap()); assert_eq!(build.signature(), Signature::nullary(Volatility::Immutable)); let build_info = common_version::build_info().to_string(); let vector = build.eval(&FunctionContext::default(), &[]).unwrap(); diff --git a/src/common/function/src/system/database.rs b/src/common/function/src/system/database.rs index 4eda1bc560..d28975afb4 100644 --- a/src/common/function/src/system/database.rs +++ b/src/common/function/src/system/database.rs @@ -16,8 +16,9 @@ use std::fmt::{self}; use std::sync::Arc; use common_query::error::Result; +use datafusion::arrow::datatypes::DataType; use datafusion_expr::{Signature, Volatility}; -use datatypes::prelude::{ConcreteDataType, ScalarVector}; +use datatypes::prelude::ScalarVector; use datatypes::vectors::{StringVector, UInt32Vector, VectorRef}; use derive_more::Display; @@ -53,8 +54,8 @@ impl Function for DatabaseFunction { DATABASE_FUNCTION_NAME } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::string_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Utf8) } fn signature(&self) -> Signature { @@ -75,8 +76,8 @@ impl Function for CurrentSchemaFunction { CURRENT_SCHEMA_FUNCTION_NAME } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::string_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Utf8) } fn signature(&self) -> Signature { @@ -95,8 +96,8 @@ impl Function for SessionUserFunction { SESSION_USER_FUNCTION_NAME } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::string_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Utf8) } fn signature(&self) -> Signature { @@ -115,8 +116,8 @@ impl Function for ReadPreferenceFunction { READ_PREFERENCE_FUNCTION_NAME } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::string_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Utf8) } fn signature(&self) -> Signature { @@ -135,8 +136,8 @@ impl Function for PgBackendPidFunction { PG_BACKEND_PID } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::uint64_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::UInt64) } fn signature(&self) -> Signature { @@ -155,8 +156,8 @@ impl Function for ConnectionIdFunction { CONNECTION_ID } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::uint64_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::UInt64) } fn signature(&self) -> Signature { @@ -205,10 +206,7 @@ mod tests { fn test_build_function() { let build = DatabaseFunction; assert_eq!("database", build.name()); - assert_eq!( - ConcreteDataType::string_datatype(), - build.return_type(&[]).unwrap() - ); + assert_eq!(DataType::Utf8, build.return_type(&[]).unwrap()); assert_eq!(build.signature(), Signature::nullary(Volatility::Immutable)); let query_ctx = QueryContextBuilder::default() diff --git a/src/common/function/src/system/pg_catalog/pg_get_userbyid.rs b/src/common/function/src/system/pg_catalog/pg_get_userbyid.rs index c6bf1ee1ab..9b32009a1a 100644 --- a/src/common/function/src/system/pg_catalog/pg_get_userbyid.rs +++ b/src/common/function/src/system/pg_catalog/pg_get_userbyid.rs @@ -16,8 +16,9 @@ use std::fmt::{self}; use std::sync::Arc; use common_query::error::Result; +use datafusion::arrow::datatypes::DataType as ArrowDataType; use datafusion_expr::{Signature, Volatility}; -use datatypes::prelude::{ConcreteDataType, DataType, VectorRef}; +use datatypes::prelude::{DataType, VectorRef}; use datatypes::types::LogicalPrimitiveType; use datatypes::with_match_primitive_type_id; use num_traits::AsPrimitive; @@ -41,8 +42,8 @@ impl Function for PGGetUserByIdFunction { NAME } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::string_datatype()) + fn return_type(&self, _: &[ArrowDataType]) -> Result { + Ok(ArrowDataType::Utf8) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/system/pg_catalog/table_is_visible.rs b/src/common/function/src/system/pg_catalog/table_is_visible.rs index 2c1d0e9bda..ba9c3891df 100644 --- a/src/common/function/src/system/pg_catalog/table_is_visible.rs +++ b/src/common/function/src/system/pg_catalog/table_is_visible.rs @@ -16,8 +16,9 @@ use std::fmt::{self}; use std::sync::Arc; use common_query::error::Result; +use datafusion::arrow::datatypes::DataType as ArrowDataType; use datafusion_expr::{Signature, Volatility}; -use datatypes::prelude::{ConcreteDataType, DataType, VectorRef}; +use datatypes::prelude::{DataType, VectorRef}; use datatypes::types::LogicalPrimitiveType; use datatypes::with_match_primitive_type_id; use num_traits::AsPrimitive; @@ -41,8 +42,8 @@ impl Function for PGTableIsVisibleFunction { NAME } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::boolean_datatype()) + fn return_type(&self, _: &[ArrowDataType]) -> Result { + Ok(ArrowDataType::Boolean) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/system/pg_catalog/version.rs b/src/common/function/src/system/pg_catalog/version.rs index c8c8095749..e25e840751 100644 --- a/src/common/function/src/system/pg_catalog/version.rs +++ b/src/common/function/src/system/pg_catalog/version.rs @@ -16,8 +16,8 @@ use std::fmt; use std::sync::Arc; use common_query::error::Result; +use datafusion::arrow::datatypes::DataType; use datafusion_expr::{Signature, Volatility}; -use datatypes::data_type::ConcreteDataType; use datatypes::vectors::{StringVector, VectorRef}; use crate::function::{Function, FunctionContext}; @@ -36,8 +36,8 @@ impl Function for PGVersionFunction { crate::pg_catalog_func_fullname!("version") } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::string_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Utf8) } fn signature(&self) -> Signature { diff --git a/src/common/function/src/system/timezone.rs b/src/common/function/src/system/timezone.rs index 923895c9f7..22bc395453 100644 --- a/src/common/function/src/system/timezone.rs +++ b/src/common/function/src/system/timezone.rs @@ -16,8 +16,9 @@ use std::fmt::{self}; use std::sync::Arc; use common_query::error::Result; +use datafusion::arrow::datatypes::DataType; use datafusion_expr::{Signature, Volatility}; -use datatypes::prelude::{ConcreteDataType, ScalarVector}; +use datatypes::prelude::ScalarVector; use datatypes::vectors::{StringVector, VectorRef}; use crate::function::{Function, FunctionContext}; @@ -33,8 +34,8 @@ impl Function for TimezoneFunction { NAME } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::string_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Utf8) } fn signature(&self) -> Signature { @@ -65,10 +66,7 @@ mod tests { fn test_build_function() { let build = TimezoneFunction; assert_eq!("timezone", build.name()); - assert_eq!( - ConcreteDataType::string_datatype(), - build.return_type(&[]).unwrap() - ); + assert_eq!(DataType::Utf8, build.return_type(&[]).unwrap()); assert_eq!(build.signature(), Signature::nullary(Volatility::Immutable)); let query_ctx = QueryContextBuilder::default().build().into(); diff --git a/src/common/function/src/system/version.rs b/src/common/function/src/system/version.rs index cd2db99765..d1ef84baea 100644 --- a/src/common/function/src/system/version.rs +++ b/src/common/function/src/system/version.rs @@ -16,8 +16,8 @@ use std::fmt; use std::sync::Arc; use common_query::error::Result; +use datafusion::arrow::datatypes::DataType; use datafusion_expr::{Signature, Volatility}; -use datatypes::data_type::ConcreteDataType; use datatypes::vectors::{StringVector, VectorRef}; use session::context::Channel; @@ -37,8 +37,8 @@ impl Function for VersionFunction { "version" } - fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result { - Ok(ConcreteDataType::string_datatype()) + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Utf8) } fn signature(&self) -> Signature { diff --git a/src/flow/src/transform.rs b/src/flow/src/transform.rs index f3ebf0ca71..8d2c1ec8a4 100644 --- a/src/flow/src/transform.rs +++ b/src/flow/src/transform.rs @@ -18,9 +18,9 @@ use std::sync::Arc; use common_error::ext::BoxedError; use common_function::function::{FunctionContext, FunctionRef}; +use datafusion::arrow::datatypes::{DataType, TimeUnit}; use datafusion_expr::{Signature, Volatility}; use datafusion_substrait::extensions::Extensions; -use datatypes::data_type::ConcreteDataType as CDT; use query::QueryEngine; use serde::{Deserialize, Serialize}; use snafu::ResultExt; @@ -142,8 +142,8 @@ impl common_function::function::Function for TumbleFunction { &self.name } - fn return_type(&self, _input_types: &[CDT]) -> common_query::error::Result { - Ok(CDT::timestamp_millisecond_datatype()) + fn return_type(&self, _: &[DataType]) -> common_query::error::Result { + Ok(DataType::Timestamp(TimeUnit::Millisecond, None)) } fn signature(&self) -> Signature { @@ -170,6 +170,7 @@ mod test { use catalog::RegisterTableRequest; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, NUMBERS_TABLE_ID}; + use datatypes::data_type::ConcreteDataType as CDT; use datatypes::prelude::*; use datatypes::schema::Schema; use datatypes::timestamp::TimestampMillisecond; diff --git a/src/flow/src/transform/aggr.rs b/src/flow/src/transform/aggr.rs index d0b9945f05..82d814ed2c 100644 --- a/src/flow/src/transform/aggr.rs +++ b/src/flow/src/transform/aggr.rs @@ -387,6 +387,7 @@ mod test { use bytes::BytesMut; use common_time::{IntervalMonthDayNano, Timestamp}; + use datatypes::data_type::ConcreteDataType as CDT; use datatypes::prelude::ConcreteDataType; use datatypes::value::Value; use pretty_assertions::assert_eq; @@ -395,7 +396,6 @@ mod test { use crate::expr::{BinaryFunc, DfScalarFunction, GlobalId, RawDfScalarFn}; use crate::plan::{Plan, TypedPlan}; use crate::repr::{ColumnType, RelationType}; - use crate::transform::CDT; use crate::transform::test::{create_test_ctx, create_test_query_engine, sql_to_substrait}; #[tokio::test] diff --git a/src/flow/src/transform/plan.rs b/src/flow/src/transform/plan.rs index dd8d5bb334..a87dce0d26 100644 --- a/src/flow/src/transform/plan.rs +++ b/src/flow/src/transform/plan.rs @@ -232,6 +232,7 @@ impl TypedPlan { #[cfg(test)] mod test { + use datatypes::data_type::ConcreteDataType as CDT; use datatypes::prelude::ConcreteDataType; use pretty_assertions::assert_eq; @@ -239,7 +240,6 @@ mod test { use crate::expr::GlobalId; use crate::plan::{Plan, TypedPlan}; use crate::repr::{ColumnType, RelationType}; - use crate::transform::CDT; use crate::transform::test::{create_test_ctx, create_test_query_engine, sql_to_substrait}; #[tokio::test] diff --git a/src/servers/src/mysql/writer.rs b/src/servers/src/mysql/writer.rs index f571995b99..59540e58e7 100644 --- a/src/servers/src/mysql/writer.rs +++ b/src/servers/src/mysql/writer.rs @@ -23,6 +23,7 @@ use datatypes::prelude::{ConcreteDataType, Value}; use datatypes::schema::SchemaRef; use datatypes::types::json_type_value_to_string; use futures::StreamExt; +use itertools::Itertools; use opensrv_mysql::{ Column, ColumnFlags, ColumnType, ErrorKind, OkResponse, QueryResultWriter, RowWriter, }; @@ -30,7 +31,7 @@ use session::context::QueryContextRef; use snafu::prelude::*; use tokio::io::AsyncWrite; -use crate::error::{self, ConvertSqlValueSnafu, Error, Result}; +use crate::error::{self, ConvertSqlValueSnafu, Result}; use crate::metrics::*; /// Try to write multiple output to the writer if possible. @@ -232,14 +233,10 @@ impl<'a, W: AsyncWrite + Unpin> MysqlResultWriter<'a, W> { row_writer.write_col(v.to_iso8601_string())? } Value::Duration(v) => row_writer.write_col(v.to_std_duration())?, - Value::List(_) => { - return Err(Error::Internal { - err_msg: format!( - "cannot write value {:?} in mysql protocol: unimplemented", - &value - ), - }); - } + Value::List(v) => row_writer.write_col(format!( + "[{}]", + v.items().iter().map(|x| x.to_string()).join(", ") + ))?, Value::Time(v) => row_writer .write_col(v.to_timezone_aware_string(Some(&query_context.timezone())))?, Value::Decimal128(v) => row_writer.write_col(v.to_string())?, @@ -295,6 +292,7 @@ pub(crate) fn create_mysql_column( ConcreteDataType::Decimal128(_) => Ok(ColumnType::MYSQL_TYPE_DECIMAL), ConcreteDataType::Json(_) => Ok(ColumnType::MYSQL_TYPE_JSON), ConcreteDataType::Vector(_) => Ok(ColumnType::MYSQL_TYPE_BLOB), + ConcreteDataType::List(_) => Ok(ColumnType::MYSQL_TYPE_VARCHAR), _ => error::UnsupportedDataTypeSnafu { data_type, reason: "not implemented", diff --git a/src/servers/tests/mysql/mysql_writer_test.rs b/src/servers/tests/mysql/mysql_writer_test.rs index dff541a762..8737fdf584 100644 --- a/src/servers/tests/mysql/mysql_writer_test.rs +++ b/src/servers/tests/mysql/mysql_writer_test.rs @@ -14,8 +14,7 @@ use std::sync::Arc; -use datatypes::prelude::*; -use datatypes::schema::{ColumnSchema, Schema}; +use datatypes::schema::Schema; use servers::mysql::writer::create_mysql_column_def; use crate::mysql::{TestingData, all_datatype_testing_data}; @@ -37,12 +36,4 @@ fn test_create_mysql_column_def() { let expected_coltype = mysql_columns_def[i]; assert_eq!(column_def.coltype, expected_coltype); } - - let column_schemas = vec![ColumnSchema::new( - "lists", - ConcreteDataType::list_datatype(ConcreteDataType::string_datatype()), - true, - )]; - let schema = Arc::new(Schema::new(column_schemas)); - assert!(create_mysql_column_def(&schema).is_err()); }