Skip to content

Commit 58b850e

Browse files
madhav-dbmadhav-db
authored
Adds fallback for complex types [ES-1469237] (#834)
* adds fallback for complex types * Add null check for requiredType in getArrowMetadataWithComplexTypeFallback method * Refactor ArrowStreamResult to simplify complex type handling and remove unused constants * undo * Enhance ArrowStreamResult with complex type handling and add unit tests for validation * Add complex type handling in DatabricksResultSet and improve exception messages for disabled support - Introduced a private method `isComplexType` to streamline complex type checks. - Updated exception messages in `getArray`, `getMap`, and `getStruct` methods to provide clearer guidance on enabling complex datatype support. - Added unit tests to verify exception handling when complex datatype support is disabled. * Refactor exception handling in DatabricksResultSet to use DatabricksSQLException - Updated methods getArray, getStruct, and getMap to throw DatabricksSQLException with specific error codes when complex datatype support is disabled. - Enhanced logging for error conditions to improve debugging and user guidance. - Adjusted unit tests to verify the new exception types and error codes. * - Added a logger to ArrowStreamResult to provide debug information when complex datatype support is disabled. - Improved clarity in handling complex types by logging the conversion to STRING when necessary.
1 parent d0a6274 commit 58b850e

File tree

4 files changed

+125
-10
lines changed

4 files changed

+125
-10
lines changed

‎src/main/java/com/databricks/jdbc/api/impl/DatabricksResultSet.java‎

Lines changed: 38 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -472,6 +472,16 @@ public ResultSetMetaData getMetaData() throws SQLException {
472472
return resultSetMetaData;
473473
}
474474

475+
/**
476+
* Checks if the given type name represents a complex type (ARRAY, MAP, or STRUCT).
477+
*
478+
* @param typeName The type name to check
479+
* @return true if the type name starts with ARRAY, MAP, or STRUCT, false otherwise
480+
*/
481+
private static boolean isComplexType(String typeName) {
482+
return typeName.startsWith(ARRAY) || typeName.startsWith(MAP) || typeName.startsWith(STRUCT);
483+
}
484+
475485
@Override
476486
public Object getObject(int columnIndex) throws SQLException {
477487
checkIfClosed();
@@ -482,9 +492,7 @@ public Object getObject(int columnIndex) throws SQLException {
482492
int columnType = resultSetMetaData.getColumnType(columnIndex);
483493
String columnTypeName = resultSetMetaData.getColumnTypeName(columnIndex);
484494
// separate handling for complex data types
485-
if (columnTypeName.startsWith(ARRAY)
486-
|| columnTypeName.startsWith(MAP)
487-
|| columnTypeName.startsWith(STRUCT)) {
495+
if (isComplexType(columnTypeName)) {
488496
return handleComplexDataTypes(obj, columnTypeName);
489497
}
490498
// TODO: Add separate handling for INTERVAL JSON_ARRAY result format.
@@ -1151,11 +1159,18 @@ public Clob getClob(int columnIndex) throws SQLException {
11511159
public Array getArray(int columnIndex) throws SQLException {
11521160
LOGGER.debug("Getting Array from column index: {}", columnIndex);
11531161
if (!complexDatatypeSupport) {
1154-
throw new SQLException("Complex datatype support support is disabled");
1162+
LOGGER.error(
1163+
"Complex datatype support support is disabled. Use connection parameter `EnableComplexDatatypeSupport=1` to enable it.");
1164+
throw new DatabricksSQLException(
1165+
"Complex datatype support support is disabled. Use connection parameter `EnableComplexDatatypeSupport=1` to enable it.",
1166+
DatabricksDriverErrorCode.COMPLEX_DATA_TYPE_ARRAY_CONVERSION_ERROR);
11551167
}
11561168
if (this.resultSetType.equals(ResultSetType.THRIFT_INLINE)
11571169
|| this.resultSetType.equals(ResultSetType.SEA_INLINE)) {
1158-
throw new SQLException("Complex data types are not supported in inline mode");
1170+
LOGGER.error("Complex data types are not supported in inline mode");
1171+
throw new DatabricksSQLException(
1172+
"Complex data types are not supported in inline mode",
1173+
DatabricksDriverErrorCode.COMPLEX_DATA_TYPE_ARRAY_CONVERSION_ERROR);
11591174
}
11601175
checkIfClosed();
11611176
Object obj = getObjectInternal(columnIndex);
@@ -1174,11 +1189,18 @@ public Array getArray(int columnIndex) throws SQLException {
11741189
public Struct getStruct(int columnIndex) throws SQLException {
11751190
LOGGER.debug("Getting Struct from column index: {}", columnIndex);
11761191
if (!complexDatatypeSupport) {
1177-
throw new SQLException("Complex datatype support support is disabled");
1192+
LOGGER.error(
1193+
"Complex datatype support support is disabled. Use connection parameter `EnableComplexDatatypeSupport=1` to enable it.");
1194+
throw new DatabricksSQLException(
1195+
"Complex datatype support support is disabled. Use connection parameter `EnableComplexDatatypeSupport=1` to enable it.",
1196+
DatabricksDriverErrorCode.COMPLEX_DATA_TYPE_STRUCT_CONVERSION_ERROR);
11781197
}
11791198
if (this.resultSetType.equals(ResultSetType.THRIFT_INLINE)
11801199
|| this.resultSetType.equals(ResultSetType.SEA_INLINE)) {
1181-
throw new SQLException("Complex data types are not supported in inline mode");
1200+
LOGGER.error("Complex data types are not supported in inline mode");
1201+
throw new DatabricksSQLException(
1202+
"Complex data types are not supported in inline mode",
1203+
DatabricksDriverErrorCode.COMPLEX_DATA_TYPE_STRUCT_CONVERSION_ERROR);
11821204
}
11831205
checkIfClosed();
11841206
Object obj = getObjectInternal(columnIndex);
@@ -1197,11 +1219,18 @@ public Struct getStruct(int columnIndex) throws SQLException {
11971219
public Map getMap(int columnIndex) throws SQLException {
11981220
LOGGER.debug("Getting Map from column index: {}", columnIndex);
11991221
if (!complexDatatypeSupport) {
1200-
throw new SQLException("Complex datatype support support is disabled");
1222+
LOGGER.error(
1223+
"Complex datatype support support is disabled. Use connection parameter `EnableComplexDatatypeSupport=1` to enable it.");
1224+
throw new DatabricksSQLException(
1225+
"Complex datatype support support is disabled. Use connection parameter `EnableComplexDatatypeSupport=1` to enable it.",
1226+
DatabricksDriverErrorCode.COMPLEX_DATA_TYPE_MAP_CONVERSION_ERROR);
12011227
}
12021228
if (this.resultSetType.equals(ResultSetType.THRIFT_INLINE)
12031229
|| this.resultSetType.equals(ResultSetType.SEA_INLINE)) {
1204-
throw new SQLException("Complex data types are not supported in inline mode");
1230+
LOGGER.error("Complex data types are not supported in inline mode");
1231+
throw new DatabricksSQLException(
1232+
"Complex data types are not supported in inline mode",
1233+
DatabricksDriverErrorCode.COMPLEX_DATA_TYPE_MAP_CONVERSION_ERROR);
12051234
}
12061235
checkIfClosed();
12071236
Object obj = getObjectInternal(columnIndex);

‎src/main/java/com/databricks/jdbc/api/impl/arrow/ArrowStreamResult.java‎

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@
1010
import com.databricks.jdbc.dbclient.impl.common.StatementId;
1111
import com.databricks.jdbc.dbclient.impl.http.DatabricksHttpClientFactory;
1212
import com.databricks.jdbc.exception.DatabricksSQLException;
13+
import com.databricks.jdbc.log.JdbcLogger;
14+
import com.databricks.jdbc.log.JdbcLoggerFactory;
1315
import com.databricks.jdbc.model.client.thrift.generated.TColumnDesc;
1416
import com.databricks.jdbc.model.client.thrift.generated.TFetchResultsResp;
1517
import com.databricks.jdbc.model.client.thrift.generated.TGetResultSetMetadataResp;
@@ -21,13 +23,15 @@
2123
import java.util.ArrayList;
2224
import java.util.List;
2325

26+
/** Result container for Arrow-based query results. */
2427
public class ArrowStreamResult implements IExecutionResult {
25-
28+
private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(ArrowStreamResult.class);
2629
private final ChunkProvider chunkProvider;
2730
private long currentRowIndex = -1;
2831
private boolean isClosed;
2932
private ArrowResultChunk.ArrowResultChunkIterator chunkIterator;
3033
private List<ColumnInfo> columnInfos;
34+
private final IDatabricksSession session;
3135

3236
public ArrowStreamResult(
3337
ResultManifest resultManifest,
@@ -51,6 +55,7 @@ public ArrowStreamResult(
5155
IDatabricksSession session,
5256
IDatabricksHttpClient httpClient)
5357
throws DatabricksSQLException {
58+
this.session = session;
5459
// Check if the result data contains the arrow data inline
5560
boolean isInlineArrow = resultData.getAttachment() != null;
5661

@@ -94,6 +99,7 @@ public ArrowStreamResult(
9499
IDatabricksSession session,
95100
IDatabricksHttpClient httpClient)
96101
throws DatabricksSQLException {
102+
this.session = session;
97103
setColumnInfo(resultsResp.getResultSetMetadata());
98104
if (isInlineArrow) {
99105
this.chunkProvider = new InlineChunkProvider(resultsResp, parentStatement, session);
@@ -126,9 +132,32 @@ public Object getObject(int columnIndex) throws DatabricksSQLException {
126132
if (arrowMetadata == null) {
127133
arrowMetadata = columnInfos.get(columnIndex).getTypeText();
128134
}
135+
136+
// Handle complex type conversion when complex datatype support is disabled
137+
boolean isComplexDatatypeSupportEnabled =
138+
this.session.getConnectionContext().isComplexDatatypeSupportEnabled();
139+
if (!isComplexDatatypeSupportEnabled && isComplexType(requiredType)) {
140+
LOGGER.debug("Complex datatype support is disabled, converting complex type to STRING");
141+
requiredType = ColumnInfoTypeName.STRING;
142+
arrowMetadata = "STRING";
143+
}
144+
129145
return chunkIterator.getColumnObjectAtCurrentRow(columnIndex, requiredType, arrowMetadata);
130146
}
131147

148+
/**
149+
* Checks if the given type is a complex type (ARRAY, MAP, or STRUCT).
150+
*
151+
* @param type The type to check
152+
* @return true if the type is a complex type, false otherwise
153+
*/
154+
@VisibleForTesting
155+
public static boolean isComplexType(ColumnInfoTypeName type) {
156+
return type == ColumnInfoTypeName.ARRAY
157+
|| type == ColumnInfoTypeName.MAP
158+
|| type == ColumnInfoTypeName.STRUCT;
159+
}
160+
132161
/** {@inheritDoc} */
133162
@Override
134163
public long getCurrentRow() {

‎src/test/java/com/databricks/jdbc/api/impl/DatabricksResultSetTest.java‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import com.databricks.jdbc.exception.DatabricksSQLFeatureNotSupportedException;
2323
import com.databricks.jdbc.model.client.thrift.generated.*;
2424
import com.databricks.jdbc.model.core.StatementStatus;
25+
import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode;
2526
import com.databricks.sdk.service.sql.ServiceError;
2627
import com.databricks.sdk.service.sql.StatementState;
2728
import java.io.*;
@@ -1205,4 +1206,46 @@ void testDefaultValuesForNullFields() throws SQLException {
12051206
resultSet.getString(1);
12061207
assertTrue(resultSet.wasNull());
12071208
}
1209+
1210+
@Test
1211+
void testIsComplexTypeThrowsExceptionWhenComplexDatatypeSupportIsDisabled() throws SQLException {
1212+
// Create a ResultSet with complex datatype support disabled
1213+
DatabricksResultSet resultSet = getResultSet(StatementState.SUCCEEDED, null);
1214+
1215+
// Test that getArray throws an exception when complex datatype support is disabled
1216+
DatabricksSQLException arrayException =
1217+
assertThrows(DatabricksSQLException.class, () -> resultSet.getArray(1));
1218+
assertTrue(
1219+
arrayException
1220+
.getMessage()
1221+
.contains(
1222+
"Complex datatype support support is disabled. Use connection parameter `EnableComplexDatatypeSupport=1` to enable it."));
1223+
assertEquals(
1224+
DatabricksDriverErrorCode.COMPLEX_DATA_TYPE_ARRAY_CONVERSION_ERROR.name(),
1225+
arrayException.getSQLState());
1226+
1227+
// Test that getMap throws an exception when complex datatype support is disabled
1228+
DatabricksSQLException mapException =
1229+
assertThrows(DatabricksSQLException.class, () -> resultSet.getMap(1));
1230+
assertTrue(
1231+
mapException
1232+
.getMessage()
1233+
.contains(
1234+
"Complex datatype support support is disabled. Use connection parameter `EnableComplexDatatypeSupport=1` to enable it."));
1235+
assertEquals(
1236+
DatabricksDriverErrorCode.COMPLEX_DATA_TYPE_MAP_CONVERSION_ERROR.name(),
1237+
mapException.getSQLState());
1238+
1239+
// Test that getStruct throws an exception when complex datatype support is disabled
1240+
DatabricksSQLException structException =
1241+
assertThrows(DatabricksSQLException.class, () -> resultSet.getStruct(1));
1242+
assertTrue(
1243+
structException
1244+
.getMessage()
1245+
.contains(
1246+
"Complex datatype support support is disabled. Use connection parameter `EnableComplexDatatypeSupport=1` to enable it."));
1247+
assertEquals(
1248+
DatabricksDriverErrorCode.COMPLEX_DATA_TYPE_STRUCT_CONVERSION_ERROR.name(),
1249+
structException.getSQLState());
1250+
}
12081251
}

‎src/test/java/com/databricks/jdbc/api/impl/arrow/ArrowStreamResultTest.java‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -209,6 +209,20 @@ public void testGetObject() throws Exception {
209209
assertInstanceOf(Double.class, objectInSecondColumn);
210210
}
211211

212+
@Test
213+
public void testComplexTypeHandling() {
214+
assertTrue(ArrowStreamResult.isComplexType(ColumnInfoTypeName.ARRAY));
215+
assertTrue(ArrowStreamResult.isComplexType(ColumnInfoTypeName.MAP));
216+
assertTrue(ArrowStreamResult.isComplexType(ColumnInfoTypeName.STRUCT));
217+
218+
// Non-complex types should return false
219+
assertFalse(ArrowStreamResult.isComplexType(ColumnInfoTypeName.INT));
220+
assertFalse(ArrowStreamResult.isComplexType(ColumnInfoTypeName.STRING));
221+
assertFalse(ArrowStreamResult.isComplexType(ColumnInfoTypeName.DOUBLE));
222+
assertFalse(ArrowStreamResult.isComplexType(ColumnInfoTypeName.BOOLEAN));
223+
assertFalse(ArrowStreamResult.isComplexType(ColumnInfoTypeName.TIMESTAMP));
224+
}
225+
212226
private List<ExternalLink> getChunkLinks(long chunkIndex, boolean isLast) {
213227
List<ExternalLink> chunkLinks = new ArrayList<>();
214228
ExternalLink chunkLink =

0 commit comments

Comments
 (0)