diff --git a/src/main/java/com/databricks/jdbc/api/impl/ComplexDataTypeParser.java b/src/main/java/com/databricks/jdbc/api/impl/ComplexDataTypeParser.java index dabe8c965b..b145b94be5 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/ComplexDataTypeParser.java +++ b/src/main/java/com/databricks/jdbc/api/impl/ComplexDataTypeParser.java @@ -198,4 +198,97 @@ private Object convertPrimitive(String text, String type) { return text; } } + + /** + * Formats a complex type JSON string into a consistent string format. This is primarily used when + * complex datatype support is disabled. + * + * @param jsonString The JSON string representation of the complex type + * @param complexType The type of complex data (MAP, ARRAY, STRUCT) + * @param typeMetadata The metadata for the type (e.g., "MAP") + * @return A consistently formatted string representation + */ + public String formatComplexTypeString( + String jsonString, String complexType, String typeMetadata) { + if (jsonString == null || complexType == null) { + return jsonString; + } + + try { + if (complexType.equals(DatabricksTypeUtil.MAP)) { + return formatMapString(jsonString, typeMetadata); + } + } catch (Exception e) { + LOGGER.warn("Failed to format complex type representation: {}", e.getMessage()); + } + + return jsonString; + } + + /** + * Formats a map JSON string into the standard {key:value} format. + * + * @param jsonString The JSON string representation of the map + * @param mapMetadata The metadata for the map type (e.g., "MAP") + * @return A map string in the format {key:value,key:value} + */ + public String formatMapString(String jsonString, String mapMetadata) { + try { + JsonNode node = JsonUtil.getMapper().readTree(jsonString); + if (node.isArray() && node.size() > 0 && node.get(0).has("key")) { + String[] kv = new String[] {"STRING", "STRING"}; + if (mapMetadata != null && mapMetadata.startsWith(DatabricksTypeUtil.MAP)) { + kv = MetadataParser.parseMapMetadata(mapMetadata).split(",", 2); + } + + String keyType = kv[0].trim(); + String valueType = kv[1].trim(); + boolean isStringKey = keyType.equalsIgnoreCase(DatabricksTypeUtil.STRING); + boolean isStringValue = valueType.equalsIgnoreCase(DatabricksTypeUtil.STRING); + + StringBuilder result = new StringBuilder("{"); + + for (int i = 0; i < node.size(); i++) { + JsonNode entry = node.get(i); + JsonNode keyNode = entry.get("key"); + JsonNode valueNode = entry.get("value"); + + // Throw error if keyNode is null + if (keyNode == null || keyNode.isNull()) { + throw new DatabricksParsingException( + "Map entry found with null key in JSON: " + entry.toString(), + DatabricksDriverErrorCode.JSON_PARSING_ERROR); + } + + if (i > 0) { + result.append(","); + } + + if (isStringKey) { + result.append("\"").append(keyNode.asText()).append("\""); + } else { + result.append(keyNode.asText()); + } + + result.append(":"); + + // Handle null valueNode + if (valueNode == null || valueNode.isNull()) { + result.append("null"); + } else if (isStringValue) { + result.append("\"").append(valueNode.asText()).append("\""); + } else { + result.append(valueNode.asText()); + } + } + + result.append("}"); + return result.toString(); + } + } catch (Exception e) { + LOGGER.warn("Failed to format map representation: {}", e.getMessage()); + } + + return jsonString; + } } diff --git a/src/main/java/com/databricks/jdbc/api/impl/arrow/ArrowStreamResult.java b/src/main/java/com/databricks/jdbc/api/impl/arrow/ArrowStreamResult.java index 6000267387..712cef1849 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/arrow/ArrowStreamResult.java +++ b/src/main/java/com/databricks/jdbc/api/impl/arrow/ArrowStreamResult.java @@ -2,6 +2,7 @@ import static com.databricks.jdbc.common.util.DatabricksThriftUtil.getTypeFromTypeDesc; +import com.databricks.jdbc.api.impl.ComplexDataTypeParser; import com.databricks.jdbc.api.impl.IExecutionResult; import com.databricks.jdbc.api.internal.IDatabricksSession; import com.databricks.jdbc.api.internal.IDatabricksStatementInternal; @@ -138,8 +139,12 @@ public Object getObject(int columnIndex) throws DatabricksSQLException { this.session.getConnectionContext().isComplexDatatypeSupportEnabled(); if (!isComplexDatatypeSupportEnabled && isComplexType(requiredType)) { LOGGER.debug("Complex datatype support is disabled, converting complex type to STRING"); - requiredType = ColumnInfoTypeName.STRING; - arrowMetadata = "STRING"; + + Object result = + chunkIterator.getColumnObjectAtCurrentRow( + columnIndex, ColumnInfoTypeName.STRING, "STRING"); + ComplexDataTypeParser parser = new ComplexDataTypeParser(); + return parser.formatComplexTypeString(result.toString(), requiredType.name(), arrowMetadata); } return chunkIterator.getColumnObjectAtCurrentRow(columnIndex, requiredType, arrowMetadata); diff --git a/src/test/java/com/databricks/jdbc/api/impl/ComplexDataTypeParserTest.java b/src/test/java/com/databricks/jdbc/api/impl/ComplexDataTypeParserTest.java index 4da815b9c9..1082aa5d52 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/ComplexDataTypeParserTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/ComplexDataTypeParserTest.java @@ -129,4 +129,48 @@ void testComplexPrimitiveConversions() throws DatabricksParsingException { fail("Should not throw: " + e.getMessage()); } } + + @Test + void testFormatComplexTypeString_withMapType() { + String jsonString = "[{\"key\":1,\"value\":2},{\"key\":3,\"value\":4}]"; + String expected = "{1:2,3:4}"; + + String result = parser.formatComplexTypeString(jsonString, "MAP", "MAP"); + assertEquals(expected, result); + } + + @Test + void testFormatComplexTypeString_withNonMapType() { + String jsonString = "[1,2,3]"; + + String result = parser.formatComplexTypeString(jsonString, "ARRAY", "ARRAY"); + assertEquals(jsonString, result); + } + + @Test + void testFormatMapString_withIntKeyAndValue() throws DatabricksParsingException { + String jsonString = "[{\"key\":1,\"value\":2},{\"key\":3,\"value\":4}]"; + String expected = "{1:2,3:4}"; + + String result = parser.formatMapString(jsonString, "MAP"); + assertEquals(expected, result); + } + + @Test + void testFormatMapString_withStringKeyAndValue() throws DatabricksParsingException { + String jsonString = "[{\"key\":\"a\",\"value\":\"b\"},{\"key\":\"c\",\"value\":\"d\"}]"; + String expected = "{\"a\":\"b\",\"c\":\"d\"}"; + + String result = parser.formatMapString(jsonString, "MAP"); + assertEquals(expected, result); + } + + @Test + void testFormatMapString_withMixedTypes() throws DatabricksParsingException { + String jsonString = "[{\"key\":\"a\",\"value\":100},{\"key\":\"b\",\"value\":200}]"; + String expected = "{\"a\":100,\"b\":200}"; + + String result = parser.formatMapString(jsonString, "MAP"); + assertEquals(expected, result); + } }