Skip to content

Commit

Permalink
Add serialization of ScalarValue::IntervalMonthDayNano (#3535)
Browse files Browse the repository at this point in the history
  • Loading branch information
@alamb
alamb authored Sep 22, 2022
1 parent 7033c0e commit 52ff7ca
Show file tree
Hide file tree
Showing 4 changed files with 72 additions and 9 deletions.
15 changes: 12 additions & 3 deletions datafusion/proto/proto/datafusion.proto
View file
Original file line number Diff line number Diff line change
Expand Up @@ -733,9 +733,18 @@ message ScalarDictionaryValue {
ScalarValue value = 2;
}

message IntervalMonthDayNanoValue {
int32 months = 1;
int32 days = 2;
int64 nanos = 3;
}


message ScalarValue{
oneof value {
// Null value of any type (type is encoded)
PrimitiveScalarType null_value = 19;

bool bool_value = 1;
string utf8_value = 2;
string large_utf8_value = 3;
Expand All @@ -754,7 +763,6 @@ message ScalarValue{
ScalarListValue list_value = 17;
ScalarType null_list_value = 18;

PrimitiveScalarType null_value = 19;
Decimal128 decimal128_value = 20;
int64 date_64_value = 21;
int32 interval_yearmonth_value = 24;
Expand All @@ -764,6 +772,7 @@ message ScalarValue{
bytes binary_value = 28;
bytes large_binary_value = 29;
int64 time64_value = 30;
IntervalMonthDayNanoValue interval_month_day_nano = 31;
}
}

Expand Down Expand Up @@ -794,13 +803,13 @@ enum PrimitiveScalarType{
TIMESTAMP_MICROSECOND = 14;
TIMESTAMP_NANOSECOND = 15;
NULL = 16;

DECIMAL128 = 17;
DATE64 = 20;
TIMESTAMP_SECOND = 21;
TIMESTAMP_MILLISECOND = 22;
INTERVAL_YEARMONTH = 23;
INTERVAL_DAYTIME = 24;
INTERVAL_MONTHDAYNANO = 28;

BINARY = 25;
LARGE_BINARY = 26;
Expand All @@ -822,7 +831,7 @@ message ScalarListType{

// Broke out into multiple message types so that type
// metadata did not need to be in separate message
//All types that are of the empty message types contain no additional metadata
// All types that are of the empty message types contain no additional metadata
// about the type
message ArrowType{
oneof arrow_type_enum{
Expand Down
24 changes: 23 additions & 1 deletion datafusion/proto/src/from_proto.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,9 @@ use crate::protobuf::{
CubeNode, GroupingSetNode, OptimizedLogicalPlanType, OptimizedPhysicalPlanType,
RollupNode,
};
use arrow::datatypes::{DataType, Field, IntervalUnit, Schema, TimeUnit, UnionMode};
use arrow::datatypes::{
DataType, Field, IntervalMonthDayNanoType, IntervalUnit, Schema, TimeUnit, UnionMode,
};
use datafusion::logical_plan::FunctionRegistry;
use datafusion_common::{
Column, DFField, DFSchema, DFSchemaRef, DataFusionError, ScalarValue,
Expand Down Expand Up @@ -245,6 +247,9 @@ impl From<protobuf::PrimitiveScalarType> for DataType {
protobuf::PrimitiveScalarType::IntervalDaytime => {
DataType::Interval(IntervalUnit::DayTime)
}
protobuf::PrimitiveScalarType::IntervalMonthdaynano => {
DataType::Interval(IntervalUnit::MonthDayNano)
}
}
}
}
Expand Down Expand Up @@ -666,6 +671,7 @@ impl TryFrom<&protobuf::PrimitiveScalarType> for ScalarValue {
}
PrimitiveScalarType::IntervalYearmonth => Self::IntervalYearMonth(None),
PrimitiveScalarType::IntervalDaytime => Self::IntervalDayTime(None),
PrimitiveScalarType::IntervalMonthdaynano => Self::IntervalMonthDayNano(None),
})
}
}
Expand Down Expand Up @@ -805,6 +811,9 @@ impl TryFrom<&protobuf::ScalarValue> for ScalarValue {
}
Value::BinaryValue(v) => Self::Binary(Some(v.clone())),
Value::LargeBinaryValue(v) => Self::LargeBinary(Some(v.clone())),
Value::IntervalMonthDayNano(v) => Self::IntervalMonthDayNano(Some(
IntervalMonthDayNanoType::make_value(v.months, v.days, v.nanos),
)),
})
}
}
Expand Down Expand Up @@ -1510,6 +1519,9 @@ fn typechecked_scalar_value_conversion(
PrimitiveScalarType::IntervalDaytime => {
ScalarValue::IntervalDayTime(None)
}
PrimitiveScalarType::IntervalMonthdaynano => {
ScalarValue::IntervalMonthDayNano(None)
}
PrimitiveScalarType::Binary => ScalarValue::Binary(None),
PrimitiveScalarType::LargeBinary => ScalarValue::LargeBinary(None),
};
Expand All @@ -1535,6 +1547,16 @@ fn typechecked_scalar_value_conversion(
(Value::IntervalDaytimeValue(v), PrimitiveScalarType::IntervalDaytime) => {
ScalarValue::IntervalDayTime(Some(*v))
}
(Value::IntervalMonthDayNano(v), PrimitiveScalarType::IntervalMonthdaynano) => {
let protobuf::IntervalMonthDayNanoValue {
months,
days,
nanos,
} = v;
ScalarValue::IntervalMonthDayNano(Some(IntervalMonthDayNanoType::make_value(
*months, *days, *nanos,
)))
}
_ => return Err(proto_error("Could not convert to the proper type")),
})
}
Expand Down
22 changes: 21 additions & 1 deletion datafusion/proto/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,10 @@ mod roundtrip_tests {
use arrow::datatypes::Schema;
use arrow::{
array::ArrayRef,
datatypes::{DataType, Field, IntervalUnit, TimeUnit, UnionMode},
datatypes::{
DataType, Field, IntervalDayTimeType, IntervalMonthDayNanoType, IntervalUnit,
TimeUnit, UnionMode,
},
};
use datafusion::logical_plan::create_udaf;
use datafusion::physical_plan::functions::make_scalar_function;
Expand Down Expand Up @@ -452,6 +455,23 @@ mod roundtrip_tests {
ScalarValue::TimestampSecond(Some(i64::MAX), None),
ScalarValue::TimestampSecond(Some(0), Some("UTC".to_string())),
ScalarValue::TimestampSecond(None, None),
ScalarValue::IntervalDayTime(Some(IntervalDayTimeType::make_value(0, 0))),
ScalarValue::IntervalDayTime(Some(IntervalDayTimeType::make_value(1, 2))),
ScalarValue::IntervalDayTime(Some(IntervalDayTimeType::make_value(
i32::MAX,
i32::MAX,
))),
ScalarValue::IntervalDayTime(None),
ScalarValue::IntervalMonthDayNano(Some(
IntervalMonthDayNanoType::make_value(0, 0, 0),
)),
ScalarValue::IntervalMonthDayNano(Some(
IntervalMonthDayNanoType::make_value(1, 2, 3),
)),
ScalarValue::IntervalMonthDayNano(Some(
IntervalMonthDayNanoType::make_value(i32::MAX, i32::MAX, i64::MAX),
)),
ScalarValue::IntervalMonthDayNano(None),
ScalarValue::new_list(
Some(vec![
ScalarValue::Float32(Some(-213.1)),
Expand Down
20 changes: 16 additions & 4 deletions datafusion/proto/src/to_proto.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,8 @@ use crate::protobuf::{
OptimizedPhysicalPlanType, RollupNode,
};
use arrow::datatypes::{
DataType, Field, IntervalUnit, Schema, SchemaRef, TimeUnit, UnionMode,
DataType, Field, IntervalMonthDayNanoType, IntervalUnit, Schema, SchemaRef, TimeUnit,
UnionMode,
};
use datafusion_common::{Column, DFField, DFSchemaRef, ScalarValue};
use datafusion_expr::expr::GroupingSet;
Expand Down Expand Up @@ -1197,9 +1198,20 @@ impl TryFrom<&ScalarValue> for protobuf::ScalarValue {
})
}

datafusion::scalar::ScalarValue::IntervalMonthDayNano(_) => {
// not yet implemented (TODO file ticket)
return Err(Error::invalid_scalar_value(val));
datafusion::scalar::ScalarValue::IntervalMonthDayNano(v) => {
let value = if let Some(v) = v {
let (months, days, nanos) = IntervalMonthDayNanoType::to_parts(*v);
Value::IntervalMonthDayNano(protobuf::IntervalMonthDayNanoValue {
months,
days,
nanos,
})
} else {
let null_arrow_type = PrimitiveScalarType::IntervalMonthdaynano;
protobuf::scalar_value::Value::NullValue(null_arrow_type as i32)
};

protobuf::ScalarValue { value: Some(value) }
}

datafusion::scalar::ScalarValue::Struct(_, _) => {
Expand Down

0 comments on commit 52ff7ca

Please sign in to comment.