diff --git a/src/main/scala/org/grapheco/lynx/dataframe/DataFrame.scala b/src/main/scala/org/grapheco/lynx/dataframe/DataFrame.scala index b7ea3166..602d5174 100644 --- a/src/main/scala/org/grapheco/lynx/dataframe/DataFrame.scala +++ b/src/main/scala/org/grapheco/lynx/dataframe/DataFrame.scala @@ -17,9 +17,13 @@ object DataFrame { def apply(schema0: Seq[(String, LynxType)], records0: () => Iterator[Seq[LynxValue]]): DataFrame = new DataFrame { - override def schema: Seq[(String, LynxType)] = schema0 + // 使用 lazy val 对 schema 和 records 进行懒加载和缓存,避免重复计算。 + private lazy val cachedSchema = schema0 + private lazy val cachedRecords = records0() - override def records: Iterator[Seq[LynxValue]] = records0() + override def schema: Seq[(String, LynxType)] = cachedSchema + + override def records: Iterator[Seq[LynxValue]] = cachedRecords } def cached(schema0: Seq[(String, LynxType)], records: Seq[Seq[LynxValue]]): DataFrame = @@ -32,8 +36,14 @@ object DataFrame { DataFrame(schema, () => Iterator.single( columns.map(col => { - expressionEvaluator.eval(col._2)(ctx) - }))) + // 增加错误处理机制,捕获并抛出异常,提供有用的错误信息。 + try { + expressionEvaluator.eval(col._2)(ctx) + } catch { + case e: Exception => throw new RuntimeException(s"Failed to evaluate expression for column ${col._1}", e) + } + }) + )) } def updateColumns(colIndexs: Seq[Int], newColsValues: Seq[Iterator[LynxValue]], srcDF: DataFrame): DataFrame = { diff --git a/src/main/scala/org/grapheco/lynx/optimizer/ApplyPushDownRule.scala b/src/main/scala/org/grapheco/lynx/optimizer/ApplyPushDownRule.scala index 2ec7fb0e..b4dd534e 100644 --- a/src/main/scala/org/grapheco/lynx/optimizer/ApplyPushDownRule.scala +++ b/src/main/scala/org/grapheco/lynx/optimizer/ApplyPushDownRule.scala @@ -19,7 +19,45 @@ object ApplyPushDownRule extends PhysicalPlanOptimizerRule { ║ [A] */ - case apply: Apply => apply + case apply: Apply => + val A = apply.left.get + val B = apply.right.get + + // 递归检查右子树是否只包含 FromArgument + def containsOnlyFromArgument(plan: PhysicalPlan): Boolean = plan match { + case _: FromArgument => true + case p if p.children.nonEmpty => + p.children.length == 1 && containsOnlyFromArgument(p.children.head) + case _ => false + } + + // 如果右子树只包含 FromArgument,则重组查询计划 + if (containsOnlyFromArgument(B)) { + // 获取右子树中除 FromArgument 外的所有操作 + def extractOperations(plan: PhysicalPlan): Option[PhysicalPlan] = plan match { + case _: FromArgument => None + case p if p.children.nonEmpty => + val childOp = extractOperations(p.children.head) + childOp match { + case Some(child) => + p.withChildren(Some(child), None) + Some(p) + case None => Some(p) + } + case _ => Some(plan) + } + + // 重组查询计划 + val operations = extractOperations(B) + operations match { + case Some(ops) => + ops.withChildren(Some(A), None) + ops + case None => A + } + } else { + apply + } } private val APPLY_PUSH_DOWN: PartialFunction[PhysicalPlan, PhysicalPlan] = { diff --git a/src/main/scala/org/grapheco/lynx/optimizer/RemoveApplyRule.scala b/src/main/scala/org/grapheco/lynx/optimizer/RemoveApplyRule.scala index 96e71ad4..227aa847 100644 --- a/src/main/scala/org/grapheco/lynx/optimizer/RemoveApplyRule.scala +++ b/src/main/scala/org/grapheco/lynx/optimizer/RemoveApplyRule.scala @@ -5,8 +5,7 @@ import org.grapheco.lynx.physical.plans.{Apply, FromArgument, Unwind, PhysicalPl object RemoveApplyRule extends PhysicalPlanOptimizerRule { override def apply(plan: PhysicalPlan, ppc: PhysicalPlannerContext): PhysicalPlan = optimizeBottomUp(plan, { - case apply: Apply => apply.right match { - /* Case 1: Combine Unwind, eg: +/* Case 1: Combine Unwind, eg: Apply ╟──[A] ╙──[Unwind] @@ -14,14 +13,15 @@ object RemoveApplyRule extends PhysicalPlanOptimizerRule { [Unwind] ║ [A] - */ - case Some(uw:Unwind) => { - uw.withChildren(apply.left) +*/ + case apply: Apply => + // attempt Case 1: Unwind optimization + val afterUnwindOptimize = apply.right match { + case Some(uw: Unwind) => uw.withChildren(apply.left) + case _ => apply } - case _ => apply - } - /* Case 2: Combine FromArgument, eg: +/* Case 2: Combine FromArgument, eg: Apply ╟──[A] ╙──[B*]──[FromArgument] @@ -29,9 +29,42 @@ object RemoveApplyRule extends PhysicalPlanOptimizerRule { [B*] ║ [A] - */ -// case apply: Apply => apply.right.get.leaves match { -// case Seq(FromArgument(variableName)) => apply.right -// } +*/ + + // 如果 Unwind 优化没有生效,尝试 Case 2: FromArgument 优化 + afterUnwindOptimize match { + case a: Apply => + a.right match { + case Some(rightPlan) => + // 获取右子树的所有叶子节点 + val leaves = collectLeaves(rightPlan) + // 检查是否只包含 FromArgument + if (leaves.forall(_.isInstanceOf[FromArgument])) { + // 获取右子树中除 FromArgument 外的所有操作 + val operations = removeFromArgument(rightPlan) + operations match { + case Some(ops) => ops.withChildren(a.left, None) + case None => a.left.getOrElse(a) + } + } else a + case None => a + } + case other => other + } }) + + // 收集所有叶子节点 + private def collectLeaves(plan: PhysicalPlan): Seq[PhysicalPlan] = { + if (plan.children.isEmpty) Seq(plan) + else plan.children.flatMap(p => collectLeaves(p)) + } + + // 移除 FromArgument 并重构查询计划 + private def removeFromArgument(plan: PhysicalPlan): Option[PhysicalPlan] = plan match { + case _: FromArgument => None + case p => + val newChildren = p.children.flatMap(child => removeFromArgument(child)) + if (newChildren.isEmpty) Some(p.withChildren()) + else Some(p.withChildren(newChildren.map(Some(_)): _*)) + } }