Skip to content

Commit 2cbe9fe

Browse files
plusplusjiajiaplusplusjiajia
committed
[flink] Preserve query authorization when narrowing splits in the Flink source
1 parent 2dcc4da commit 2cbe9fe

25 files changed

Lines changed: 1890 additions & 75 deletions

‎paimon-core/src/main/java/org/apache/paimon/table/ChainGroupReadTable.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@
3939
import org.apache.paimon.table.source.InnerTableRead;
4040
import org.apache.paimon.table.source.QueryAuthSplit;
4141
import org.apache.paimon.table.source.Split;
42+
import org.apache.paimon.table.source.Splits;
4243
import org.apache.paimon.table.source.TableRead;
4344
import org.apache.paimon.types.RowType;
4445
import org.apache.paimon.utils.ChainPartitionProjector;
@@ -610,7 +611,7 @@ public InnerTableRead withReadBatchSizer(ReadBatchSizer sizer) {
610611
@Override
611612
public RecordReader<InternalRow> createReader(Split split) throws IOException {
612613
// fallbackRead unwraps and applies the rules itself, so it gets the wrapper untouched.
613-
Split inner = QueryAuthSplit.unwrap(split);
614+
Split inner = Splits.underlying(split);
614615
if (inner instanceof ChainSplit || inner instanceof DataSplit) {
615616
return fallbackRead.createReader(split);
616617
}
@@ -640,7 +641,7 @@ List<DataSplit> unwrap(List<Split> splits) {
640641
"Branch scans of the same chain table returned different authorization rules.");
641642
authResult = next;
642643
}
643-
dataSplits.add((DataSplit) QueryAuthSplit.unwrap(split));
644+
dataSplits.add((DataSplit) Splits.underlying(split));
644645
}
645646
return dataSplits;
646647
}

‎paimon-core/src/main/java/org/apache/paimon/table/ChainTableFileStoreTable.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,8 @@
2929
import org.apache.paimon.table.source.ChainSplit;
3030
import org.apache.paimon.table.source.DataSplit;
3131
import org.apache.paimon.table.source.InnerTableRead;
32-
import org.apache.paimon.table.source.QueryAuthSplit;
3332
import org.apache.paimon.table.source.Split;
33+
import org.apache.paimon.table.source.Splits;
3434
import org.apache.paimon.table.source.StreamDataTableScan;
3535
import org.apache.paimon.table.source.TableRead;
3636
import org.apache.paimon.types.RowType;
@@ -219,7 +219,7 @@ public RecordReader<InternalRow> createReader(Split split) throws IOException {
219219
return fallbackRead.createReader(split);
220220
}
221221
// Route on what the wrapper carries, so the wrapper still reaches the read.
222-
Split inner = QueryAuthSplit.unwrap(split);
222+
Split inner = Splits.underlying(split);
223223
if (inner instanceof ChainSplit || inner instanceof DataSplit) {
224224
return chainGroupRead.createReader(split);
225225
}

‎paimon-core/src/main/java/org/apache/paimon/table/source/QueryAuthSplit.java‎

Lines changed: 0 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -54,16 +54,6 @@ public Split split() {
5454
return split;
5555
}
5656

