Add indentation parameter to json_format function

Author: aayushsh06

core/trino-main/src/main/java/io/trino/operator/scalar/JsonFunctions.java

@@ -14,11 +14,18 @@
 package io.trino.operator.scalar;
 
 import com.fasterxml.jackson.core.JsonFactory;
+import com.fasterxml.jackson.core.JsonGenerator;
 import com.fasterxml.jackson.core.JsonParser;
 import com.fasterxml.jackson.core.JsonToken;
+import com.fasterxml.jackson.core.util.DefaultIndenter;
+import com.fasterxml.jackson.core.util.DefaultPrettyPrinter;
 import com.fasterxml.jackson.databind.MappingJsonFactory;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.primitives.Doubles;
+import io.airlift.json.ObjectMapperProvider;
+import io.airlift.slice.DynamicSliceOutput;
 import io.airlift.slice.Slice;
+import io.airlift.slice.SliceOutput;
 import io.trino.plugin.base.util.JsonTypeUtil;
 import io.trino.spi.TrinoException;
 import io.trino.spi.function.LiteralParameter;
@@ -32,6 +39,7 @@
 import io.trino.type.JsonPathType;
 
 import java.io.IOException;
+import java.io.OutputStream;
 import java.util.LinkedList;
 import java.util.List;
 
@@ -45,6 +53,7 @@
 import static com.fasterxml.jackson.core.JsonToken.VALUE_NUMBER_INT;
 import static com.fasterxml.jackson.core.JsonToken.VALUE_STRING;
 import static com.fasterxml.jackson.core.JsonToken.VALUE_TRUE;
+import static com.fasterxml.jackson.databind.SerializationFeature.ORDER_MAP_ENTRIES_BY_KEYS;
 import static io.airlift.slice.Slices.utf8Slice;
 import static io.trino.plugin.base.util.JsonUtils.jsonFactoryBuilder;
 import static io.trino.spi.StandardErrorCode.INVALID_FUNCTION_ARGUMENT;
@@ -61,6 +70,9 @@ public final class JsonFunctions
     private static final JsonFactory MAPPING_JSON_FACTORY = new MappingJsonFactory()
             .disable(CANONICALIZE_FIELD_NAMES);
 
+    private static final ObjectMapper SORTED_MAPPER = new ObjectMapperProvider().get()
+            .configure(ORDER_MAP_ENTRIES_BY_KEYS, true);
+
     private JsonFunctions() {}
 
     @ScalarOperator(OperatorType.CAST)
@@ -124,6 +136,41 @@ public static Slice jsonFormat(@SqlType(StandardTypes.JSON) Slice slice)
         return slice;
     }
 
+    @ScalarFunction
+    @SqlType(StandardTypes.VARCHAR)
+    public static Slice jsonFormat(@SqlType(StandardTypes.JSON) Slice slice, @SqlType(StandardTypes.BIGINT) long indentSpaces)
+    {
+        if (indentSpaces < 0) {
+            throw new TrinoException(INVALID_FUNCTION_ARGUMENT, "Indentation spaces must be non-negative, got: " + indentSpaces);
+        }
+        if (indentSpaces == 0) {
+            return slice;
+        }
+
+        try {
+            Object jsonValue;
+            try (JsonParser parser = createJsonParser(JSON_FACTORY, slice)) {
+                jsonValue = SORTED_MAPPER.readValue(parser, Object.class);
+            }
+
+            DefaultPrettyPrinter prettyPrinter = new DefaultPrettyPrinter();
+            String indentString = " ".repeat((int) indentSpaces);
+            prettyPrinter.indentObjectsWith(DefaultIndenter.SYSTEM_LINEFEED_INSTANCE.withIndent(indentString));
+            prettyPrinter.indentArraysWith(DefaultIndenter.SYSTEM_LINEFEED_INSTANCE.withIndent(indentString));
+
+            SliceOutput output = new DynamicSliceOutput((int) (slice.length() * 1.5));
+            try (JsonGenerator generator = JSON_FACTORY.createGenerator((OutputStream) output)) {
+                generator.setPrettyPrinter(prettyPrinter);
+                SORTED_MAPPER.writeValue(generator, jsonValue);
+            }
+
+            return output.slice();
+        }
+        catch (IOException e) {
+            throw new TrinoException(INVALID_FUNCTION_ARGUMENT, "Failed to format JSON: " + truncateIfNecessaryForErrorMessage(slice), e);
+        }
+    }
+
     @ScalarFunction
     @LiteralParameters("x")
     @SqlType(StandardTypes.JSON)

