1+ // Copyright 2021-present StarRocks, Inc. All rights reserved.
2+ //
3+ // Licensed under the Apache License, Version 2.0 (the "License");
4+ // you may not use this file except in compliance with the License.
5+ // You may obtain a copy of the License at
6+ //
7+ // https://www.apache.org/licenses/LICENSE-2.0
8+ //
9+ // Unless required by applicable law or agreed to in writing, software
10+ // distributed under the License is distributed on an "AS IS" BASIS,
11+ // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+ // See the License for the specific language governing permissions and
13+ // limitations under the License.
14+
15+ // language: java
16+ package com .starrocks .planner ;
17+
18+ import com .google .common .collect .Maps ;
19+ import com .starrocks .catalog .IcebergTable ;
20+ import com .starrocks .connector .CatalogConnector ;
21+ import com .starrocks .connector .ConnectorMgr ;
22+ import com .starrocks .credential .CloudConfiguration ;
23+ import com .starrocks .credential .CloudConfigurationFactory ;
24+ import com .starrocks .qe .SessionVariable ;
25+ import com .starrocks .server .GlobalStateMgr ;
26+ import com .starrocks .thrift .TCompressionType ;
27+ import com .starrocks .thrift .TDataSink ;
28+ import mockit .Expectations ;
29+ import mockit .Mocked ;
30+ import org .apache .iceberg .Table ;
31+ import org .apache .iceberg .io .FileIO ;
32+ import org .junit .jupiter .api .Assertions ;
33+ import org .junit .jupiter .api .Test ;
34+
35+ import java .util .HashMap ;
36+ import java .util .Map ;
37+
38+ import static org .apache .iceberg .TableProperties .PARQUET_COMPRESSION ;
39+
40+ public class IcebergTableSinkTest {
41+
42+ @ Test
43+ public void testCompressionFromNativeTableProperty (@ Mocked GlobalStateMgr globalStateMgr , @ Mocked ConnectorMgr connectorMgr ,
44+ @ Mocked CatalogConnector connector , @ Mocked IcebergTable icebergTable ,
45+ @ Mocked Table nativeTable , @ Mocked FileIO fileIO ,
46+ @ Mocked SessionVariable sessionVariable ) {
47+ // nativeTable.properties contains parquet compression -> should use it
48+ Map <String , String > nativeProps = Maps .newHashMap ();
49+ nativeProps .put (PARQUET_COMPRESSION , "zstd" );
50+
51+ CloudConfiguration cc = CloudConfigurationFactory .buildCloudConfigurationForStorage (new HashMap <>());
52+
53+ new Expectations () {
54+ {
55+ // iceberg table / native table basic
56+ icebergTable .getNativeTable ();
57+ result = nativeTable ;
58+
59+ icebergTable .getCatalogName ();
60+ result = "catA" ;
61+
62+ icebergTable .getId ();
63+ result = 101L ;
64+
65+ icebergTable .getUUID ();
66+ result = "uuidA" ;
67+
68+ nativeTable .location ();
69+ result = "s3://bucket/a" ;
70+
71+ nativeTable .properties ();
72+ result = nativeProps ;
73+
74+ nativeTable .io ();
75+ result = fileIO ;
76+
77+ fileIO .properties ();
78+ result = new HashMap <String , String >();
79+
80+ // Global state and connector metadata fallback path
81+ GlobalStateMgr .getCurrentState ();
82+ result = globalStateMgr ;
83+
84+ globalStateMgr .getConnectorMgr ();
85+ result = connectorMgr ;
86+
87+ connectorMgr .getConnector ("catA" );
88+ result = connector ;
89+
90+ connector .getMetadata ().getCloudConfiguration ();
91+ result = cc ;
92+
93+ // session variable should not be used for compression in this case, but stub anyway
94+ sessionVariable .getConnectorSinkCompressionCodec ();
95+ result = "gzip" ;
96+
97+ sessionVariable .getConnectorSinkTargetMaxFileSize ();
98+ result = 0L ;
99+ }
100+ };
101+
102+ TupleDescriptor desc = new TupleDescriptor (new TupleId (0 ));
103+ IcebergTableSink sink = new IcebergTableSink (icebergTable , desc , false , sessionVariable , null );
104+ TDataSink t = sink .toThrift ();
105+ // compression_type should map "zstd" -> TCompressionType.ZSTD
106+ Assertions .assertEquals (TCompressionType .ZSTD , t .getIceberg_table_sink ().getCompression_type ());
107+ }
108+
109+ @ Test
110+ public void testCompressionFromSessionVariableFallback (@ Mocked GlobalStateMgr globalStateMgr ,
111+ @ Mocked ConnectorMgr connectorMgr , @ Mocked CatalogConnector connector ,
112+ @ Mocked IcebergTable icebergTable , @ Mocked Table nativeTable ,
113+ @ Mocked FileIO fileIO , @ Mocked SessionVariable sessionVariable ) {
114+ // nativeTable.properties does NOT contain parquet compression -> should fallback to sessionVariable
115+ Map <String , String > nativeProps = Maps .newHashMap (); // empty
116+
117+ CloudConfiguration cc = CloudConfigurationFactory .buildCloudConfigurationForStorage (new HashMap <>());
118+
119+ new Expectations () {
120+ {
121+ icebergTable .getNativeTable ();
122+ result = nativeTable ;
123+
124+ icebergTable .getCatalogName ();
125+ result = "catB" ;
126+
127+ icebergTable .getId ();
128+ result = 102L ;
129+
130+ icebergTable .getUUID ();
131+ result = "uuidB" ;
132+
133+ nativeTable .location ();
134+ result = "s3://bucket/b" ;
135+
136+ nativeTable .properties ();
137+ result = nativeProps ;
138+
139+ nativeTable .io ();
140+ result = fileIO ;
141+
142+ fileIO .properties ();
143+ result = new HashMap <String , String >();
144+
145+ GlobalStateMgr .getCurrentState ();
146+ result = globalStateMgr ;
147+
148+ globalStateMgr .getConnectorMgr ();
149+ result = connectorMgr ;
150+
151+ connectorMgr .getConnector ("catB" );
152+ result = connector ;
153+
154+ connector .getMetadata ().getCloudConfiguration ();
155+ result = cc ;
156+
157+ // session variable provides fallback compression codec
158+ sessionVariable .getConnectorSinkCompressionCodec ();
159+ result = "gzip" ;
160+ sessionVariable .getConnectorSinkTargetMaxFileSize ();
161+ result = 0L ;
162+ }
163+ };
164+
165+ TupleDescriptor desc = new TupleDescriptor (new TupleId (0 ));
166+ IcebergTableSink sink = new IcebergTableSink (icebergTable , desc , false , sessionVariable , null );
167+ TDataSink t = sink .toThrift ();
168+ // fallback "gzip" -> TCompressionType.GZIP
169+ Assertions .assertEquals (TCompressionType .GZIP , t .getIceberg_table_sink ().getCompression_type ());
170+ }
171+ }
0 commit comments