Skip to content

Commit a676552

Browse files
kevinjqliukevinjqliu
authored
Spark 3.4: Support recursive delegate unwrapping to find ExtendedParser in parser chains (#16306)
Backport of #14483 (and follow-up #14497) to spark/v3.4.
1 parent 1cea23e commit a676552

2 files changed

Lines changed: 279 additions & 3 deletions

File tree

‎spark/v3.4/spark-extensions/src/test/java/org/apache/iceberg/spark/TestExtendedParser.java‎

Lines changed: 236 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,236 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
package org.apache.iceberg.spark;
20+
21+
import static org.assertj.core.api.Assertions.assertThat;
22+
import static org.assertj.core.api.Assertions.assertThatThrownBy;
23+
import static org.mockito.Mockito.mock;
24+
import static org.mockito.Mockito.verify;
25+
import static org.mockito.Mockito.when;
26+
27+
import java.lang.reflect.Field;
28+
import java.util.Collections;
29+
import java.util.List;
30+
import org.apache.iceberg.NullOrder;
31+
import org.apache.iceberg.SortDirection;
32+
import org.apache.iceberg.expressions.Term;
33+
import org.apache.spark.sql.SparkSession;
34+
import org.apache.spark.sql.catalyst.parser.AbstractSqlParser;
35+
import org.apache.spark.sql.catalyst.parser.AstBuilder;
36+
import org.apache.spark.sql.catalyst.parser.ParserInterface;
37+
import org.apache.spark.sql.catalyst.parser.extensions.IcebergSparkSqlExtensionsParser;
38+
import org.junit.jupiter.api.AfterAll;
39+
import org.junit.jupiter.api.AfterEach;
40+
import org.junit.jupiter.api.BeforeAll;
41+
import org.junit.jupiter.api.BeforeEach;
42+
import org.junit.jupiter.api.Test;
43+
44+
public class TestExtendedParser {
45+
46+
private static SparkSession spark;
47+
private static final String SQL_PARSER_FIELD = "sqlParser";
48+
private ParserInterface originalParser;
49+
50+
@BeforeAll
51+
public static void before() {
52+
spark =
53+
SparkSession.builder()
54+
.master("local")
55+
.appName("TestExtendedParser")
56+
.config(TestBase.DISABLE_UI)
57+
.getOrCreate();
58+
}
59+
60+
@AfterAll
61+
public static void after() {
62+
if (spark != null) {
63+
spark.stop();
64+
}
65+
}
66+
67+
@BeforeEach
68+
public void saveOriginalParser() throws Exception {
69+
Class<?> clazz = spark.sessionState().getClass();
70+
Field parserField = null;
71+
while (clazz != null && parserField == null) {
72+
try {
73+
parserField = clazz.getDeclaredField(SQL_PARSER_FIELD);
74+
} catch (NoSuchFieldException e) {
75+
clazz = clazz.getSuperclass();
76+
}
77+
}
78+
parserField.setAccessible(true);
79+
originalParser = (ParserInterface) parserField.get(spark.sessionState());
80+
}
81+
82+
@AfterEach
83+
public void restoreOriginalParser() throws Exception {
84+
setSessionStateParser(spark.sessionState(), originalParser);
85+
}
86+
87+
/**
88+
* Tests that the Iceberg extended SQL parser can correctly parse a sort order string and return
89+
* the expected RawOrderField.
90+
*
91+
* @throws Exception if reflection access fails
92+
*/
93+
@Test
94+
public void testParseSortOrderWithRealIcebergExtendedParser() throws Exception {
95+
ParserInterface origParser = null;
96+
Class<?> clazz = spark.sessionState().getClass();
97+
while (clazz != null && origParser == null) {
98+
try {
99+
Field parserField = clazz.getDeclaredField(SQL_PARSER_FIELD);
100+
parserField.setAccessible(true);
101+
origParser = (ParserInterface) parserField.get(spark.sessionState());
102+
} catch (NoSuchFieldException e) {
103+
clazz = clazz.getSuperclass();
104+
}
105+
}
106+
assertThat(origParser).isNotNull();
107+
108+
IcebergSparkSqlExtensionsParser icebergParser = new IcebergSparkSqlExtensionsParser(origParser);
109+
110+
setSessionStateParser(spark.sessionState(), icebergParser);
111+
112+
List<ExtendedParser.RawOrderField> fields =
113+
ExtendedParser.parseSortOrder(spark, "id ASC NULLS FIRST");
114+
115+
assertThat(fields).isNotEmpty();
116+
ExtendedParser.RawOrderField first = fields.get(0);
117+
assertThat(first.direction()).isEqualTo(SortDirection.ASC);
118+
assertThat(first.nullOrder()).isEqualTo(NullOrder.NULLS_FIRST);
119+
}
120+
121+
/**
122+
* Tests that parseSortOrder can find and use an ExtendedParser that is wrapped inside another
123+
* ParserInterface implementation.
124+
*
125+
* @throws Exception if reflection access fails
126+
*/
127+
@Test
128+
public void testParseSortOrderFindsNestedExtendedParser() throws Exception {
129+
ExtendedParser icebergParser = mock(ExtendedParser.class);
130+
131+
ExtendedParser.RawOrderField field =
132+
new ExtendedParser.RawOrderField(
133+
mock(Term.class), SortDirection.ASC, NullOrder.NULLS_FIRST);
134+
List<ExtendedParser.RawOrderField> expected = Collections.singletonList(field);
135+
136+
when(icebergParser.parseSortOrder("id ASC NULLS FIRST")).thenReturn(expected);
137+
138+
ParserInterface wrapper = new WrapperParser(icebergParser);
139+
140+
setSessionStateParser(spark.sessionState(), wrapper);
141+
142+
List<ExtendedParser.RawOrderField> result =
143+
ExtendedParser.parseSortOrder(spark, "id ASC NULLS FIRST");
144+
assertThat(result).isSameAs(expected);
145+
146+
verify(icebergParser).parseSortOrder("id ASC NULLS FIRST");
147+
}
148+
149+
/**
150+
* Tests that parseSortOrder throws an exception if no ExtendedParser instance can be found in the
151+
* parser chain.
152+
*
153+
* @throws Exception if reflection access fails
154+
*/
155+
@Test
156+
public void testParseSortOrderThrowsWhenNoExtendedParserFound() throws Exception {
157+
ParserInterface dummy = mock(ParserInterface.class);
158+
setSessionStateParser(spark.sessionState(), dummy);
159+
160+
assertThatThrownBy(() -> ExtendedParser.parseSortOrder(spark, "id ASC"))
161+
.isInstanceOf(IllegalStateException.class)
162+
.hasMessageContaining("Iceberg ExtendedParser");
163+
}
164+
165+
/**
166+
* Tests that parseSortOrder can find an ExtendedParser in a parent class field of the parser.
167+
*
168+
* @throws Exception if reflection access fails
169+
*/
170+
@Test
171+
public void testParseSortOrderFindsExtendedParserInParentClassField() throws Exception {
172+
ExtendedParser icebergParser = mock(ExtendedParser.class);
173+
ExtendedParser.RawOrderField field =
174+
new ExtendedParser.RawOrderField(
175+
mock(Term.class), SortDirection.ASC, NullOrder.NULLS_FIRST);
176+
List<ExtendedParser.RawOrderField> expected = Collections.singletonList(field);
177+
when(icebergParser.parseSortOrder("id ASC NULLS FIRST")).thenReturn(expected);
178+
ParserInterface parser = new GrandChildParser(icebergParser);
179+
setSessionStateParser(spark.sessionState(), parser);
180+
181+
List<ExtendedParser.RawOrderField> result =
182+
ExtendedParser.parseSortOrder(spark, "id ASC NULLS FIRST");
183+
assertThat(result).isSameAs(expected);
184+
verify(icebergParser).parseSortOrder("id ASC NULLS FIRST");
185+
}
186+
187+
private static void setSessionStateParser(Object sessionState, ParserInterface parser)
188+
throws Exception {
189+
Class<?> clazz = sessionState.getClass();
190+
Field targetField = null;
191+
while (clazz != null && targetField == null) {
192+
try {
193+
targetField = clazz.getDeclaredField(SQL_PARSER_FIELD);
194+
} catch (NoSuchFieldException e) {
195+
clazz = clazz.getSuperclass();
196+
}
197+
}
198+
if (targetField == null) {
199+
throw new IllegalStateException(
200+
"No suitable sqlParser field found in sessionState class hierarchy!");
201+
}
202+
targetField.setAccessible(true);
203+
targetField.set(sessionState, parser);
204+
}
205+
206+
private static class WrapperParser extends AbstractSqlParser {
207+
private final ParserInterface delegate;
208+
private String name;
209+
210+
WrapperParser(ParserInterface delegate) {
211+
this.delegate = delegate;
212+
this.name = "delegate";
213+
}
214+
215+
public ParserInterface getDelegate() {
216+
return delegate;
217+
}
218+
219+
@Override
220+
public AstBuilder astBuilder() {
221+
return null;
222+
}
223+
}
224+
225+
private static class ChildParser extends WrapperParser {
226+
ChildParser(ParserInterface parent) {
227+
super(parent);
228+
}
229+
}
230+
231+
private static class GrandChildParser extends ChildParser {
232+
GrandChildParser(ParserInterface parent) {
233+
super(parent);
234+
}
235+
}
236+
}

‎spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/ExtendedParser.java‎

Lines changed: 43 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,13 +18,16 @@
1818
*/
1919
package org.apache.iceberg.spark;
2020

21+
import java.lang.reflect.Field;
2122
import java.util.List;
2223
import org.apache.iceberg.NullOrder;
2324
import org.apache.iceberg.SortDirection;
2425
import org.apache.iceberg.expressions.Term;
2526
import org.apache.spark.sql.AnalysisException;
2627
import org.apache.spark.sql.SparkSession;
2728
import org.apache.spark.sql.catalyst.parser.ParserInterface;
29+
import org.slf4j.Logger;
30+
import org.slf4j.LoggerFactory;
2831

2932
public interface ExtendedParser extends ParserInterface {
3033
class RawOrderField {
@@ -52,10 +55,10 @@ public NullOrder nullOrder() {
5255
}
5356

5457
static List<RawOrderField> parseSortOrder(SparkSession spark, String orderString) {
55-
if (spark.sessionState().sqlParser() instanceof ExtendedParser) {
56-
ExtendedParser parser = (ExtendedParser) spark.sessionState().sqlParser();
58+
ExtendedParser extParser = findParser(spark.sessionState().sqlParser(), ExtendedParser.class);
59+
if (extParser != null) {
5760
try {
58-
return parser.parseSortOrder(orderString);
61+
return extParser.parseSortOrder(orderString);
5962
} catch (AnalysisException e) {
6063
throw new IllegalArgumentException(
6164
String.format("Unable to parse sortOrder: %s", orderString), e);
@@ -66,5 +69,42 @@ static List<RawOrderField> parseSortOrder(SparkSession spark, String orderString
6669
}
6770
}
6871

72+
private static <T> T findParser(ParserInterface parser, Class<T> clazz) {
73+
ParserInterface current = parser;
74+
while (current != null) {
75+
if (clazz.isInstance(current)) {
76+
return clazz.cast(current);
77+
}
78+
79+
current = getNextDelegateParser(current);
80+
}
81+
82+
return null;
83+
}
84+
85+
private static ParserInterface getNextDelegateParser(ParserInterface parser) {
86+
try {
87+
Class<?> clazz = parser.getClass();
88+
while (clazz != null) {
89+
for (Field field : clazz.getDeclaredFields()) {
90+
field.setAccessible(true);
91+
Object value = field.get(parser);
92+
if (value instanceof ParserInterface && value != parser) {
93+
return (ParserInterface) value;
94+
}
95+
}
96+
clazz = clazz.getSuperclass();
97+
}
98+
} catch (Exception e) {
99+
log().warn("Failed to scan delegate parser in {}: ", parser.getClass().getName(), e);
100+
}
101+
102+
return null;
103+
}
104+
105+
private static Logger log() {
106+
return LoggerFactory.getLogger(ExtendedParser.class);
107+
}
108+
69109
List<RawOrderField> parseSortOrder(String orderString) throws AnalysisException;
70110
}

0 commit comments

Comments
 (0)