core/trino-main/src/test/java/io/trino/operator/scalar/TestJsonFunctions.java

@@ -699,6 +699,42 @@ public void testJsonFormat()
                 .isEqualTo("[\"a\",\"b\"]");
     }
 
+    @Test
+    public void testJsonFormatWithIndentation()
+    {
+        assertThat(assertions.function("json_format", "JSON '{\"a\":1,\"b\":2}'", "2"))
+                .hasType(VARCHAR)
+                .isEqualTo("{\n  \"a\" : 1,\n  \"b\" : 2\n}");
+
+        assertThat(assertions.function("json_format", "JSON '{\"a\":1,\"b\":2}'", "4"))
+                .hasType(VARCHAR)
+                .isEqualTo("{\n    \"a\" : 1,\n    \"b\" : 2\n}");
+
+        assertThat(assertions.function("json_format", "JSON '{\"a\":1,\"b\":2}'", "0"))
+                .hasType(VARCHAR)
+                .isEqualTo("{\"a\":1,\"b\":2}");
+
+        assertThat(assertions.function("json_format", "JSON '{\"a\":{\"x\":1},\"b\":[1,2]}'", "2"))
+                .hasType(VARCHAR)
+                .isEqualTo("{\n  \"a\" : {\n    \"x\" : 1\n  },\n  \"b\" : [\n    1,\n    2\n  ]\n}");
+
+        assertThat(assertions.function("json_format", "JSON '[1,2,3]'", "2"))
+                .hasType(VARCHAR)
+                .isEqualTo("[\n  1,\n  2,\n  3\n]");
+
+        assertThat(assertions.function("json_format", "JSON '{\"id\":9223372036854775807}'", "2"))
+                .hasType(VARCHAR)
+                .isEqualTo("{\n  \"id\" : 9223372036854775807\n}");
+    }
+
+    @Test
+    public void testJsonFormatWithIndentationInvalid()
+    {
+        assertTrinoExceptionThrownBy(assertions.function("json_format", "JSON '{\"a\":1}'", "-1")::evaluate)
+                .hasErrorCode(INVALID_FUNCTION_ARGUMENT)
+                .hasMessageContaining("Indentation spaces must be non-negative");
+    }
+
     @Test
     public void testJsonSize()
     {

docs/src/main/sphinx/functions/json.md

@@ -1894,6 +1894,32 @@ SELECT json_format(JSON '[1, 2, 3]'); -- '[1,2,3]'
 SELECT json_format(JSON '"a"');       -- '"a"'
 ```
 
+:::{function} json_format(json, indent_spaces) -> varchar
+:no-index:
+Returns the JSON text serialized from the input JSON value with specified
+indentation. The `indent_spaces` parameter specifies the number of spaces to use
+for indentation. This is useful for formatting JSON output in a human-readable
+format, especially when working with clients that cannot handle large integers
+in compact JSON strings.
+
+```
+SELECT json_format(JSON '{"a": 1, "b": 2}', 2);
+-- '{
+--   "a" : 1,
+--   "b" : 2
+-- }'
+
+SELECT json_format(JSON '{"id": 9223372036854775807}', 2);
+-- '{
+--   "id" : 9223372036854775807
+-- }'
+```
+
+When `indent_spaces` is 0, the function returns the compact format (same as
+the single-argument version). The `indent_spaces` parameter must be
+non-negative.
+:::
+
 :::{note}
 {func}`json_format` and `CAST(json AS VARCHAR)` have completely
 different semantics.