feat: Add partial support for from_json (#2934)

Author: andygrove

.github/workflows/pr_build_linux.yml

@@ -161,6 +161,7 @@ jobs:
               org.apache.comet.CometStringExpressionSuite
               org.apache.comet.CometBitwiseExpressionSuite
               org.apache.comet.CometMapExpressionSuite
+              org.apache.comet.CometJsonExpressionSuite
               org.apache.comet.expressions.conditional.CometIfSuite
               org.apache.comet.expressions.conditional.CometCoalesceSuite
               org.apache.comet.expressions.conditional.CometCaseWhenSuite

.github/workflows/pr_build_macos.yml

@@ -126,6 +126,7 @@ jobs:
               org.apache.comet.CometStringExpressionSuite
               org.apache.comet.CometBitwiseExpressionSuite
               org.apache.comet.CometMapExpressionSuite
+              org.apache.comet.CometJsonExpressionSuite
               org.apache.comet.expressions.conditional.CometIfSuite
               org.apache.comet.expressions.conditional.CometCoalesceSuite
               org.apache.comet.expressions.conditional.CometCaseWhenSuite

docs/source/user-guide/latest/configs.md

@@ -264,6 +264,7 @@ These settings can be used to determine which parts of the plan are accelerated
 | `spark.comet.expression.IsNaN.enabled` | Enable Comet acceleration for `IsNaN` | true |
 | `spark.comet.expression.IsNotNull.enabled` | Enable Comet acceleration for `IsNotNull` | true |
 | `spark.comet.expression.IsNull.enabled` | Enable Comet acceleration for `IsNull` | true |
+| `spark.comet.expression.JsonToStructs.enabled` | Enable Comet acceleration for `JsonToStructs` | true |
 | `spark.comet.expression.KnownFloatingPointNormalized.enabled` | Enable Comet acceleration for `KnownFloatingPointNormalized` | true |
 | `spark.comet.expression.Length.enabled` | Enable Comet acceleration for `Length` | true |
 | `spark.comet.expression.LessThan.enabled` | Enable Comet acceleration for `LessThan` | true |

native/Cargo.lock

@@ -25,12 +25,13 @@ use datafusion::common::ScalarValue;
 use datafusion::physical_expr::expressions::{LikeExpr, Literal};
 use datafusion::physical_expr::PhysicalExpr;
 use datafusion_comet_proto::spark_expression::Expr;
-use datafusion_comet_spark_expr::{RLike, SubstringExpr};
+use datafusion_comet_spark_expr::{FromJson, RLike, SubstringExpr};
 
 use crate::execution::{
     expressions::extract_expr,
     operators::ExecutionError,
     planner::{expression_registry::ExpressionBuilder, PhysicalPlanner},
+    serde::to_arrow_datatype,
 };
 
 /// Builder for Substring expressions
@@ -98,3 +99,27 @@ impl ExpressionBuilder for RlikeBuilder {
         }
     }
 }
+
+pub struct FromJsonBuilder;
+
+impl ExpressionBuilder for FromJsonBuilder {
+    fn build(
+        &self,
+        spark_expr: &Expr,
+        input_schema: SchemaRef,
+        planner: &PhysicalPlanner,
+    ) -> Result<Arc<dyn PhysicalExpr>, ExecutionError> {
+        let expr = extract_expr!(spark_expr, FromJson);
+        let child = planner.create_expr(
+            expr.child.as_ref().ok_or_else(|| {
+                ExecutionError::GeneralError("FromJson missing child".to_string())
+            })?,
+            input_schema,
+        )?;
+        let schema =
+            to_arrow_datatype(expr.schema.as_ref().ok_or_else(|| {
+                ExecutionError::GeneralError("FromJson missing schema".to_string())
+            })?);
+        Ok(Arc::new(FromJson::new(child, schema, &expr.timezone)))
+    }
+}

native/core/src/execution/expressions/strings.rs

@@ -94,6 +94,7 @@ pub enum ExpressionType {
     CreateNamedStruct,
     GetStructField,
     ToJson,
+    FromJson,
     ToPrettyString,
     ListExtract,
     GetArrayStructFields,
@@ -281,6 +282,8 @@ impl ExpressionRegistry {
             .insert(ExpressionType::Like, Box::new(LikeBuilder));
         self.builders
             .insert(ExpressionType::Rlike, Box::new(RlikeBuilder));
+        self.builders
+            .insert(ExpressionType::FromJson, Box::new(FromJsonBuilder));
     }
 
     /// Extract expression type from Spark protobuf expression
@@ -336,6 +339,7 @@ impl ExpressionRegistry {
             Some(ExprStruct::CreateNamedStruct(_)) => Ok(ExpressionType::CreateNamedStruct),
             Some(ExprStruct::GetStructField(_)) => Ok(ExpressionType::GetStructField),
             Some(ExprStruct::ToJson(_)) => Ok(ExpressionType::ToJson),
+            Some(ExprStruct::FromJson(_)) => Ok(ExpressionType::FromJson),
             Some(ExprStruct::ToPrettyString(_)) => Ok(ExpressionType::ToPrettyString),
             Some(ExprStruct::ListExtract(_)) => Ok(ExpressionType::ListExtract),
             Some(ExprStruct::GetArrayStructFields(_)) => Ok(ExpressionType::GetArrayStructFields),

native/core/src/execution/planner/expression_registry.rs

@@ -85,6 +85,7 @@ message Expr {
     Rand randn = 62;
     EmptyExpr spark_partition_id = 63;
     EmptyExpr monotonically_increasing_id = 64;
+    FromJson from_json = 89;
   }
 }
 
@@ -268,6 +269,12 @@ message ToJson {
   bool ignore_null_fields = 6;
 }
 
+message FromJson {
+  Expr child = 1;
+  DataType schema = 2;
+  string timezone = 3;
+}
+
 enum BinaryOutputStyle {
   UTF8 = 0;
   BASIC = 1;

native/proto/src/proto/expr.proto

@@ -33,6 +33,7 @@ datafusion = { workspace = true }
 chrono-tz = { workspace = true }
 num = { workspace = true }
 regex = { workspace = true }
+serde_json = "1.0"
 thiserror = { workspace = true }
 futures = { workspace = true }
 twox-hash = "2.1.2"