Skip to content

Commit b16a607

Browse files
committed
case ByteType =>
case ShortType => -> case ByteType | ShortType =>
1 parent e9d5625 commit b16a607

File tree

2 files changed

+14
-38
lines changed

2 files changed

+14
-38
lines changed

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilters.scala

Lines changed: 12 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -42,14 +42,10 @@ private[parquet] class ParquetFilters(pushDownDate: Boolean, pushDownStartWith:
4242
private val makeEq: PartialFunction[DataType, (String, Any) => FilterPredicate] = {
4343
case BooleanType =>
4444
(n: String, v: Any) => FilterApi.eq(booleanColumn(n), v.asInstanceOf[java.lang.Boolean])
45-
case ByteType =>
45+
case ByteType | ShortType =>
4646
(n: String, v: Any) => FilterApi.eq(
4747
intColumn(n),
48-
Option(v).map(b => b.asInstanceOf[java.lang.Byte].toInt.asInstanceOf[Integer]).orNull)
49-
case ShortType =>
50-
(n: String, v: Any) => FilterApi.eq(
51-
intColumn(n),
52-
Option(v).map(b => b.asInstanceOf[java.lang.Short].toInt.asInstanceOf[Integer]).orNull)
48+
Option(v).map(_.asInstanceOf[Number].intValue.asInstanceOf[Integer]).orNull)
5349
case IntegerType =>
5450
(n: String, v: Any) => FilterApi.eq(intColumn(n), v.asInstanceOf[Integer])
5551
case LongType =>
@@ -77,14 +73,10 @@ private[parquet] class ParquetFilters(pushDownDate: Boolean, pushDownStartWith:
7773
private val makeNotEq: PartialFunction[DataType, (String, Any) => FilterPredicate] = {
7874
case BooleanType =>
7975
(n: String, v: Any) => FilterApi.notEq(booleanColumn(n), v.asInstanceOf[java.lang.Boolean])
80-
case ByteType =>
81-
(n: String, v: Any) => FilterApi.notEq(
82-
intColumn(n),
83-
Option(v).map(b => b.asInstanceOf[java.lang.Byte].toInt.asInstanceOf[Integer]).orNull)
84-
case ShortType =>
76+
case ByteType | ShortType =>
8577
(n: String, v: Any) => FilterApi.notEq(
8678
intColumn(n),
87-
Option(v).map(b => b.asInstanceOf[java.lang.Short].toInt.asInstanceOf[Integer]).orNull)
79+
Option(v).map(_.asInstanceOf[Number].intValue.asInstanceOf[Integer]).orNull)
8880
case IntegerType =>
8981
(n: String, v: Any) => FilterApi.notEq(intColumn(n), v.asInstanceOf[Integer])
9082
case LongType =>
@@ -109,14 +101,10 @@ private[parquet] class ParquetFilters(pushDownDate: Boolean, pushDownStartWith:
109101
}
110102

111103
private val makeLt: PartialFunction[DataType, (String, Any) => FilterPredicate] = {
112-
case ByteType =>
113-
(n: String, v: Any) => FilterApi.lt(
114-
intColumn(n),
115-
Option(v).map(b => b.asInstanceOf[java.lang.Byte].toInt.asInstanceOf[Integer]).orNull)
116-
case ShortType =>
104+
case ByteType | ShortType =>
117105
(n: String, v: Any) => FilterApi.lt(
118106
intColumn(n),
119-
Option(v).map(b => b.asInstanceOf[java.lang.Short].toInt.asInstanceOf[Integer]).orNull)
107+
v.asInstanceOf[Number].intValue.asInstanceOf[Integer])
120108
case IntegerType =>
121109
(n: String, v: Any) => FilterApi.lt(intColumn(n), v.asInstanceOf[Integer])
122110
case LongType =>
@@ -140,14 +128,10 @@ private[parquet] class ParquetFilters(pushDownDate: Boolean, pushDownStartWith:
140128
}
141129

142130
private val makeLtEq: PartialFunction[DataType, (String, Any) => FilterPredicate] = {
143-
case ByteType =>
131+
case ByteType | ShortType =>
144132
(n: String, v: Any) => FilterApi.ltEq(
145133
intColumn(n),
146-
Option(v).map(b => b.asInstanceOf[java.lang.Byte].toInt.asInstanceOf[Integer]).orNull)
147-
case ShortType =>
148-
(n: String, v: Any) => FilterApi.ltEq(
149-
intColumn(n),
150-
Option(v).map(b => b.asInstanceOf[java.lang.Short].toInt.asInstanceOf[Integer]).orNull)
134+
v.asInstanceOf[Number].intValue.asInstanceOf[Integer])
151135
case IntegerType =>
152136
(n: String, v: Any) => FilterApi.ltEq(intColumn(n), v.asInstanceOf[java.lang.Integer])
153137
case LongType =>
@@ -171,14 +155,10 @@ private[parquet] class ParquetFilters(pushDownDate: Boolean, pushDownStartWith:
171155
}
172156

173157
private val makeGt: PartialFunction[DataType, (String, Any) => FilterPredicate] = {
174-
case ByteType =>
175-
(n: String, v: Any) => FilterApi.gt(
176-
intColumn(n),
177-
Option(v).map(b => b.asInstanceOf[java.lang.Byte].toInt.asInstanceOf[Integer]).orNull)
178-
case ShortType =>
158+
case ByteType | ShortType =>
179159
(n: String, v: Any) => FilterApi.gt(
180160
intColumn(n),
181-
Option(v).map(b => b.asInstanceOf[java.lang.Short].toInt.asInstanceOf[Integer]).orNull)
161+
v.asInstanceOf[Number].intValue.asInstanceOf[Integer])
182162
case IntegerType =>
183163
(n: String, v: Any) => FilterApi.gt(intColumn(n), v.asInstanceOf[java.lang.Integer])
184164
case LongType =>
@@ -202,14 +182,10 @@ private[parquet] class ParquetFilters(pushDownDate: Boolean, pushDownStartWith:
202182
}
203183

204184
private val makeGtEq: PartialFunction[DataType, (String, Any) => FilterPredicate] = {
205-
case ByteType =>
206-
(n: String, v: Any) => FilterApi.gtEq(
207-
intColumn(n),
208-
Option(v).map(b => b.asInstanceOf[java.lang.Byte].toInt.asInstanceOf[Integer]).orNull)
209-
case ShortType =>
185+
case ByteType | ShortType =>
210186
(n: String, v: Any) => FilterApi.gtEq(
211187
intColumn(n),
212-
Option(v).map(b => b.asInstanceOf[java.lang.Short].toInt.asInstanceOf[Integer]).orNull)
188+
v.asInstanceOf[Number].intValue.asInstanceOf[Integer])
213189
case IntegerType =>
214190
(n: String, v: Any) => FilterApi.gtEq(intColumn(n), v.asInstanceOf[java.lang.Integer])
215191
case LongType =>

sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -178,7 +178,7 @@ class ParquetFilterSuite extends QueryTest with ParquetTest with SharedSQLContex
178178
}
179179
}
180180

181-
test(s"filter pushdown - ByteType") {
181+
test("filter pushdown - tinyint") {
182182
withParquetDataFrame((1 to 4).map(i => Tuple1(Option(i.toByte)))) { implicit df =>
183183
assert(df.schema.head.dataType === ByteType)
184184
checkFilterPredicate('_1.isNull, classOf[Eq[_]], Seq.empty[Row])
@@ -206,7 +206,7 @@ class ParquetFilterSuite extends QueryTest with ParquetTest with SharedSQLContex
206206
}
207207
}
208208

209-
test(s"filter pushdown - ShortType") {
209+
test("filter pushdown - smallint") {
210210
withParquetDataFrame((1 to 4).map(i => Tuple1(Option(i.toShort)))) { implicit df =>
211211
assert(df.schema.head.dataType === ShortType)
212212
checkFilterPredicate('_1.isNull, classOf[Eq[_]], Seq.empty[Row])

0 commit comments

Comments
 (0)