From b711454cd78cb51263170c0118fb2bd1420c38b4 Mon Sep 17 00:00:00 2001 From: Lennon Chin Date: Fri, 12 Sep 2025 16:09:11 +0800 Subject: [PATCH 1/3] [AUTHZ] select(*)/select(1) should check sub plans' privileges --- .../spark/authz/PrivilegesBuilder.scala | 5 +- .../authz/ranger/RangerSparkExtension.scala | 5 +- .../rule/RuleEliminateChildOutputHolder.scala | 31 ++++++ .../authz/rule/plan/ChildOutputHolder.scala | 31 ++++++ .../rule/plan/RuleChildOutputMarker.scala | 32 ++++++ .../ranger/RangerSparkExtensionSuite.scala | 104 +++++++++++++++++- 6 files changed, 201 insertions(+), 7 deletions(-) create mode 100644 extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/RuleEliminateChildOutputHolder.scala create mode 100644 extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/plan/ChildOutputHolder.scala create mode 100644 extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/plan/RuleChildOutputMarker.scala diff --git a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/PrivilegesBuilder.scala b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/PrivilegesBuilder.scala index 01266eb2c85..b0449b8b993 100644 --- a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/PrivilegesBuilder.scala +++ b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/PrivilegesBuilder.scala @@ -28,6 +28,7 @@ import org.slf4j.LoggerFactory import org.apache.kyuubi.plugin.spark.authz.OperationType.OperationType import org.apache.kyuubi.plugin.spark.authz.PrivilegeObjectActionType._ import org.apache.kyuubi.plugin.spark.authz.rule.Authorization._ +import org.apache.kyuubi.plugin.spark.authz.rule.plan.ChildOutputHolder import org.apache.kyuubi.plugin.spark.authz.rule.rowfilter._ import org.apache.kyuubi.plugin.spark.authz.serde._ import org.apache.kyuubi.plugin.spark.authz.util.AuthZUtils._ @@ -100,13 +101,15 @@ object PrivilegesBuilder { privilegeObjects += PrivilegeObject(table) case p => + val existsChildOutputHolder = p.exists(_.isInstanceOf[ChildOutputHolder]) for (child <- p.children) { // If current plan's references don't have relation to it's input, have two cases // 1. `MapInPandas`, `ScriptTransformation` // 2. `Project` output only have constant value if (columnPrune(p.references.toSeq ++ p.output, p.inputSet).isEmpty) { // If plan is project and output don't have relation to input, can ignore. - if (!p.isInstanceOf[Project]) { + // If plan tree exists ChildOutputHolder, we should build child logic plan. + if (!p.isInstanceOf[Project] || existsChildOutputHolder) { buildQuery( child, privilegeObjects, diff --git a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtension.scala b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtension.scala index 288719f07bf..8e61783b6a8 100644 --- a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtension.scala +++ b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtension.scala @@ -19,11 +19,12 @@ package org.apache.kyuubi.plugin.spark.authz.ranger import org.apache.spark.sql.SparkSessionExtensions -import org.apache.kyuubi.plugin.spark.authz.rule.{RuleEliminateMarker, RuleEliminatePermanentViewMarker, RuleEliminateTypeOf} +import org.apache.kyuubi.plugin.spark.authz.rule.{RuleEliminateChildOutputHolder, RuleEliminateMarker, RuleEliminatePermanentViewMarker, RuleEliminateTypeOf} import org.apache.kyuubi.plugin.spark.authz.rule.config.AuthzConfigurationChecker import org.apache.kyuubi.plugin.spark.authz.rule.datamasking.{RuleApplyDataMaskingStage0, RuleApplyDataMaskingStage1} import org.apache.kyuubi.plugin.spark.authz.rule.expression.RuleApplyTypeOfMarker import org.apache.kyuubi.plugin.spark.authz.rule.permanentview.RuleApplyPermanentViewMarker +import org.apache.kyuubi.plugin.spark.authz.rule.plan.RuleChildOutputMarker import org.apache.kyuubi.plugin.spark.authz.rule.rowfilter.{FilterDataSourceV2Strategy, RuleApplyRowFilter, RuleReplaceShowObjectCommands} /** @@ -48,6 +49,7 @@ class RangerSparkExtension extends (SparkSessionExtensions => Unit) { v1.injectResolutionRule(_ => RuleReplaceShowObjectCommands) v1.injectResolutionRule(_ => RuleApplyPermanentViewMarker) v1.injectResolutionRule(_ => RuleApplyTypeOfMarker) + v1.injectResolutionRule(_ => RuleChildOutputMarker) v1.injectResolutionRule(RuleApplyRowFilter) v1.injectResolutionRule(RuleApplyDataMaskingStage0) v1.injectResolutionRule(RuleApplyDataMaskingStage1) @@ -55,6 +57,7 @@ class RangerSparkExtension extends (SparkSessionExtensions => Unit) { v1.injectOptimizerRule(RuleAuthorization) v1.injectOptimizerRule(RuleEliminatePermanentViewMarker) v1.injectOptimizerRule(_ => RuleEliminateTypeOf) + v1.injectOptimizerRule(_ => RuleEliminateChildOutputHolder) v1.injectPlannerStrategy(FilterDataSourceV2Strategy) } } diff --git a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/RuleEliminateChildOutputHolder.scala b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/RuleEliminateChildOutputHolder.scala new file mode 100644 index 00000000000..aecd0d54322 --- /dev/null +++ b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/RuleEliminateChildOutputHolder.scala @@ -0,0 +1,31 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kyuubi.plugin.spark.authz.rule + +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan +import org.apache.spark.sql.catalyst.rules.Rule + +import org.apache.kyuubi.plugin.spark.authz.rule.plan.ChildOutputHolder + +object RuleEliminateChildOutputHolder extends Rule[LogicalPlan] { + override def apply(plan: LogicalPlan): LogicalPlan = { + plan.transform { + case p@ChildOutputHolder(child, _) => child + } + } +} diff --git a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/plan/ChildOutputHolder.scala b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/plan/ChildOutputHolder.scala new file mode 100644 index 00000000000..60e0c6f0700 --- /dev/null +++ b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/plan/ChildOutputHolder.scala @@ -0,0 +1,31 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kyuubi.plugin.spark.authz.rule.plan + +import org.apache.spark.sql.catalyst.expressions.Attribute +import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, UnaryNode} + +import org.apache.kyuubi.plugin.spark.authz.util.WithInternalChild + +case class ChildOutputHolder(child: LogicalPlan, fixedOutput: Seq[Attribute]) + extends UnaryNode with WithInternalChild { + + val output: Seq[Attribute] = fixedOutput + + override def withNewChildInternal(newChild: LogicalPlan): LogicalPlan = copy(child = newChild) +} diff --git a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/plan/RuleChildOutputMarker.scala b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/plan/RuleChildOutputMarker.scala new file mode 100644 index 00000000000..09c45c987db --- /dev/null +++ b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/plan/RuleChildOutputMarker.scala @@ -0,0 +1,32 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kyuubi.plugin.spark.authz.rule.plan + +import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan} +import org.apache.spark.sql.catalyst.rules.Rule + +object RuleChildOutputMarker extends Rule[LogicalPlan] { + override def apply(plan: LogicalPlan): LogicalPlan = { + plan.transform { + case a @ Aggregate(_, _, child) + if !child.isInstanceOf[ChildOutputHolder] && + child.outputSet.intersect(a.references).isEmpty => + a.copy(child = ChildOutputHolder(child, child.output)) + } + } +} diff --git a/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtensionSuite.scala b/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtensionSuite.scala index 1fdea0ed969..1c78e16ce9d 100644 --- a/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtensionSuite.scala +++ b/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtensionSuite.scala @@ -961,11 +961,36 @@ class HiveCatalogRangerSparkExtensionSuite extends RangerSparkExtensionSuite { |AS |SELECT count(*) as cnt, sum(id) as sum_id FROM $db1.$table1 """.stripMargin)) - checkAnswer(someone, s"SELECT count(*) FROM $db1.$table1", Row(0) :: Nil) - checkAnswer(someone, s"SELECT count(*) FROM $db1.$view1", Row(0) :: Nil) + checkAnswer(admin, s"SELECT count(*) FROM $db1.$table1", Row(0) :: Nil) - checkAnswer(someone, s"SELECT count(*) FROM $db1.$view2", Row(1) :: Nil) + checkAnswer(admin, s"SELECT count(*) FROM $db1.$view1", Row(0) :: Nil) + + checkAnswer(admin, s"SELECT count(*) FROM $db1.$view2", Row(1) :: Nil) + + interceptEndsWith[AccessControlException]( + doAs(someone, sql(s"SELECT count(*) FROM $db1.$table1").show()))( + s"does not have [select] privilege on [$db1/$table1/id,$db1/$table1/scope]") + + interceptEndsWith[AccessControlException]( + doAs(someone, sql(s"SELECT count(1) FROM $db1.$table1").show()))( + s"does not have [select] privilege on [$db1/$table1/id,$db1/$table1/scope]") + + interceptEndsWith[AccessControlException]( + doAs(someone, sql(s"SELECT count(*) FROM $db1.$view1").show()))( + s"does not have [select] privilege on [$db1/$view1/id,$db1/$view1/scope]") + + interceptEndsWith[AccessControlException]( + doAs(someone, sql(s"SELECT count(1) FROM $db1.$view1").show()))( + s"does not have [select] privilege on [$db1/$view1/id,$db1/$view1/scope]") + + interceptEndsWith[AccessControlException]( + doAs(someone, sql(s"SELECT count(*) FROM $db1.$view2").show()))( + s"does not have [select] privilege on [$db1/$view2/cnt,$db1/$view2/sum_id]") + + interceptEndsWith[AccessControlException]( + doAs(someone, sql(s"SELECT count(1) FROM $db1.$view2").show()))( + s"does not have [select] privilege on [$db1/$view2/cnt,$db1/$view2/sum_id]") interceptEndsWith[AccessControlException]( doAs(someone, sql(s"SELECT count(id) FROM $db1.$table1 WHERE id > 10").show()))( @@ -1500,13 +1525,33 @@ class HiveCatalogRangerSparkExtensionSuite extends RangerSparkExtensionSuite { doAs(admin, sql(s"CREATE TABLE IF NOT EXISTS $db1.$table1 (id int, scope int)")) doAs(admin, sql(s"CREATE VIEW $db1.$view1 AS SELECT * FROM $db1.$table1")) checkAnswer( - someone, + admin, s"SELECT count(*) FROM $db1.$table1 WHERE id > 1", Row(0) :: Nil) checkAnswer( - someone, + admin, s"SELECT count(*) FROM $db1.$view1 WHERE id > 1", Row(0) :: Nil) + interceptContains[AccessControlException]( + doAs( + someone, + sql(s"SELECT count(*) FROM $db1.$table1 WHERE id > 1").collect()))( + s"does not have [select] privilege on [$db1/$table1/id,$db1/$table1/scope]") + interceptContains[AccessControlException]( + doAs( + someone, + sql(s"SELECT count(1) FROM $db1.$table1 WHERE id > 1").collect()))( + s"does not have [select] privilege on [$db1/$table1/id,$db1/$table1/scope]") + interceptContains[AccessControlException]( + doAs( + someone, + sql(s"SELECT count(*) FROM $db1.$view1 WHERE id > 1").collect()))( + s"does not have [select] privilege on [$db1/$view1/id,$db1/$view1/scope]") + interceptContains[AccessControlException]( + doAs( + someone, + sql(s"SELECT count(1) FROM $db1.$view1 WHERE id > 1").collect()))( + s"does not have [select] privilege on [$db1/$view1/id,$db1/$view1/scope]") interceptContains[AccessControlException]( doAs( someone, @@ -1542,4 +1587,53 @@ class HiveCatalogRangerSparkExtensionSuite extends RangerSparkExtensionSuite { } } } + + test("select count(*)/count(1) should not ignore privileges") { + val db1 = defaultDb + val table1 = "table1" + withSingleCallEnabled { + withCleanTmpResources(Seq((s"$db1.$table1", "table"))) { + doAs( + admin, + sql( + s"""CREATE TABLE IF NOT EXISTS $db1.$table1 + |(id int, scope int, part string) + |PARTITIONED BY(part) + |""".stripMargin)) + + interceptContains[AccessControlException]( + doAs(someone, sql(s"select count(*) from $db1.$table1").show()))( + s"does not have [select] privilege on " + + s"[$db1/$table1/id,$db1/$table1/scope,$db1/$table1/part]") + + interceptContains[AccessControlException]( + doAs(someone, sql(s"select count(id) from $db1.$table1").show()))( + s"does not have [select] privilege on [$db1/$table1/id]") + + interceptContains[AccessControlException]( + doAs(someone, sql( + s"""select count(1) from ( + | select id, part from $db1.$table1 + |) t where part = 'part-1' + |""".stripMargin).show()))( + s"does not have [select] privilege on [$db1/$table1/id,$db1/$table1/part]") + + interceptContains[AccessControlException]( + doAs(someone, sql( + s"""select cnt from ( + | select count(1) as cnt from ( + | select id from $db1.$table1 where part = 'part-1' + | ) t1 + |) t2 + |""".stripMargin).show()))( + s"does not have [select] privilege on [$db1/$table1/id,$db1/$table1/part]") + + // df.count() same with select count(*) + interceptContains[AccessControlException]( + doAs(someone, sql(s"select id from $db1.$table1 where part = 'part-1'").toDF().count()))( + s"does not have [select] privilege on [$db1/$table1/id,$db1/$table1/part]") + } + + } + } } From b4d999275d5dad8057eac438dd4c0d8a2f508cc0 Mon Sep 17 00:00:00 2001 From: Lennon Chin Date: Fri, 12 Sep 2025 16:48:03 +0800 Subject: [PATCH 2/3] style fixed --- .../rule/RuleEliminateChildOutputHolder.scala | 2 +- .../ranger/RangerSparkExtensionSuite.scala | 28 +++++++++++-------- 2 files changed, 17 insertions(+), 13 deletions(-) diff --git a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/RuleEliminateChildOutputHolder.scala b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/RuleEliminateChildOutputHolder.scala index aecd0d54322..abd9702fdaa 100644 --- a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/RuleEliminateChildOutputHolder.scala +++ b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/RuleEliminateChildOutputHolder.scala @@ -25,7 +25,7 @@ import org.apache.kyuubi.plugin.spark.authz.rule.plan.ChildOutputHolder object RuleEliminateChildOutputHolder extends Rule[LogicalPlan] { override def apply(plan: LogicalPlan): LogicalPlan = { plan.transform { - case p@ChildOutputHolder(child, _) => child + case p @ ChildOutputHolder(child, _) => child } } } diff --git a/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtensionSuite.scala b/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtensionSuite.scala index 1c78e16ce9d..e378c40bfc9 100644 --- a/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtensionSuite.scala +++ b/extensions/spark/kyuubi-spark-authz/src/test/scala/org/apache/kyuubi/plugin/spark/authz/ranger/RangerSparkExtensionSuite.scala @@ -1611,21 +1611,25 @@ class HiveCatalogRangerSparkExtensionSuite extends RangerSparkExtensionSuite { s"does not have [select] privilege on [$db1/$table1/id]") interceptContains[AccessControlException]( - doAs(someone, sql( - s"""select count(1) from ( - | select id, part from $db1.$table1 - |) t where part = 'part-1' - |""".stripMargin).show()))( + doAs( + someone, + sql( + s"""select count(1) from ( + | select id, part from $db1.$table1 + |) t where part = 'part-1' + |""".stripMargin).show()))( s"does not have [select] privilege on [$db1/$table1/id,$db1/$table1/part]") interceptContains[AccessControlException]( - doAs(someone, sql( - s"""select cnt from ( - | select count(1) as cnt from ( - | select id from $db1.$table1 where part = 'part-1' - | ) t1 - |) t2 - |""".stripMargin).show()))( + doAs( + someone, + sql( + s"""select cnt from ( + | select count(1) as cnt from ( + | select id from $db1.$table1 where part = 'part-1' + | ) t1 + |) t2 + |""".stripMargin).show()))( s"does not have [select] privilege on [$db1/$table1/id,$db1/$table1/part]") // df.count() same with select count(*) From 952548f0770ee4512af076163499bdd7ee783412 Mon Sep 17 00:00:00 2001 From: Lennon Chin Date: Mon, 15 Sep 2025 09:35:41 +0800 Subject: [PATCH 3/3] optimize --- .../kyuubi/plugin/spark/authz/PrivilegesBuilder.scala | 7 +++++-- .../spark/authz/rule/RuleEliminateChildOutputHolder.scala | 3 +++ .../plugin/spark/authz/rule/plan/ChildOutputHolder.scala | 4 ++-- .../spark/authz/rule/plan/RuleChildOutputMarker.scala | 7 +++++++ 4 files changed, 17 insertions(+), 4 deletions(-) diff --git a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/PrivilegesBuilder.scala b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/PrivilegesBuilder.scala index b0449b8b993..4aaeaccd7c0 100644 --- a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/PrivilegesBuilder.scala +++ b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/PrivilegesBuilder.scala @@ -107,8 +107,11 @@ object PrivilegesBuilder { // 1. `MapInPandas`, `ScriptTransformation` // 2. `Project` output only have constant value if (columnPrune(p.references.toSeq ++ p.output, p.inputSet).isEmpty) { - // If plan is project and output don't have relation to input, can ignore. - // If plan tree exists ChildOutputHolder, we should build child logic plan. + // 1. If plan is project and output don't have relation to input, can ignore. + // 2. If sub logic plan tree exists ChildOutputHolder node, it means that the output of + // some nodes in the tree is fixed by RuleChildOutputMarker in some special + // scenarios, such as the Aggregate(count(*)) child node. To avoid missing child node + // permissions, we need to continue checking down. if (!p.isInstanceOf[Project] || existsChildOutputHolder) { buildQuery( child, diff --git a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/RuleEliminateChildOutputHolder.scala b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/RuleEliminateChildOutputHolder.scala index abd9702fdaa..faa95a7b40e 100644 --- a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/RuleEliminateChildOutputHolder.scala +++ b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/RuleEliminateChildOutputHolder.scala @@ -22,6 +22,9 @@ import org.apache.spark.sql.catalyst.rules.Rule import org.apache.kyuubi.plugin.spark.authz.rule.plan.ChildOutputHolder +/** + * Transforming down [[ChildOutputHolder]] + */ object RuleEliminateChildOutputHolder extends Rule[LogicalPlan] { override def apply(plan: LogicalPlan): LogicalPlan = { plan.transform { diff --git a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/plan/ChildOutputHolder.scala b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/plan/ChildOutputHolder.scala index 60e0c6f0700..739cde12ab2 100644 --- a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/plan/ChildOutputHolder.scala +++ b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/plan/ChildOutputHolder.scala @@ -22,10 +22,10 @@ import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, UnaryNode} import org.apache.kyuubi.plugin.spark.authz.util.WithInternalChild -case class ChildOutputHolder(child: LogicalPlan, fixedOutput: Seq[Attribute]) +case class ChildOutputHolder(child: LogicalPlan, childOutput: Seq[Attribute]) extends UnaryNode with WithInternalChild { - val output: Seq[Attribute] = fixedOutput + val output: Seq[Attribute] = childOutput override def withNewChildInternal(newChild: LogicalPlan): LogicalPlan = copy(child = newChild) } diff --git a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/plan/RuleChildOutputMarker.scala b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/plan/RuleChildOutputMarker.scala index 09c45c987db..a4a6bee28f7 100644 --- a/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/plan/RuleChildOutputMarker.scala +++ b/extensions/spark/kyuubi-spark-authz/src/main/scala/org/apache/kyuubi/plugin/spark/authz/rule/plan/RuleChildOutputMarker.scala @@ -20,6 +20,13 @@ package org.apache.kyuubi.plugin.spark.authz.rule.plan import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan} import org.apache.spark.sql.catalyst.rules.Rule +/** + * Rule to add [[ChildOutputHolder]] used to fixed logic plan's child node's output, + * now used for following cases: + * + * 1. Aggregate(count(*)/count(1)), it's child node will be pruned in Spark optimizer + * rule [[org.apache.spark.sql.catalyst.optimizer.ColumnPruning]]. + */ object RuleChildOutputMarker extends Rule[LogicalPlan] { override def apply(plan: LogicalPlan): LogicalPlan = { plan.transform {