<commit_before>package mimir.optimizer;

import java.sql._;

import mimir.algebra._;
import mimir.ctables._;

object ProjectRedundantColumns {

  private def projectIfNeeded(oper: Operator, dependencies: Set[String]) = 
  {
    val existingColumns = oper.schema.map(_._1)
    if(existingColumns.forall( dependencies contains _ )){ oper }
    else {
      OperatorUtils.projectColumns(dependencies.toList, oper)
    }
  }

  def apply(o: Operator): Operator = 
    apply(o, o.schema.map(_._1).toSet)

  def apply(o: Operator, dependencies: Set[String]): Operator =
  {
    o match {
      case Project(args, source) => {
        val newArgs = 
          args.filter( arg => dependencies contains arg.name )

        val childDependencies = 
          newArgs.map( _.expression ).
            flatMap( ExpressionUtils.getColumns(_) ).
            toSet

        Project( newArgs, apply(source, childDependencies) )
      }

      case Select(condition, source) => {
        val childDependencies = 
          ExpressionUtils.getColumns( condition ) ++ dependencies

        Select(condition, apply(source, childDependencies))
      }

      case Union(lhs, rhs) => {
        Union(apply(lhs, dependencies), apply(rhs, dependencies))
      }

      case Join(lhs, rhs) => {
        val lhsDeps = lhs.schema.map(_._1).toSet & dependencies
        val rhsDeps = rhs.schema.map(_._1).toSet & dependencies
<commit_msg>ProjectRedundantColumns now respects Union schemas<commit_after>package mimir.optimizer;

import java.sql._;

import mimir.algebra._;
import mimir.ctables._;

object ProjectRedundantColumns {

  private def projectIfNeeded(oper: Operator, dependencies: Set[String]) = 
  {
    val existingColumns = oper.schema.map(_._1)
    if(existingColumns.forall( dependencies contains _ )){ oper }
    else {
      OperatorUtils.projectColumns(dependencies.toList, oper)
    }
  }

  def apply(o: Operator): Operator = 
    apply(o, o.schema.map(_._1).toSet)

  def apply(o: Operator, dependencies: Set[String]): Operator =
  {
    o match {
      case Project(args, source) => {
        val newArgs = 
          args.filter( arg => dependencies contains arg.name )

        val childDependencies = 
          newArgs.map( _.expression ).
            flatMap( ExpressionUtils.getColumns(_) ).
            toSet

        Project( newArgs, apply(source, childDependencies) )
      }

      case Select(condition, source) => {
        val childDependencies = 
          ExpressionUtils.getColumns( condition ) ++ dependencies

        Select(condition, apply(source, childDependencies))
      }

      case Union(lhs, rhs) => {
        // Naively, we could just push down the dependencies, but
        // in some cases we'll get back a superset of the desired 
        // schema, which would result in an invalid union
        // For example Union(Table(...), Project[...](Table(...)))
        //
        // We could introduce spurious projections around tables if
        // ABSOLUTELY needed... but let's instead see if it's really
        // critical

        val newLhs = apply(lhs, dependencies)
        val newRhs = apply(rhs, dependencies)
        val newLhsSchema = lhs.schema.map(_._1).toSet
        val newRhsSchema = rhs.schema.map(_._1).toSet

        if(   (!(dependencies -- newLhsSchema).isEmpty)
            ||(!(dependencies -- newRhsSchema).isEmpty))
        {
          if(!(newLhsSchema).equals(newRhsSchema)){
            Union(
              projectIfNeeded(newLhs, dependencies),
              projectIfNeeded(newRhs, dependencies)
            )
          } else {
            Union(newLhs,newRhs)
          }
        } else {
          Union(newLhs,newRhs)
        }
      }

      case Join(lhs, rhs) => {
        val lhsDeps = lhs.schema.map(_._1).toSet & dependencies
        val rhsDeps = rhs.schema.map(_._1).toSet & dependencies

{'fraction_non_alphanumeric': 0.11073651452282157, 'fraction_numerical': 0.0025933609958506223, 'mean_word_length': 2.3597560975609757, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7658484', 'n_tokens_mistral': 1140, 'n_tokens_neox': 1074, 'n_words': 259}