Skip to content

Commit 6d32472

Browse files
zhengruifengHyukjinKwon
authored andcommitted
[SPARK-49112][CONNECT][TEST] Make createLocalRelationProto support TimestampType
### What changes were proposed in this pull request? Make `createLocalRelationProto` support relation with `TimestampType` ### Why are the changes needed? existing helper function `createLocalRelationProto` cannot create table with `TimestampType`: ``` org.apache.spark.SparkException: [INTERNAL_ERROR] Missing timezoneId where it is mandatory. SQLSTATE: XX000 at org.apache.spark.SparkException$.internalError(SparkException.scala:99) at org.apache.spark.SparkException$.internalError(SparkException.scala:103) at org.apache.spark.sql.util.ArrowUtils$.toArrowType(ArrowUtils.scala:57) at org.apache.spark.sql.util.ArrowUtils$.toArrowField(ArrowUtils.scala:139) at org.apache.spark.sql.util.ArrowUtils$.$anonfun$toArrowSchema$1(ArrowUtils.scala:181) at scala.collection.TraversableLike.$anonfun$map$1(TraversableLike.scala:286) at scala.collection.Iterator.foreach(Iterator.scala:943) at scala.collection.Iterator.foreach$(Iterator.scala:943) at scala.collection.AbstractIterator.foreach(Iterator.scala:1431) at scala.collection.IterableLike.foreach(IterableLike.scala:74) ``` ### Does this PR introduce _any_ user-facing change? No, test-only ### How was this patch tested? added ut ### Was this patch authored or co-authored using generative AI tooling? no Closes #47608 from zhengruifeng/create_timestamp_localrel. Authored-by: Ruifeng Zheng <ruifengz@apache.org> Signed-off-by: Hyukjin Kwon <gurwls223@apache.org>
1 parent de8ee94 commit 6d32472

2 files changed

Lines changed: 20 additions & 3 deletions

File tree

sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectPlannerSuite.scala

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,8 @@ trait SparkConnectPlanTest extends SharedSparkSession {
8787
*/
8888
def createLocalRelationProto(
8989
attrs: Seq[AttributeReference],
90-
data: Seq[InternalRow]): proto.Relation = {
90+
data: Seq[InternalRow],
91+
timeZoneId: String = "UTC"): proto.Relation = {
9192
val localRelationBuilder = proto.LocalRelation.newBuilder()
9293

9394
val bytes = ArrowConverters
@@ -96,7 +97,7 @@ trait SparkConnectPlanTest extends SharedSparkSession {
9697
DataTypeUtils.fromAttributes(attrs.map(_.toAttribute)),
9798
Long.MaxValue,
9899
Long.MaxValue,
99-
null,
100+
timeZoneId,
100101
true)
101102
.next()
102103

sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectProtoSuite.scala

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ import org.apache.spark.sql.connector.catalog.{Identifier, InMemoryTableCatalog,
4242
import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.CatalogHelper
4343
import org.apache.spark.sql.execution.arrow.ArrowConverters
4444
import org.apache.spark.sql.functions._
45-
import org.apache.spark.sql.types.{ArrayType, BooleanType, ByteType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, Metadata, ShortType, StringType, StructField, StructType}
45+
import org.apache.spark.sql.types._
4646
import org.apache.spark.unsafe.types.UTF8String
4747
import org.apache.spark.util.Utils
4848

@@ -64,6 +64,11 @@ class SparkConnectProtoSuite extends PlanTest with SparkConnectPlanTest {
6464
Seq(AttributeReference("id", IntegerType)(), AttributeReference("name", StringType)()),
6565
Seq.empty)
6666

67+
lazy val connectTestRelation3 =
68+
createLocalRelationProto(
69+
Seq(AttributeReference("id", IntegerType)(), AttributeReference("date", TimestampType)()),
70+
Seq.empty)
71+
6772
lazy val connectTestRelationMap =
6873
createLocalRelationProto(
6974
Seq(AttributeReference("id", MapType(StringType, StringType))()),
@@ -79,6 +84,11 @@ class SparkConnectProtoSuite extends PlanTest with SparkConnectPlanTest {
7984
new java.util.ArrayList[Row](),
8085
StructType(Seq(StructField("id", IntegerType), StructField("name", StringType))))
8186

87+
lazy val sparkTestRelation3: DataFrame =
88+
spark.createDataFrame(
89+
new java.util.ArrayList[Row](),
90+
StructType(Seq(StructField("id", IntegerType), StructField("date", TimestampType))))
91+
8292
lazy val sparkTestRelationMap: DataFrame =
8393
spark.createDataFrame(
8494
new java.util.ArrayList[Row](),
@@ -93,6 +103,12 @@ class SparkConnectProtoSuite extends PlanTest with SparkConnectPlanTest {
93103
comparePlans(connectPlan, sparkPlan)
94104
}
95105

106+
test("Basic select timestamp") {
107+
val connectPlan = connectTestRelation3.select("date".protoAttr)
108+
val sparkPlan = sparkTestRelation3.select("date")
109+
comparePlans(connectPlan, sparkPlan)
110+
}
111+
96112
test("Test select expression in strings") {
97113
val connectPlan = connectTestRelation.selectExpr("abs(id)", "name")
98114
val sparkPlan = sparkTestRelation.selectExpr("abs(id)", "name")

0 commit comments

Comments
 (0)