[SPARK-58307][SQL][PYTHON] Add json_typeof function

python/docs/source/reference/pyspark.sql/functions.rst

@@ -587,6 +587,7 @@ JSON Functions
     json_array_length
     json_object_keys
     json_tuple
+    json_typeof
     schema_of_json
     to_json
 

python/pyspark/sql/connect/functions/builtin.py

@@ -2060,6 +2060,13 @@ def json_object_keys(col: "ColumnOrName") -> Column:
 json_object_keys.__doc__ = pysparkfuncs.json_object_keys.__doc__
 
 
+def json_typeof(col: "ColumnOrName") -> Column:
+    return _invoke_function_over_columns("json_typeof", col)
+
+
+json_typeof.__doc__ = pysparkfuncs.json_typeof.__doc__
+
+
 def inline(col: "ColumnOrName") -> Column:
     return _invoke_function_over_columns("inline", col)
 

python/pyspark/sql/functions/__init__.py

@@ -471,6 +471,7 @@
     "json_array_length",
     "json_object_keys",
     "json_tuple",
+    "json_typeof",
     "schema_of_json",
     "to_json",
     # VARIANT Functions

python/pyspark/sql/functions/builtin.py

@@ -23375,6 +23375,40 @@ def json_object_keys(col: "ColumnOrName") -> Column:
     return _invoke_function_over_columns("json_object_keys", col)
 
 
+@_try_remote_functions
+def json_typeof(col: "ColumnOrName") -> Column:
+    """
+    Returns the type of the outermost JSON value as a string: one of 'object', 'array',
+    'string', 'number', 'boolean', or 'null'. Returns null if the input is not a valid JSON
+    string or is an empty string.
+
+    .. versionadded:: 4.3.0
+
+    Parameters
+    ----------
+    col: :class:`~pyspark.sql.Column` or str
+        target column to compute on.
+        A column that evaluates to a string.
+
+    Returns
+    -------
+    :class:`~pyspark.sql.Column`
+        the type of the outermost JSON value.
+        Returns a column that evaluates to a string.
+
+    See Also
+    --------
+    :meth:`pyspark.sql.functions.json_object_keys`
+
+    Examples
+    --------
+    >>> df = spark.createDataFrame([('{"a": 1}',), ('[1, 2, 3]',), ('123',), ('',)], ['data'])
+    >>> df.select(json_typeof(df.data).alias('r')).collect()
+    [Row(r='object'), Row(r='array'), Row(r='number'), Row(r=None)]
+    """
+    return _invoke_function_over_columns("json_typeof", col)
+
+
 # TODO: Fix and add an example for StructType with Spark Connect
 #   e.g., StructType([StructField("a", IntegerType())])
 @_try_remote_functions

sql/api/src/main/scala/org/apache/spark/sql/functions.scala

