Skip to content

Commit 34cd6f9

Browse files
heyihongheyihong
committed
Remove flags
1 parent 78db6fd commit 34cd6f9

File tree

4 files changed

+95
-139
lines changed

4 files changed

+95
-139
lines changed

‎python/pyspark/sql/connect/expressions.py

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -29,10 +29,9 @@
2929
Optional,
3030
)
3131

32-
import os
33-
import json
34-
import decimal
3532
import datetime
33+
import decimal
34+
import json
3635
import warnings
3736
from threading import Lock
3837

@@ -85,9 +84,6 @@
8584
from pyspark.sql.connect.window import WindowSpec
8685
from pyspark.sql.connect.plan import LogicalPlan
8786

88-
_optional_data_types_for_map_and_array_literals_enabled = (
89-
os.getenv("CONNECT_OPTIONAL_DATATYPE_FOR_MAP_AND_ARRAY_LITERALS_ENABLED", "false") == "true"
90-
)
9187

9288
class Expression:
9389
"""
@@ -494,8 +490,10 @@ def to_plan(self, session: "SparkConnectClient") -> "proto.Expression":
494490
elif isinstance(self._dataType, DayTimeIntervalType):
495491
expr.literal.day_time_interval = int(self._value)
496492
elif isinstance(self._dataType, ArrayType):
497-
if not _optional_data_types_for_map_and_array_literals_enabled or len(self._value) == 0:
498-
expr.literal.array.element_type.CopyFrom(pyspark_types_to_proto_types(self._dataType.elementType))
493+
if len(self._value) == 0:
494+
expr.literal.array.element_type.CopyFrom(
495+
pyspark_types_to_proto_types(self._dataType.elementType)
496+
)
499497
for v in self._value:
500498
expr.literal.array.elements.append(
501499
LiteralExpression(v, self._dataType.elementType).to_plan(session).literal

‎python/pyspark/sql/tests/connect/test_connect_plan.py

Lines changed: 47 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@
1919
import datetime
2020
import decimal
2121
import math
22-
from contextlib import contextmanager
2322

2423
from pyspark.testing.connectutils import (
2524
PlanOnlyTestFixture,
@@ -1011,66 +1010,54 @@ def test_literal_expression_with_arrays(self):
10111010
self.assertEqual(len(l4.array.elements[1].array.elements), 2)
10121011
self.assertEqual(len(l4.array.elements[2].array.elements), 0)
10131012

1014-
@contextmanager
1015-
def _optional_data_types_enabled(self, enabled: bool):
1016-
from pyspark.sql.connect.expressions import _optional_data_types_for_map_and_array_literals_enabled
1017-
previous_value = _optional_data_types_for_map_and_array_literals_enabled
1018-
try:
1019-
_optional_data_types_for_map_and_array_literals_enabled = enabled
1020-
yield
1021-
finally:
1022-
_optional_data_types_for_map_and_array_literals_enabled = previous_value
1023-
10241013
def test_literal_to_any_conversion(self):
1025-
for optional_data_types_enabled in [True, False]:
1026-
with self._optional_data_types_enabled(optional_data_types_enabled):
1027-
for value in [
1028-
b"binary\0\0asas",
1029-
True,
1030-
False,
1031-
0,
1032-
12,
1033-
-1,
1034-
0.0,
1035-
1.234567,
1036-
decimal.Decimal(0.0),
1037-
decimal.Decimal(1.234567),
1038-
"sss",
1039-
datetime.date(2022, 12, 13),
1040-
datetime.datetime.now(),
1041-
datetime.timedelta(1, 2, 3),
1042-
[1, 2, 3, 4, 5, 6],
1043-
[-1.0, 2.0, 3.0],
1044-
["x", "y", "z"],
1045-
[[1.0, 2.0, 3.0], [4.0, 5.0], [6.0]],
1046-
]:
1047-
lit = LiteralExpression._from_value(value)
1048-
proto_lit = lit.to_plan(None).literal
1049-
value2 = LiteralExpression._to_value(proto_lit)
1050-
self.assertEqual(value, value2)
1051-
1052-
with self.assertRaises(AssertionError):
1053-
lit = LiteralExpression._from_value(1.234567)
1054-
proto_lit = lit.to_plan(None).literal
1055-
LiteralExpression._to_value(proto_lit, StringType())
1056-
1057-
with self.assertRaises(AssertionError):
1058-
lit = LiteralExpression._from_value("1.234567")
1059-
proto_lit = lit.to_plan(None).literal
1060-
LiteralExpression._to_value(proto_lit, DoubleType())
1061-
1062-
with self.assertRaises(AssertionError):
1063-
# build a array<string> proto literal, but with incorrect elements
1064-
proto_lit = proto.Expression().literal
1065-
proto_lit.array.element_type.CopyFrom(pyspark_types_to_proto_types(StringType()))
1066-
proto_lit.array.elements.append(
1067-
LiteralExpression("string", StringType()).to_plan(None).literal
1068-
)
1069-
proto_lit.array.elements.append(
1070-
LiteralExpression(1.234, DoubleType()).to_plan(None).literal
1071-
)
1072-
1073-
LiteralExpression._to_value(proto_lit, DoubleType)
1014+
for value in [
1015+
b"binary\0\0asas",
1016+
True,
1017+
False,
1018+
0,
1019+
12,
1020+
-1,
1021+
0.0,
1022+
1.234567,
1023+
decimal.Decimal(0.0),
1024+
decimal.Decimal(1.234567),
1025+
"sss",
1026+
datetime.date(2022, 12, 13),
1027+
datetime.datetime.now(),
1028+
datetime.timedelta(1, 2, 3),
1029+
[1, 2, 3, 4, 5, 6],
1030+
[-1.0, 2.0, 3.0],
1031+
["x", "y", "z"],
1032+
[[1.0, 2.0, 3.0], [4.0, 5.0], [6.0]],
1033+
]:
1034+
lit = LiteralExpression._from_value(value)
1035+
proto_lit = lit.to_plan(None).literal
1036+
value2 = LiteralExpression._to_value(proto_lit)
1037+
self.assertEqual(value, value2)
1038+
1039+
with self.assertRaises(AssertionError):
1040+
lit = LiteralExpression._from_value(1.234567)
1041+
proto_lit = lit.to_plan(None).literal
1042+
LiteralExpression._to_value(proto_lit, StringType())
1043+
1044+
with self.assertRaises(AssertionError):
1045+
lit = LiteralExpression._from_value("1.234567")
1046+
proto_lit = lit.to_plan(None).literal
1047+
LiteralExpression._to_value(proto_lit, DoubleType())
1048+
1049+
with self.assertRaises(AssertionError):
1050+
# build a array<string> proto literal, but with incorrect elements
1051+
proto_lit = proto.Expression().literal
1052+
proto_lit.array.element_type.CopyFrom(pyspark_types_to_proto_types(StringType()))
1053+
proto_lit.array.elements.append(
1054+
LiteralExpression("string", StringType()).to_plan(None).literal
1055+
)
1056+
proto_lit.array.elements.append(
1057+
LiteralExpression(1.234, DoubleType()).to_plan(None).literal
1058+
)
1059+
1060+
LiteralExpression._to_value(proto_lit, DoubleType)
10741061

10751062

10761063
if __name__ == "__main__":

‎sql/connect/common/src/main/scala/org/apache/spark/sql/connect/common/LiteralValueProtoConverter.scala

Lines changed: 4 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@ import java.lang.{Boolean => JBoolean, Byte => JByte, Character => JChar, Double
2121
import java.math.{BigDecimal => JBigDecimal}
2222
import java.sql.{Date, Timestamp}
2323
import java.time._
24-
import java.util.concurrent.atomic.AtomicBoolean
2524

2625
import scala.collection.{immutable, mutable}
2726
import scala.jdk.CollectionConverters._
@@ -41,22 +40,6 @@ import org.apache.spark.util.SparkClassUtils
4140

4241
object LiteralValueProtoConverter {
4342

44-
private var optionalDataTypeEnabled = new AtomicBoolean(
45-
sys.env.getOrElse(
46-
"CONNECT_OPTIONAL_DATATYPE_FOR_MAP_AND_ARRAY_LITERALS_ENABLED",
47-
"false") == "true")
48-
49-
object forTest {
50-
def withOptionalDataType(enabled: Boolean)(f: => Unit): Unit = {
51-
val previousValue = optionalDataTypeEnabled.getAndSet(enabled)
52-
try {
53-
f
54-
} finally {
55-
optionalDataTypeEnabled.set(previousValue)
56-
}
57-
}
58-
}
59-
6043
/**
6144
* Transforms literal value to the `proto.Expression.Literal.Builder`.
6245
*
@@ -81,7 +64,7 @@ object LiteralValueProtoConverter {
8164
def arrayBuilder(array: Array[_]) = {
8265
val ab = builder.getArrayBuilder
8366
array.foreach(x => ab.addElements(toLiteralProto(x)))
84-
if (!optionalDataTypeEnabled.get() || ab.getElementsCount == 0) {
67+
if (ab.getElementsCount == 0) {
8568
ab.setElementType(toConnectProtoType(toDataType(array.getClass.getComponentType)))
8669
}
8770
ab
@@ -146,7 +129,7 @@ object LiteralValueProtoConverter {
146129
throw new IllegalArgumentException(s"literal $other not supported (yet).")
147130
}
148131

149-
if (!optionalDataTypeEnabled.get() || ab.getElementsCount == 0) {
132+
if (ab.getElementsCount == 0) {
150133
ab.setElementType(toConnectProtoType(elementType))
151134
}
152135

@@ -168,13 +151,11 @@ object LiteralValueProtoConverter {
168151
throw new IllegalArgumentException(s"literal $other not supported (yet).")
169152
}
170153

171-
val enabled = optionalDataTypeEnabled.get()
172-
173-
if (!enabled || mb.getKeysCount == 0) {
154+
if (mb.getKeysCount == 0) {
174155
mb.setKeyType(toConnectProtoType(keyType))
175156
}
176157

177-
if (!enabled || mb.getValuesCount == 0) {
158+
if (mb.getValuesCount == 0) {
178159
mb.setValueType(toConnectProtoType(valueType))
179160
}
180161

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

Lines changed: 38 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -34,57 +34,47 @@ class LiteralExpressionProtoConverterSuite extends AnyFunSuite { // scalastyle:i
3434
}
3535
}
3636

37-
for (optionalDataTypeEnabled <- Seq(false, true)) {
38-
LiteralValueProtoConverter.forTest.withOptionalDataType(optionalDataTypeEnabled) {
39-
Seq(
40-
(Array(true, false, true), ArrayType(BooleanType)),
41-
(Array(1.toByte, 2.toByte, 3.toByte), ArrayType(ByteType)),
42-
(Array(1.toShort, 2.toShort, 3.toShort), ArrayType(ShortType)),
43-
(Array(1, 2, 3), ArrayType(IntegerType)),
44-
(Array(1L, 2L, 3L), ArrayType(LongType)),
45-
(Array(1.1d, 2.1d, 3.1d), ArrayType(DoubleType)),
46-
(Array(1.1f, 2.1f, 3.1f), ArrayType(FloatType)),
47-
(Array(Array[Int](), Array(1, 2, 3), Array(4, 5, 6)), ArrayType(ArrayType(IntegerType))),
48-
(Array(Array(1, 2, 3), Array(4, 5, 6), Array[Int]()), ArrayType(ArrayType(IntegerType))),
49-
(
50-
Array(Array(Array(Array(Array(Array(1, 2, 3)))))),
51-
ArrayType(ArrayType(ArrayType(ArrayType(ArrayType(ArrayType(IntegerType))))))),
52-
(Array(Map(1 -> 2)), ArrayType(MapType(IntegerType, IntegerType))),
53-
(Map[Int, Int](), MapType(IntegerType, IntegerType)),
54-
(Map(1 -> 2, 3 -> 4, 5 -> 6), MapType(IntegerType, IntegerType))).zipWithIndex.foreach {
55-
case ((v, t), idx) =>
56-
test(
57-
s"complex proto value and catalyst value conversion #$idx - " +
58-
s"optionalDataTypeEnabled = $optionalDataTypeEnabled") {
59-
assertResult(v)(
60-
LiteralValueProtoConverter.toCatalystValue(
61-
LiteralValueProtoConverter.toLiteralProto(v, t)))
62-
}
37+
Seq(
38+
(Array(true, false, true), ArrayType(BooleanType)),
39+
(Array(1.toByte, 2.toByte, 3.toByte), ArrayType(ByteType)),
40+
(Array(1.toShort, 2.toShort, 3.toShort), ArrayType(ShortType)),
41+
(Array(1, 2, 3), ArrayType(IntegerType)),
42+
(Array(1L, 2L, 3L), ArrayType(LongType)),
43+
(Array(1.1d, 2.1d, 3.1d), ArrayType(DoubleType)),
44+
(Array(1.1f, 2.1f, 3.1f), ArrayType(FloatType)),
45+
(Array(Array[Int](), Array(1, 2, 3), Array(4, 5, 6)), ArrayType(ArrayType(IntegerType))),
46+
(Array(Array(1, 2, 3), Array(4, 5, 6), Array[Int]()), ArrayType(ArrayType(IntegerType))),
47+
(
48+
Array(Array(Array(Array(Array(Array(1, 2, 3)))))),
49+
ArrayType(ArrayType(ArrayType(ArrayType(ArrayType(ArrayType(IntegerType))))))),
50+
(Array(Map(1 -> 2)), ArrayType(MapType(IntegerType, IntegerType))),
51+
(Map[Int, Int](), MapType(IntegerType, IntegerType)),
52+
(Map(1 -> 2, 3 -> 4, 5 -> 6), MapType(IntegerType, IntegerType))).zipWithIndex.foreach {
53+
case ((v, t), idx) =>
54+
test(s"complex proto value and catalyst value conversion #$idx") {
55+
assertResult(v)(
56+
LiteralValueProtoConverter.toCatalystValue(
57+
LiteralValueProtoConverter.toLiteralProto(v, t)))
6358
}
59+
}
6460

65-
test(
66-
"invalid array literal - empty array - " +
67-
s"optionalDataTypeEnabled = $optionalDataTypeEnabled") {
68-
val literalProto = proto.Expression.Literal
69-
.newBuilder()
70-
.setArray(proto.Expression.Literal.Array.newBuilder())
71-
.build()
72-
intercept[InvalidPlanInput] {
73-
LiteralValueProtoConverter.toCatalystValue(literalProto)
74-
}
75-
}
61+
test("invalid array literal - empty array") {
62+
val literalProto = proto.Expression.Literal
63+
.newBuilder()
64+
.setArray(proto.Expression.Literal.Array.newBuilder())
65+
.build()
66+
intercept[InvalidPlanInput] {
67+
LiteralValueProtoConverter.toCatalystValue(literalProto)
68+
}
69+
}
7670

77-
test(
78-
"invalid map literal - " +
79-
s"optionalDataTypeEnabled = $optionalDataTypeEnabled") {
80-
val literalProto = proto.Expression.Literal
81-
.newBuilder()
82-
.setMap(proto.Expression.Literal.Map.newBuilder())
83-
.build()
84-
intercept[InvalidPlanInput] {
85-
LiteralValueProtoConverter.toCatalystValue(literalProto)
86-
}
87-
}
71+
test("invalid map literal") {
72+
val literalProto = proto.Expression.Literal
73+
.newBuilder()
74+
.setMap(proto.Expression.Literal.Map.newBuilder())
75+
.build()
76+
intercept[InvalidPlanInput] {
77+
LiteralValueProtoConverter.toCatalystValue(literalProto)
8878
}
8979
}
9080
}

0 commit comments

Comments
 (0)