57-
/**
58-
* The split this wrapper carries, or {@code split} itself when it is not one. Narrowing a
59-
* planned split to a concrete type must go through this, or the cast fails once the table
60-
* carries a rule. A caller handing the split to a read passes the wrapper on instead, or the
61-
* authorization never reaches the reader.
62-
*/
63-
public static Split unwrap(Split split) {
64-
return split instanceof QueryAuthSplit ? ((QueryAuthSplit) split).split() : split;
65-
}
66-
6757
@Nullable
6858
public TableQueryAuthResult authResult() {
6959
return authResult;

‎paimon-core/src/main/java/org/apache/paimon/table/source/Splits.java‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
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, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
19+
package org.apache.paimon.table.source;
20+
21+
import org.apache.paimon.catalog.TableQueryAuthResult;
22+
import org.apache.paimon.table.FallbackReadFileStoreTable.FallbackSplit;
23+
24+
import javax.annotation.Nullable;
25+
26+
/** Utilities for looking at what a planned {@link Split} is built from. */
27+
public class Splits {
28+
29+
private Splits() {}
30+
31+
/**
32+
* The split a planned one is built from, looked at through every wrapper. One layer is not
33+
* enough: an authorized fallback table hands out a {@link FallbackSplit} holding a {@link
34+
* QueryAuthSplit}.
35+
*
36+
* <p>Stops at a {@link DataSplit}. {@code FallbackDataSplit} is one and wraps itself, so that
37+
* condition ends the loop as much as it preserves what a plan carrying no rules resolved to.
38+
*
39+
* <p>Only for reading what the split is made of. A caller handing it to a read passes the
40+
* planned split on as it is, or the authorization never reaches the reader.
41+
*/
42+
public static Split underlying(Split split) {
43+
Split current = split;
44+
while (!(current instanceof DataSplit)) {
45+
if (current instanceof QueryAuthSplit) {
46+
current = ((QueryAuthSplit) current).split();
47+
} else if (current instanceof FallbackSplit) {
48+
current = ((FallbackSplit) current).wrapped();
49+
} else {
50+
return current;
51+
}
52+
}
53+
return current;
54+
}
55+
56+
/**
57+
* The authorization a planned split carries, looked at through every wrapper, or null when the
58+
* plan carries none. The companion of {@link #underlying(Split)}, which drops the wrapper the
59+
* rules live on.
60+
*
61+
* <p>A {@code FallbackDataSplit} wraps itself and ends the walk the way an unwrapped split
62+
* does: neither one carries rules.
63+
*/
64+
@Nullable
65+
public static TableQueryAuthResult authResult(Split split) {
66+
Split current = split;
67+
while (!(current instanceof QueryAuthSplit)) {
68+
if (!(current instanceof FallbackSplit)) {
69+
return null;
70+
}
71+
Split wrapped = ((FallbackSplit) current).wrapped();
72+
if (wrapped == current) {
73+
return null;
74+
}
75+
current = wrapped;
76+
}
77+
return ((QueryAuthSplit) current).authResult();
78+
}
79+
}

‎paimon-core/src/test/java/org/apache/paimon/table/source/SplitsTest.java‎

Lines changed: 167 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,167 @@
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, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
19+
package org.apache.paimon.table.source;
20+
21+
import org.apache.paimon.catalog.TableQueryAuthResult;
22+
import org.apache.paimon.data.BinaryRow;
23+
import org.apache.paimon.table.FallbackReadFileStoreTable;
24+
import org.apache.paimon.table.FallbackReadFileStoreTable.FallbackDataSplit;
25+
import org.apache.paimon.table.FallbackReadFileStoreTable.FallbackSplit;
26+
import org.apache.paimon.table.FallbackReadFileStoreTable.FallbackSplitImpl;
27+
28+
import org.junit.jupiter.api.Test;
29+
30+
import java.util.Collections;
31+
32+
import static org.apache.paimon.io.DataFileTestUtils.newFile;
33+
import static org.assertj.core.api.Assertions.assertThat;
34+
35+
/** Test for {@link Splits}. */
36+
public class SplitsTest {
37+
38+
@Test
39+
public void testDataSplitIsItsOwnUnderlyingSplit() {
40+
DataSplit split = dataSplit();
41+
42+
assertThat(Splits.underlying(split)).isSameAs(split);
43+
}
44+
45+
@Test
46+
public void testAuthorizationWrapperIsLookedThrough() {
47+
DataSplit split = dataSplit();
48+
49+
assertThat(Splits.underlying(withQueryAuth(split))).isSameAs(split);
50+
}
51+
52+
@Test
53+
public void testFallbackWrapperAroundAuthorizationWrapperIsLookedThrough() {
54+
DataSplit split = dataSplit();
55+
Split nested = toFallbackSplit(withQueryAuth(split));
56+
57+
assertThat(nested).isInstanceOf(FallbackSplitImpl.class);
58+
assertThat(((FallbackSplit) nested).wrapped()).isInstanceOf(QueryAuthSplit.class);
59+
assertThat(Splits.underlying(nested)).isSameAs(split);
60+
}
61+
62+
@Test
63+
public void testFallbackDataSplitIsItsOwnUnderlyingSplit() {
64+
DataSplit split = dataSplit();
65+
Split fallback = toFallbackSplit(split);
66+
67+
assertThat(fallback).isInstanceOf(FallbackDataSplit.class);
68+
assertThat(Splits.underlying(fallback)).isSameAs(fallback);
69+
}
70+
71+
@Test
72+
public void testSplitThatIsNeitherADataSplitNorAWrapperIsItsOwnUnderlyingSplit() {
73+
ChainSplit split = chainSplit();
74+
75+
assertThat(Splits.underlying(split)).isSameAs(split);
76+
}
77+
78+
@Test
79+
public void testFallbackWrapperAroundASplitThatIsNotADataSplitIsLookedThrough() {
80+
ChainSplit split = chainSplit();
81+
82+
assertThat(Splits.underlying(toFallbackSplit(split))).isSameAs(split);
83+
}
84+
85+
@Test
86+
public void testAuthorizationWrapperAroundASplitThatIsNotADataSplitIsLookedThrough() {
87+
ChainSplit split = chainSplit();
88+
89+
assertThat(Splits.underlying(withQueryAuth(split))).isSameAs(split);
90+
}
91+
92+
@Test
93+
public void testAuthorizationWrapperAroundAFallbackDataSplitStopsAtTheFallbackDataSplit() {
94+
Split fallback = toFallbackSplit(dataSplit());
95+
96+
assertThat(Splits.underlying(withQueryAuth(fallback))).isSameAs(fallback);
97+
}
98+
99+
@Test
100+
public void testSplitCarryingNoAuthorizationHasNoAuthResult() {
101+
assertThat(Splits.authResult(dataSplit())).isNull();
102+
assertThat(Splits.authResult(chainSplit())).isNull();
103+
assertThat(Splits.authResult(toFallbackSplit(dataSplit()))).isNull();
104+
assertThat(Splits.authResult(toFallbackSplit(chainSplit()))).isNull();
105+
}
106+
107+
@Test
108+
public void testAuthResultIsReadThroughTheAuthorizationWrapper() {
109+
TableQueryAuthResult authResult = authResult();
110+
111+
assertThat(Splits.authResult(new QueryAuthSplit(dataSplit(), authResult)))
112+
.isSameAs(authResult);
113+
}
114+
115+
@Test
116+
public void testAuthResultIsReadThroughAFallbackWrapper() {
117+
TableQueryAuthResult authResult = authResult();
118+
Split nested = toFallbackSplit(new QueryAuthSplit(dataSplit(), authResult));
119+
120+
assertThat(Splits.authResult(nested)).isSameAs(authResult);
121+
}
122+
123+
@Test
124+
public void testAuthResultIsReadFromAWrapperAroundAFallbackDataSplit() {
125+
TableQueryAuthResult authResult = authResult();
126+
Split split = new QueryAuthSplit(toFallbackSplit(dataSplit()), authResult);
127+
128+
assertThat(Splits.authResult(split)).isSameAs(authResult);
129+
}
130+
131+
@Test
132+
public void testAuthorizationWrapperWithoutRulesHasNoAuthResult() {
133+
assertThat(Splits.authResult(new QueryAuthSplit(dataSplit(), null))).isNull();
134+
}
135+
136+
private static TableQueryAuthResult authResult() {
137+
return new TableQueryAuthResult(
138+
null, Collections.singletonMap("name", "{\"name\":\"NULL\"}"));
139+
}
140+
141+
private static DataSplit dataSplit() {
142+
return DataSplit.builder()
143+
.withSnapshot(1)
144+
.withPartition(BinaryRow.EMPTY_ROW)
145+
.withBucket(0)
146+
.withBucketPath("bucket-0")
147+
.withDataFiles(Collections.singletonList(newFile(0L, 1L)))
148+
.build();
149+
}
150+
151+
private static ChainSplit chainSplit() {
152+
return new ChainSplit(
153+
BinaryRow.EMPTY_ROW,
154+
Collections.singletonList(newFile(0L, 1L)),
155+
Collections.emptyMap(),
156+
Collections.emptyMap(),
157+
null);
158+
}
159+
160+
private static Split withQueryAuth(Split split) {
161+
return new QueryAuthSplit(split, new TableQueryAuthResult(null, null));
162+
}
163+
164+
private static Split toFallbackSplit(Split split) {
165+
return FallbackReadFileStoreTable.toFallbackSplit(split, true);
166+
}
167+
}

‎paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/ContinuousFileSplitEnumerator.java‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,8 @@
3131
import org.apache.paimon.table.source.EndOfScanException;
3232
import org.apache.paimon.table.source.IncrementalSplit;
3333
import org.apache.paimon.table.source.SnapshotNotExistPlan;
34+
import org.apache.paimon.table.source.Split;
35+
import org.apache.paimon.table.source.Splits;
3436
import org.apache.paimon.table.source.StreamTableScan;
3537
import org.apache.paimon.table.source.TableScan;
3638

@@ -329,12 +331,14 @@ protected synchronized void assignSplits() {
329331

330332
protected int assignSuggestedTask(FileStoreSourceSplit split) {
331333
int task;
332-
if (split.split() instanceof DataSplit) {
333-
task = assignSuggestedTask((DataSplit) split.split());
334-
} else if (split.split() instanceof ChainSplit) {
335-
task = assignSuggestedTask((ChainSplit) split.split());
334+
// Reads the underlying split, assigns the original.
335+
Split inner = Splits.underlying(split.split());
336+
if (inner instanceof DataSplit) {
337+
task = assignSuggestedTask((DataSplit) inner);
338+
} else if (inner instanceof ChainSplit) {
339+
task = assignSuggestedTask((ChainSplit) inner);
336340
} else {
337-
task = assignSuggestedTask((IncrementalSplit) split.split());
341+
task = assignSuggestedTask((IncrementalSplit) inner);
338342
}
339343

340344
// Split assigners keep splits in a map keyed by task, but only ever hand out splits for

‎paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/DynamicPartitionFilteringInfo.java‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,20 @@
1818

1919
package org.apache.paimon.flink.source;
2020

21+
import org.apache.paimon.catalog.TableQueryAuthResult;
2122
import org.apache.paimon.codegen.CodeGenUtils;
2223
import org.apache.paimon.codegen.Projection;
24+
import org.apache.paimon.data.BinaryRow;
25+
import org.apache.paimon.flink.FlinkRowData;
26+
import org.apache.paimon.table.source.DataSplit;
27+
import org.apache.paimon.table.source.Split;
28+
import org.apache.paimon.table.source.Splits;
2329
import org.apache.paimon.types.RowType;
2430

31+
import org.apache.flink.table.connector.source.DynamicFilteringData;
32+
2533
import java.io.Serializable;
34+
import java.util.Collections;
2635
import java.util.List;
2736

2837
/** Manage dynamic partition filtering fields and table partition row. */
@@ -49,4 +58,25 @@ public Projection getPartitionRowProjection() {
4958

5059
return partitionRowProjection;
5160
}
61+
62+
/**
63+
* Whether {@code split} may still hold a row {@code dynamicFilteringData} asks for. Answered
64+
* from the partition the plan recorded, so pruning happens without opening the split.
65+
*
66+
* <p>A mask on a filtering field breaks that shortcut: the reader hands out a value other than
67+
* the recorded one, and comparing the recorded one drops splits the join does want. Such a
68+
* split is kept, and the join filters it on the masked value instead.
69+
*/
70+
public boolean mayMatch(DynamicFilteringData dynamicFilteringData, Split split) {
71+
TableQueryAuthResult authResult = Splits.authResult(split);
72+
if (authResult != null
73+
&& !Collections.disjoint(
74+
authResult.extractColumnMasking().keySet(),
75+
dynamicPartitionFilteringFields)) {
76+
return true;
77+
}
78+
BinaryRow partition = ((DataSplit) Splits.underlying(split)).partition();
79+
return dynamicFilteringData.contains(
80+
new FlinkRowData(getPartitionRowProjection().apply(partition)));
81+
}
5282
}

0 commit comments

Comments
 (0)