@@ -14576,6 +14576,19 @@ object functions {
    */
   def json_object_keys(e: Column): Column = Column.fn("json_object_keys", e)
 
+  /**
+   * Returns the type of the outermost JSON value as a string: one of 'object', 'array', 'string',
+   * 'number', 'boolean', or 'null'. Returns null for invalid or empty input.
+   *
+   * @param e
+   *   the JSON string column. A column that evaluates to a string.
+   * @group json_funcs
+   * @since 4.3.0
+   * @return
+   *   Returns a column that evaluates to a string.
+   */
+  def json_typeof(e: Column): Column = Column.fn("json_typeof", e)
+
   // scalastyle:off line.size.limit
   /**
    * (Scala-specific) Converts a column containing a `StructType`, `ArrayType` or a `MapType` into

sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/json/JsonExpressionUtils.java

@@ -77,4 +77,32 @@ public static GenericArrayData jsonObjectKeys(UTF8String json) {
       return null;
     }
   }
+
+  public static UTF8String jsonTypeof(UTF8String json) {
+    try (JsonParser jsonParser =
+        CreateJacksonParser.utf8String(SharedFactory.jsonFactory(), json)) {
+      JsonToken token = jsonParser.nextToken();
+      if (token == null) {
+        return null;
+      }
+      String type = switch (token) {
+        case START_OBJECT -> "object";
+        case START_ARRAY -> "array";
+        case VALUE_STRING -> "string";
+        case VALUE_NUMBER_INT, VALUE_NUMBER_FLOAT -> "number";
+        case VALUE_TRUE, VALUE_FALSE -> "boolean";
+        case VALUE_NULL -> "null";
+        default -> null;
+      };
+      if (type == null) {
+        return null;
+      }
+      // Consume the value so malformed input surfaces as a parse error and returns null,
+      // matching json_object_keys and json_array_length.
+      jsonParser.skipChildren();
+      return UTF8String.fromString(type);
+    } catch (IOException e) {
+      return null;
+    }
+  }
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala

@@ -975,6 +975,7 @@ object FunctionRegistry {
     expression[SchemaOfJson]("schema_of_json"),
     expression[LengthOfJsonArray]("json_array_length"),
     expression[JsonObjectKeys]("json_object_keys"),
+    expression[JsonTypeof]("json_typeof"),
 
     // Variant
     expressionBuilder("parse_json", ParseJsonExpressionBuilder),

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala

@@ -709,3 +709,49 @@ case class JsonObjectKeys(child: Expression)
   override protected def withNewChildInternal(newChild: Expression): JsonObjectKeys =
     copy(child = newChild)
 }
+
+/**
+ * A function which returns the type of the outermost JSON value as a string.
+ */
+@ExpressionDescription(
+  usage = "_FUNC_(json) - Returns the type of the outermost JSON value, or null if invalid.",
+  arguments = """
+    Arguments:
+      * json - A JSON string. Returns the type of the outermost value ('object', 'array',
+          'string', 'number', 'boolean', 'null'), or null for an invalid or empty string.
+        An expression that evaluates to a string.
+  """,
+  examples = """
+    Examples:
+      > SELECT _FUNC_('{"a": 1}');
+        object
+      > SELECT _FUNC_('[1, 2, 3]');
+        array
+      > SELECT _FUNC_('123');
+        number
+  """,
+  group = "json_funcs",
+  since = "4.3.0"
+)
+case class JsonTypeof(child: Expression)
+  extends UnaryExpression
+  with ExpectsInputTypes
+  with RuntimeReplaceable
+  with DefaultStringProducingExpression {
+
+  override def inputTypes: Seq[AbstractDataType] =
+    Seq(StringTypeWithCollation(supportsTrimCollation = true))
+  override def nullable: Boolean = true
+  override def prettyName: String = "json_typeof"
+
+  override def replacement: Expression = StaticInvoke(
+    classOf[JsonExpressionUtils],
+    dataType,
+    "jsonTypeof",
+    Seq(child),
+    inputTypes
+  )
+
+  override protected def withNewChildInternal(newChild: Expression): JsonTypeof =
+    copy(child = newChild)
+}

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/JsonExpressionsSuite.scala

@@ -909,6 +909,29 @@ class JsonExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
     }
   }
 
