Add ExpressionPlacement::Leaf for JSON UDFs to enable scan-level pushdown

friendlymatthew · claude · friendlymatthew · commit 3a907c576e1f · 2026-03-20T16:23:09.000-04:00
This allows the ExtractLeafExpressions optimizer to push JSON accessor
functions (json_as_text, json_get, etc.) into the Parquet scan projection,
where custom PhysicalExprAdapters can rewrite them for shredded/reduced
column fallback.

Co-Authored-By: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/src/json_as_text.rs b/src/json_as_text.rs
@@ -56,6 +56,22 @@ impl ScalarUDFImpl for JsonAsText {
     fn aliases(&self) -> &[String] {
         &self.aliases
     }
+
+    fn placement(
+        &self,
+        args: &[datafusion::logical_expr::ExpressionPlacement],
+    ) -> datafusion::logical_expr::ExpressionPlacement {
+        // If the first argument is a column and the remaining arguments are literals (a path)
+        // then we can push this UDF down to the leaf nodes.
+        if args.len() >= 2
+            && matches!(args[0], datafusion::logical_expr::ExpressionPlacement::Column)
+            && args[1..].iter().all(|arg| matches!(arg, datafusion::logical_expr::ExpressionPlacement::Literal))
+        {
+            datafusion::logical_expr::ExpressionPlacement::MoveTowardsLeafNodes
+        } else {
+            datafusion::logical_expr::ExpressionPlacement::KeepInPlace
+        }
+    }
 }
 
 impl InvokeResult for StringArray {
diff --git a/src/json_contains.rs b/src/json_contains.rs
@@ -60,6 +60,22 @@ impl ScalarUDFImpl for JsonContains {
     fn aliases(&self) -> &[String] {
         &self.aliases
     }
+
+    fn placement(
+        &self,
+        args: &[datafusion::logical_expr::ExpressionPlacement],
+    ) -> datafusion::logical_expr::ExpressionPlacement {
+        // If the first argument is a column and the remaining arguments are literals (a path)
+        // then we can push this UDF down to the leaf nodes.
+        if args.len() >= 2
+            && matches!(args[0], datafusion::logical_expr::ExpressionPlacement::Column)
+            && args[1..].iter().all(|arg| matches!(arg, datafusion::logical_expr::ExpressionPlacement::Literal))
+        {
+            datafusion::logical_expr::ExpressionPlacement::MoveTowardsLeafNodes
+        } else {
+            datafusion::logical_expr::ExpressionPlacement::KeepInPlace
+        }
+    }
 }
 
 impl InvokeResult for BooleanArray {
diff --git a/src/json_get.rs b/src/json_get.rs
@@ -62,6 +62,22 @@ impl ScalarUDFImpl for JsonGet {
     fn aliases(&self) -> &[String] {
         &self.aliases
     }
+
+    fn placement(
+        &self,
+        args: &[datafusion::logical_expr::ExpressionPlacement],
+    ) -> datafusion::logical_expr::ExpressionPlacement {
+        // If the first argument is a column and the remaining arguments are literals (a path)
+        // then we can push this UDF down to the leaf nodes.
+        if args.len() >= 2
+            && matches!(args[0], datafusion::logical_expr::ExpressionPlacement::Column)
+            && args[1..].iter().all(|arg| matches!(arg, datafusion::logical_expr::ExpressionPlacement::Literal))
+        {
+            datafusion::logical_expr::ExpressionPlacement::MoveTowardsLeafNodes
+        } else {
+            datafusion::logical_expr::ExpressionPlacement::KeepInPlace
+        }
+    }
 }
 
 impl InvokeResult for JsonUnion {
diff --git a/src/json_get_array.rs b/src/json_get_array.rs
@@ -64,6 +64,22 @@ impl ScalarUDFImpl for JsonGetArray {
     fn aliases(&self) -> &[String] {
         &self.aliases
     }
+
+    fn placement(
+        &self,
+        args: &[datafusion::logical_expr::ExpressionPlacement],
+    ) -> datafusion::logical_expr::ExpressionPlacement {
+        // If the first argument is a column and the remaining arguments are literals (a path)
+        // then we can push this UDF down to the leaf nodes.
+        if args.len() >= 2
+            && matches!(args[0], datafusion::logical_expr::ExpressionPlacement::Column)
+            && args[1..].iter().all(|arg| matches!(arg, datafusion::logical_expr::ExpressionPlacement::Literal))
+        {
+            datafusion::logical_expr::ExpressionPlacement::MoveTowardsLeafNodes
+        } else {
+            datafusion::logical_expr::ExpressionPlacement::KeepInPlace
+        }
+    }
 }
 
 #[derive(Debug)]
diff --git a/src/json_get_bool.rs b/src/json_get_bool.rs
@@ -55,6 +55,22 @@ impl ScalarUDFImpl for JsonGetBool {
     fn aliases(&self) -> &[String] {
         &self.aliases
     }
+
+    fn placement(
+        &self,
+        args: &[datafusion::logical_expr::ExpressionPlacement],
+    ) -> datafusion::logical_expr::ExpressionPlacement {
+        // If the first argument is a column and the remaining arguments are literals (a path)
+        // then we can push this UDF down to the leaf nodes.
+        if args.len() >= 2
+            && matches!(args[0], datafusion::logical_expr::ExpressionPlacement::Column)
+            && args[1..].iter().all(|arg| matches!(arg, datafusion::logical_expr::ExpressionPlacement::Literal))
+        {
+            datafusion::logical_expr::ExpressionPlacement::MoveTowardsLeafNodes
+        } else {
+            datafusion::logical_expr::ExpressionPlacement::KeepInPlace
+        }
+    }
 }
 
 fn jiter_json_get_bool(json_data: Option<&str>, path: &[JsonPath]) -> Result<bool, GetError> {
diff --git a/src/json_get_float.rs b/src/json_get_float.rs
@@ -56,6 +56,22 @@ impl ScalarUDFImpl for JsonGetFloat {
     fn aliases(&self) -> &[String] {
         &self.aliases
     }
+
+    fn placement(
+        &self,
+        args: &[datafusion::logical_expr::ExpressionPlacement],
+    ) -> datafusion::logical_expr::ExpressionPlacement {
+        // If the first argument is a column and the remaining arguments are literals (a path)
+        // then we can push this UDF down to the leaf nodes.
+        if args.len() >= 2
+            && matches!(args[0], datafusion::logical_expr::ExpressionPlacement::Column)
+            && args[1..].iter().all(|arg| matches!(arg, datafusion::logical_expr::ExpressionPlacement::Literal))
+        {
+            datafusion::logical_expr::ExpressionPlacement::MoveTowardsLeafNodes
+        } else {
+            datafusion::logical_expr::ExpressionPlacement::KeepInPlace
+        }
+    }
 }
 
 impl InvokeResult for Float64Array {
diff --git a/src/json_get_int.rs b/src/json_get_int.rs
@@ -56,6 +56,22 @@ impl ScalarUDFImpl for JsonGetInt {
     fn aliases(&self) -> &[String] {
         &self.aliases
     }
+
+    fn placement(
+        &self,
+        args: &[datafusion::logical_expr::ExpressionPlacement],
+    ) -> datafusion::logical_expr::ExpressionPlacement {
+        // If the first argument is a column and the remaining arguments are literals (a path)
+        // then we can push this UDF down to the leaf nodes.
+        if args.len() >= 2
+            && matches!(args[0], datafusion::logical_expr::ExpressionPlacement::Column)
+            && args[1..].iter().all(|arg| matches!(arg, datafusion::logical_expr::ExpressionPlacement::Literal))
+        {
+            datafusion::logical_expr::ExpressionPlacement::MoveTowardsLeafNodes
+        } else {
+            datafusion::logical_expr::ExpressionPlacement::KeepInPlace
+        }
+    }
 }
 
 impl InvokeResult for Int64Array {
diff --git a/src/json_get_json.rs b/src/json_get_json.rs
@@ -54,6 +54,22 @@ impl ScalarUDFImpl for JsonGetJson {
     fn aliases(&self) -> &[String] {
         &self.aliases
     }
+
+    fn placement(
+        &self,
+        args: &[datafusion::logical_expr::ExpressionPlacement],
+    ) -> datafusion::logical_expr::ExpressionPlacement {
+        // If the first argument is a column and the remaining arguments are literals (a path)
+        // then we can push this UDF down to the leaf nodes.
+        if args.len() >= 2
+            && matches!(args[0], datafusion::logical_expr::ExpressionPlacement::Column)
+            && args[1..].iter().all(|arg| matches!(arg, datafusion::logical_expr::ExpressionPlacement::Literal))
+        {
+            datafusion::logical_expr::ExpressionPlacement::MoveTowardsLeafNodes
+        } else {
+            datafusion::logical_expr::ExpressionPlacement::KeepInPlace
+        }
+    }
 }
 
 fn jiter_json_get_json(opt_json: Option<&str>, path: &[JsonPath]) -> Result<String, GetError> {
diff --git a/src/json_get_str.rs b/src/json_get_str.rs
@@ -55,6 +55,22 @@ impl ScalarUDFImpl for JsonGetStr {
     fn aliases(&self) -> &[String] {
         &self.aliases
     }
+
+    fn placement(
+        &self,
+        args: &[datafusion::logical_expr::ExpressionPlacement],
+    ) -> datafusion::logical_expr::ExpressionPlacement {
+        // If the first argument is a column and the remaining arguments are literals (a path)
+        // then we can push this UDF down to the leaf nodes.
+        if args.len() >= 2
+            && matches!(args[0], datafusion::logical_expr::ExpressionPlacement::Column)
+            && args[1..].iter().all(|arg| matches!(arg, datafusion::logical_expr::ExpressionPlacement::Literal))
+        {
+            datafusion::logical_expr::ExpressionPlacement::MoveTowardsLeafNodes
+        } else {
+            datafusion::logical_expr::ExpressionPlacement::KeepInPlace
+        }
+    }
 }
 
 fn jiter_json_get_str(json_data: Option<&str>, path: &[JsonPath]) -> Result<String, GetError> {
diff --git a/src/json_length.rs b/src/json_length.rs
@@ -56,6 +56,22 @@ impl ScalarUDFImpl for JsonLength {
     fn aliases(&self) -> &[String] {
         &self.aliases
     }
+
+    fn placement(
+        &self,
+        args: &[datafusion::logical_expr::ExpressionPlacement],
+    ) -> datafusion::logical_expr::ExpressionPlacement {
+        // If the first argument is a column and the remaining arguments are literals (a path)
+        // then we can push this UDF down to the leaf nodes.
+        if args.len() >= 2
+            && matches!(args[0], datafusion::logical_expr::ExpressionPlacement::Column)
+            && args[1..].iter().all(|arg| matches!(arg, datafusion::logical_expr::ExpressionPlacement::Literal))
+        {
+            datafusion::logical_expr::ExpressionPlacement::MoveTowardsLeafNodes
+        } else {
+            datafusion::logical_expr::ExpressionPlacement::KeepInPlace
+        }
+    }
 }
 
 impl InvokeResult for UInt64Array {
diff --git a/src/json_object_keys.rs b/src/json_object_keys.rs
@@ -60,6 +60,22 @@ impl ScalarUDFImpl for JsonObjectKeys {
     fn aliases(&self) -> &[String] {
         &self.aliases
     }
+
+    fn placement(
+        &self,
+        args: &[datafusion::logical_expr::ExpressionPlacement],
+    ) -> datafusion::logical_expr::ExpressionPlacement {
+        // If the first argument is a column and the remaining arguments are literals (a path)
+        // then we can push this UDF down to the leaf nodes.
+        if args.len() >= 2
+            && matches!(args[0], datafusion::logical_expr::ExpressionPlacement::Column)
+            && args[1..].iter().all(|arg| matches!(arg, datafusion::logical_expr::ExpressionPlacement::Literal))
+        {
+            datafusion::logical_expr::ExpressionPlacement::MoveTowardsLeafNodes
+        } else {
+            datafusion::logical_expr::ExpressionPlacement::KeepInPlace
+        }
+    }
 }
 
 /// Struct used to build a `ListArray` from the result of `jiter_json_object_keys`.
diff --git a/tests/main.rs b/tests/main.rs
@@ -775,9 +775,9 @@ async fn test_plan_json_get_cte() {
         select name, json_get(j, 0) v from t
     ";
     let expected = [
-        "Projection: t.name, json_get(t.j, Int64(0)) AS v",
+        "Projection: t.name, __datafusion_extracted_1 AS v",
         "  SubqueryAlias: t",
-        "    Projection: test.name, json_get(test.json_data, Utf8(\"foo\")) AS j",
+        "    Projection: test.name, json_get(json_get(test.json_data, Utf8(\"foo\")), Int64(0)) AS __datafusion_extracted_1",
         "      TableScan: test projection=[name, json_data]",
     ];
 
@@ -1268,7 +1268,7 @@ async fn test_plan_double_arrow_double_nested_cast() {
 
     // NB: json_as_text(..)::int is NOT the same as `json_get_int(..)`, hence the cast is not rewritten
     let expected = [
-        "Projection: CAST(json_as_text(test.json_data, Utf8(\"foo\"), Int64(0)) AS json_data ->> 'foo' ->> 0 AS Int32)",
+        "Projection: CAST(json_as_text(test.json_data, Utf8(\"foo\"), Int64(0)) AS Int32) AS json_data ->> 'foo' ->> 0",
         "  TableScan: test projection=[json_data]",
     ];
 
