From ee5cb95d690fcb2378993c9fedb7f620cd8ad152 Mon Sep 17 00:00:00 2001 From: Madhav Sainanee Date: Wed, 28 May 2025 11:18:30 +0000 Subject: [PATCH 1/5] Enhance ComplexDataTypeParser with formatting methods for complex types - Added `formatComplexTypeString` method to standardize JSON string representation for complex types when support is disabled. - Introduced `formatMapString` method to convert map JSON strings into a consistent {key:value} format. - Updated `ArrowStreamResult` to utilize the new formatting methods for handling complex types when complex datatype support is disabled. --- .../jdbc/api/impl/ComplexDataTypeParser.java | 85 +++++++++++++++++++ .../api/impl/arrow/ArrowStreamResult.java | 12 ++- 2 files changed, 95 insertions(+), 2 deletions(-) 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..d94015e49d 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,89 @@ 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); + // [{\"key\":1,\"value\":2},{\"key\":3,\"value\":4}] + if (node.isArray() && node.size() > 0 && node.get(0).has("key")) { + // Parse the map type information + 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"); + + if (i > 0) { + result.append(","); + } + + if (isStringKey) { + result.append("\"").append(keyNode.asText()).append("\""); + } else { + result.append(keyNode.asText()); + } + + result.append(":"); + + 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..5965d2885d 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,15 @@ 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"); + if (result != null) { + ComplexDataTypeParser parser = new ComplexDataTypeParser(); + return parser.formatComplexTypeString( + result.toString(), requiredType.name(), arrowMetadata); + } } return chunkIterator.getColumnObjectAtCurrentRow(columnIndex, requiredType, arrowMetadata); From 13452938cc79d446da2912b5d7b1813390110d2e Mon Sep 17 00:00:00 2001 From: Madhav Sainanee Date: Wed, 28 May 2025 11:38:07 +0000 Subject: [PATCH 2/5] unit test --- .../api/impl/ComplexDataTypeParserTest.java | 44 +++++++++++++++++++ 1 file changed, 44 insertions(+) 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..9ccf5057fd 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() { + 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() { + 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() { + 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); + } } From 7fbdd82dadd6704e9397168e508404d87b8117b1 Mon Sep 17 00:00:00 2001 From: Madhav Sainanee Date: Wed, 28 May 2025 11:43:54 +0000 Subject: [PATCH 3/5] Refactor formatMapString method in ComplexDataTypeParser to improve map metadata handling - Changed the initialization of the key-value array in the formatMapString method to use a dynamic array instead of a fixed size. - Removed commented-out code for clarity and improved readability. --- .../com/databricks/jdbc/api/impl/ComplexDataTypeParser.java | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) 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 d94015e49d..2ad3dc180c 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/ComplexDataTypeParser.java +++ b/src/main/java/com/databricks/jdbc/api/impl/ComplexDataTypeParser.java @@ -235,10 +235,8 @@ public String formatComplexTypeString( public String formatMapString(String jsonString, String mapMetadata) { try { JsonNode node = JsonUtil.getMapper().readTree(jsonString); - // [{\"key\":1,\"value\":2},{\"key\":3,\"value\":4}] if (node.isArray() && node.size() > 0 && node.get(0).has("key")) { - // Parse the map type information - String[] kv = new String[] {"STRING", "STRING"}; + String[] kv = new String[2]; if (mapMetadata != null && mapMetadata.startsWith(DatabricksTypeUtil.MAP)) { kv = MetadataParser.parseMapMetadata(mapMetadata).split(",", 2); } From fc7f770133254f216e2fd12d1954545b79da758f Mon Sep 17 00:00:00 2001 From: Madhav Sainanee Date: Wed, 28 May 2025 13:22:08 +0000 Subject: [PATCH 4/5] fix --- .../databricks/jdbc/api/impl/arrow/ArrowStreamResult.java | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) 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 5965d2885d..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 @@ -143,11 +143,8 @@ public Object getObject(int columnIndex) throws DatabricksSQLException { Object result = chunkIterator.getColumnObjectAtCurrentRow( columnIndex, ColumnInfoTypeName.STRING, "STRING"); - if (result != null) { - ComplexDataTypeParser parser = new ComplexDataTypeParser(); - return parser.formatComplexTypeString( - result.toString(), requiredType.name(), arrowMetadata); - } + ComplexDataTypeParser parser = new ComplexDataTypeParser(); + return parser.formatComplexTypeString(result.toString(), requiredType.name(), arrowMetadata); } return chunkIterator.getColumnObjectAtCurrentRow(columnIndex, requiredType, arrowMetadata); From c6b5609e3857fedef09c3c0161ea2ab71e30cd57 Mon Sep 17 00:00:00 2001 From: Madhav Sainanee Date: Thu, 29 May 2025 10:35:19 +0000 Subject: [PATCH 5/5] addresses comments --- .../jdbc/api/impl/ComplexDataTypeParser.java | 14 ++++++++++++-- .../jdbc/api/impl/ComplexDataTypeParserTest.java | 6 +++--- 2 files changed, 15 insertions(+), 5 deletions(-) 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 2ad3dc180c..b145b94be5 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/ComplexDataTypeParser.java +++ b/src/main/java/com/databricks/jdbc/api/impl/ComplexDataTypeParser.java @@ -236,7 +236,7 @@ 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[2]; + String[] kv = new String[] {"STRING", "STRING"}; if (mapMetadata != null && mapMetadata.startsWith(DatabricksTypeUtil.MAP)) { kv = MetadataParser.parseMapMetadata(mapMetadata).split(",", 2); } @@ -253,6 +253,13 @@ public String formatMapString(String jsonString, String mapMetadata) { 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(","); } @@ -265,7 +272,10 @@ public String formatMapString(String jsonString, String mapMetadata) { result.append(":"); - if (isStringValue) { + // 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()); 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 9ccf5057fd..1082aa5d52 100644 --- a/src/test/java/com/databricks/jdbc/api/impl/ComplexDataTypeParserTest.java +++ b/src/test/java/com/databricks/jdbc/api/impl/ComplexDataTypeParserTest.java @@ -148,7 +148,7 @@ void testFormatComplexTypeString_withNonMapType() { } @Test - void testFormatMapString_withIntKeyAndValue() { + void testFormatMapString_withIntKeyAndValue() throws DatabricksParsingException { String jsonString = "[{\"key\":1,\"value\":2},{\"key\":3,\"value\":4}]"; String expected = "{1:2,3:4}"; @@ -157,7 +157,7 @@ void testFormatMapString_withIntKeyAndValue() { } @Test - void testFormatMapString_withStringKeyAndValue() { + void testFormatMapString_withStringKeyAndValue() throws DatabricksParsingException { String jsonString = "[{\"key\":\"a\",\"value\":\"b\"},{\"key\":\"c\",\"value\":\"d\"}]"; String expected = "{\"a\":\"b\",\"c\":\"d\"}"; @@ -166,7 +166,7 @@ void testFormatMapString_withStringKeyAndValue() { } @Test - void testFormatMapString_withMixedTypes() { + void testFormatMapString_withMixedTypes() throws DatabricksParsingException { String jsonString = "[{\"key\":\"a\",\"value\":100},{\"key\":\"b\",\"value\":200}]"; String expected = "{\"a\":100,\"b\":200}";