+  test("json_typeof") {
+    Seq(
+      // Invalid or empty inputs return null.
+      ("", null),
+      ("bad", null),
+      ("""{"key": 45, "random_string"}""", null),
+      // Valid JSON values return the type of the outermost value.
+      ("{}", "object"),
+      ("""{"key": 1, "arr": [1, 2]}""", "object"),
+      ("[]", "array"),
+      ("[1, 2, 3]", "array"),
+      ("\"hello\"", "string"),
+      ("123", "number"),
+      ("1.5", "number"),
+      ("true", "boolean"),
+      ("false", "boolean"),
+      ("null", "null")
+    ).foreach {
+      case (input, expected) =>
+        checkEvaluation(JsonTypeof(Literal(input)), expected)
+    }
+  }
+
   test("SPARK-35320: from_json should fail with a key type different of StringType") {
     Seq(
       (MapType(IntegerType, StringType), """{"1": "test"}"""),

sql/core/src/test/resources/sql-functions/sql-expression-schema.md

@@ -187,6 +187,7 @@
 | org.apache.spark.sql.catalyst.expressions.JsonObjectKeys | json_object_keys | SELECT json_object_keys('{}') | struct<json_object_keys({}):array<string>> |
 | org.apache.spark.sql.catalyst.expressions.JsonToStructs | from_json | SELECT from_json('{"a":1, "b":0.8}', 'a INT, b DOUBLE') | struct<from_json({"a":1, "b":0.8}):struct<a:int,b:double>> |
 | org.apache.spark.sql.catalyst.expressions.JsonTuple | json_tuple | SELECT json_tuple('{"a":1, "b":2}', 'a', 'b') | struct<c0:string,c1:string> |
+| org.apache.spark.sql.catalyst.expressions.JsonTypeof | json_typeof | SELECT json_typeof('{"a": 1}') | struct<json_typeof({"a": 1}):string> |
 | org.apache.spark.sql.catalyst.expressions.KllSketchGetNBigint | kll_sketch_get_n_bigint | SELECT kll_sketch_get_n_bigint(kll_sketch_agg_bigint(col)) FROM VALUES (1), (2), (3), (4), (5) tab(col) | struct<kll_sketch_get_n_bigint(kll_sketch_agg_bigint(col)):bigint> |
 | org.apache.spark.sql.catalyst.expressions.KllSketchGetNDouble | kll_sketch_get_n_double | SELECT kll_sketch_get_n_double(kll_sketch_agg_double(col)) FROM VALUES (CAST(1.0 AS DOUBLE)), (CAST(2.0 AS DOUBLE)), (CAST(3.0 AS DOUBLE)), (CAST(4.0 AS DOUBLE)), (CAST(5.0 AS DOUBLE)) tab(col) | struct<kll_sketch_get_n_double(kll_sketch_agg_double(col)):bigint> |
 | org.apache.spark.sql.catalyst.expressions.KllSketchGetNFloat | kll_sketch_get_n_float | SELECT kll_sketch_get_n_float(kll_sketch_agg_float(col)) FROM VALUES (CAST(1.0 AS FLOAT)), (CAST(2.0 AS FLOAT)), (CAST(3.0 AS FLOAT)), (CAST(4.0 AS FLOAT)), (CAST(5.0 AS FLOAT)) tab(col) | struct<kll_sketch_get_n_float(kll_sketch_agg_float(col)):bigint> |

sql/core/src/test/resources/sql-tests/analyzer-results/json-functions.sql.out

@@ -742,6 +742,151 @@ Project [json_object_keys([1, 2, 3]) AS json_object_keys([1, 2, 3])#x]
 +- OneRowRelation
 
 
+-- !query
+select json_typeof()
+-- !query analysis
+org.apache.spark.sql.AnalysisException
+{
+  "errorClass" : "WRONG_NUM_ARGS.WITHOUT_SUGGESTION",
+  "sqlState" : "42605",
+  "messageParameters" : {
+    "actualNum" : "0",
+    "docroot" : "https://spark.apache.org/docs/latest",
+    "expectedNum" : "1",
+    "functionName" : "`json_typeof`"
+  },
+  "queryContext" : [ {
+    "objectType" : "",
+    "objectName" : "",
+    "startIndex" : 8,
+    "stopIndex" : 20,
+    "fragment" : "json_typeof()"
+  } ]
+}
+
+
+-- !query
+select json_typeof(null)
+-- !query analysis
+Project [json_typeof(null) AS json_typeof(NULL)#x]
++- OneRowRelation
+
+
+-- !query
+select json_typeof(200)
+-- !query analysis
+org.apache.spark.sql.catalyst.ExtendedAnalysisException
+{
+  "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE",
+  "sqlState" : "42K09",
+  "messageParameters" : {
+    "inputSql" : "\"200\"",
+    "inputType" : "\"INT\"",
+    "paramIndex" : "first",
+    "requiredType" : "\"STRING\"",
+    "sqlExpr" : "\"json_typeof(200)\""
+  },
+  "queryContext" : [ {
+    "objectType" : "",
+    "objectName" : "",
+    "startIndex" : 8,
+    "stopIndex" : 23,
+    "fragment" : "json_typeof(200)"
+  } ]
+}
+
+
+-- !query
+select json_typeof('')
+-- !query analysis
+Project [json_typeof() AS json_typeof()#x]
++- OneRowRelation
+
+
+-- !query
+select json_typeof('{}')
+-- !query analysis
+Project [json_typeof({}) AS json_typeof({})#x]
++- OneRowRelation
+
+
+-- !query
+select json_typeof('{"key": 1, "arr": [1, 2]}')
+-- !query analysis
+Project [json_typeof({"key": 1, "arr": [1, 2]}) AS json_typeof({"key": 1, "arr": [1, 2]})#x]
++- OneRowRelation
+
+
+-- !query
+select json_typeof('[]')
+-- !query analysis
+Project [json_typeof([]) AS json_typeof([])#x]
++- OneRowRelation
+
+
+-- !query
+select json_typeof('[1, 2, 3]')
+-- !query analysis
+Project [json_typeof([1, 2, 3]) AS json_typeof([1, 2, 3])#x]
++- OneRowRelation
+
+
+-- !query
+select json_typeof('"hello"')
+-- !query analysis
+Project [json_typeof("hello") AS json_typeof("hello")#x]
++- OneRowRelation
+
+
+-- !query
+select json_typeof('123')
+-- !query analysis
+Project [json_typeof(123) AS json_typeof(123)#x]
++- OneRowRelation
+
+
+-- !query
+select json_typeof('1.5')
+-- !query analysis
+Project [json_typeof(1.5) AS json_typeof(1.5)#x]
++- OneRowRelation
+
+
+-- !query
+select json_typeof('true')
+-- !query analysis
+Project [json_typeof(true) AS json_typeof(true)#x]
++- OneRowRelation
+
+
+-- !query
+select json_typeof('false')
+-- !query analysis
+Project [json_typeof(false) AS json_typeof(false)#x]
++- OneRowRelation
+
+
+-- !query
+select json_typeof('null')
+-- !query analysis
+Project [json_typeof(null) AS json_typeof(null)#x]
++- OneRowRelation
+
+
+-- !query
+select json_typeof('bad')
+-- !query analysis
+Project [json_typeof(bad) AS json_typeof(bad)#x]
++- OneRowRelation
+
+
+-- !query
+select json_typeof('{"key": 45, "random_string"}')
+-- !query analysis
+Project [json_typeof({"key": 45, "random_string"}) AS json_typeof({"key": 45, "random_string"})#x]
++- OneRowRelation
+
+
 -- !query
 DROP VIEW IF EXISTS jsonTable
 -- !query analysis

sql/core/src/test/resources/sql-tests/inputs/json-functions.sql

@@ -108,6 +108,24 @@ select json_object_keys('{[1,2]}');
 select json_object_keys('{"key": 45, "random_string"}');
 select json_object_keys('[1, 2, 3]');
 
+-- json_typeof
+select json_typeof();
+select json_typeof(null);
+select json_typeof(200);
+select json_typeof('');
+select json_typeof('{}');
+select json_typeof('{"key": 1, "arr": [1, 2]}');
+select json_typeof('[]');
+select json_typeof('[1, 2, 3]');
+select json_typeof('"hello"');
+select json_typeof('123');
+select json_typeof('1.5');
+select json_typeof('true');
+select json_typeof('false');
+select json_typeof('null');
+select json_typeof('bad');
+select json_typeof('{"key": 45, "random_string"}');
+
 -- Clean up
 DROP VIEW IF EXISTS jsonTable;