<commit_before>package org.deepdive.inference

import anorm._
import au.com.bytecode.opencsv.CSVWriter
import java.io.{ByteArrayInputStream, File, FileOutputStream, FileWriter,
  StringWriter, Reader, FileReader, InputStream, InputStreamReader}
import org.deepdive.settings.FactorFunctionVariable
import org.deepdive.datastore.PostgresDataStore
import org.deepdive.Logging
import org.deepdive.calibration._
import scala.collection.mutable.{ArrayBuffer, Set}
import scala.io.Source

/* Stores the factor graph and inference results in a postges database. */
trait PostgresInferenceDataStoreComponent extends InferenceDataStoreComponent {

  lazy val inferenceDataStore = new PostgresInferenceDataStore

  class PostgresInferenceDataStore extends InferenceDataStore with Logging {

    implicit lazy val connection = PostgresDataStore.borrowConnection()

    // Default batch size, if not overwritten by user
    val BatchSize = Some(250000)
    
    // Temorary buffer for the next batch.
    // These collections will the cleared when we write the next batch to postgres
    val variables = ArrayBuffer[Variable]()
    val factors = ArrayBuffer[Factor]()
    val weights = ArrayBuffer[Weight]()
    
    def init() : Unit = {
      
      variables.clear()
      factors.clear()
      weights.clear()

      // weights(id, initial_value, is_fixed, description)
      SQL("""drop table if exists weights CASCADE; 
        create table weights(id bigint primary key, 
        initial_value double precision, is_fixed boolean, description text);""").execute()
      
      // factors(id, weight_id, factor_function)
      SQL("""drop table if exists factors CASCADE; 
        create table factors(id bigint primary key, 
        weight_id bigint, factor_function text);""").execute()

      // variables(id, data_type, initial_value, is_evidence, is_query, mapping_relation, mapping_column)
      SQL("""drop table if exists variables CASCADE; 
        create table variables(id bigint primary key, data_type text,
        initial_value double precision, is_evidence boolean, is_query boolean,
        mapping_relation text, mapping_column text, mapping_id bigint);""").execute()
      
      // factor_variables(factor_id, variable_id, position, is_positive)
      SQL("""drop table if exists factor_variables; 
        create table factor_variables(factor_id bigint, variable_id bigint, 
        position int, is_positive boolean);""").execute()
      SQL("CREATE INDEX factor_idx ON factor_variables (factor_id);").execute()
      SQL("CREATE INDEX factor_variables_idx ON factor_variables (variable_id);").execute()
      SQL("analyze").execute()

      // inference_result(id, last_sample, probability)
      SQL("""drop table if exists inference_result CASCADE; 
        create table inference_result(id bigint primary key, last_sample boolean, 
        probability double precision);""").execute()

      // inference_result_weights(id, weight)
      SQL("""drop table if exists inference_result_weights CASCADE; 
        create table inference_result_weights(id bigint primary key, 
          weight double precision);""").execute()

      // A view for the mapped inference result.
      // The view is a join of the variables and inference result tables.
      SQL("""drop view if exists mapped_inference_result; 
      CREATE VIEW mapped_inference_result AS SELECT variables.*, inference_result.last_sample, inference_result.probability 
      FROM variables INNER JOIN inference_result ON variables.id = inference_result.id;
      """).execute()
    }

    def getLocalVariableIds(rowMap: Map[String, Any], factorVar: FactorFunctionVariable) : Array[Long] = {
      if (factorVar.isArray)
        // Postgres prefixes aggregated colimns with a dot
        rowMap(s".${factorVar.relation}.id").asInstanceOf[Array[java.lang.Long]].map(_.toLong)
      else
        Array(rowMap(s"${factorVar.relation}.id").asInstanceOf[Long])
    }

    def addFactor(factor: Factor) = { 
      factors += factor
    }

    def addVariable(variable: Variable) = {
      variables += variable
    }
    
    def addWeight(weight: Weight) = { 
      weights += weight
    }

    def dumpFactorGraph(serializer: Serializer) : Unit = {
      // Add all weights
      log.info(s"Dumping factor graph...")
      log.info("Serializing weights...")
      SQL(s"SELECT * from weights order by id asc")().iterator.foreach { row =>
        serializer.addWeight(
          row[Long]("id"),
          row[Boolean]("is_fixed"),
          row[Double]("initial_value"),
          row[String]("description")
        )
      }
      // Add all variables
      log.info("Serializing variables...")
      SQL(s"SELECT * from variables order by id asc")().iterator.foreach { row =>
        serializer.addVariable(
          row[Long]("id"),
          if (row[Boolean]("is_evidence")) row[Option[Double]]("initial_value") else None,
          row[String]("data_type")
        )
      }
      // Add all factors
      log.info("Serializing factors...")
      SQL(s"SELECT * from factors order by id asc")().iterator.foreach { row =>
        serializer.addFactor(
          row[Long]("id"),
          row[Long]("weight_id"),
          row[String]("factor_function")
        )
      }
      // Add all edges
      log.info("Serializing edges...")
      SQL(s"SELECT * from factor_variables")().iterator.foreach { row =>
        serializer.addEdge(
          row[Long]("variable_id"),
          row[Long]("factor_id"),
          row[Long]("position"),
          row[Boolean]("is_positive")
        )
      }
      log.info("Writing serialization result...")
      serializer.close()
    }

    def writebackInferenceResult(variableSchema: Map[String, String],
      variableOutputFile: String, weightsOutputFile: String) = {
      log.info("writing back inference result")

      // Copy the inference result back to the database
      PostgresDataStore.copyBatchData("COPY inference_result(id, last_sample, probability) FROM STDIN",
        Source.fromFile(variableOutputFile).reader())
      SQL("CREATE INDEX inference_result_idx ON inference_result using btree (probability);").execute()
      SQL("analyze").execute()

      // Each (relation, column) tuple is a variable in the plate model.
      // Find all (relation, column) combinations
      val relationsColumns = variableSchema.keys map (_.split('.')) map {
        case Array(relation, column) => (relation, column)
      } 

      // Generate a view for each (relation, column) combination.
      relationsColumns.foreach { case(relationName, columnName) => 
        val view_name = s"${relationName}_${columnName}_inference"
        log.info(s"creating view=${view_name}")
        SQL(s"""DROP VIEW IF EXISTS ${view_name}; CREATE VIEW ${view_name} AS
          SELECT ${relationName}.*, mir.last_sample, mir.probability FROM
          ${relationName} JOIN
            (SELECT mir.last_sample, mir.probability, mir.id, mir.mapping_id 
            FROM mapped_inference_result mir 
            WHERE mapping_relation = '${relationName}' AND mapping_column = '${columnName}'
            ORDER BY mir.probability DESC) 
          mir ON ${relationName}.id = mir.mapping_id""").execute()
      }

      // Copy the weight result back to the database
      PostgresDataStore.copyBatchData("COPY inference_result_weights(id, weight) FROM STDIN",
        Source.fromFile(weightsOutputFile).reader())
      SQL("""CREATE INDEX inference_result_weights_idx 
        ON inference_result_weights using btree (weight);""").execute()
      SQL("analyze").execute()
      
      // Create a view that maps weight descriptions to the weight values
      val mappedVeightsView = "inference_result_mapped_weights"
      log.info(s"creating view=${mappedVeightsView}")
      SQL(s"""create or replace view ${mappedVeightsView} AS
        SELECT weights.*, inference_result_weights.weight FROM
        weights JOIN inference_result_weights ON weights.id = inference_result_weights.id
        ORDER BY abs(weight) DESC""").execute()
      
    }

    def getCalibrationData(variable: String, buckets: List[Bucket]) : Map[Bucket, BucketData] = {
      val Array(relationName, columnName) = variable.split('.')
      val inference_view = s"${relationName}_${columnName}_inference"<commit_msg>Automatically create weight analytic view for each variable.<commit_after>package org.deepdive.inference

import anorm._
import au.com.bytecode.opencsv.CSVWriter
import java.io.{ByteArrayInputStream, File, FileOutputStream, FileWriter,
  StringWriter, Reader, FileReader, InputStream, InputStreamReader}
import org.deepdive.settings.FactorFunctionVariable
import org.deepdive.datastore.PostgresDataStore
import org.deepdive.Logging
import org.deepdive.calibration._
import scala.collection.mutable.{ArrayBuffer, Set}
import scala.io.Source

/* Stores the factor graph and inference results in a postges database. */
trait PostgresInferenceDataStoreComponent extends InferenceDataStoreComponent {

  lazy val inferenceDataStore = new PostgresInferenceDataStore

  class PostgresInferenceDataStore extends InferenceDataStore with Logging {

    implicit lazy val connection = PostgresDataStore.borrowConnection()

    // Default batch size, if not overwritten by user
    val BatchSize = Some(250000)
    
    // Temorary buffer for the next batch.
    // These collections will the cleared when we write the next batch to postgres
    val variables = ArrayBuffer[Variable]()
    val factors = ArrayBuffer[Factor]()
    val weights = ArrayBuffer[Weight]()
    
    def init() : Unit = {
      
      variables.clear()
      factors.clear()
      weights.clear()

      // weights(id, initial_value, is_fixed, description)
      SQL("""drop table if exists weights CASCADE; 
        create table weights(id bigint primary key, 
        initial_value double precision, is_fixed boolean, description text);""").execute()
      
      // factors(id, weight_id, factor_function)
      SQL("""drop table if exists factors CASCADE; 
        create table factors(id bigint primary key, 
        weight_id bigint, factor_function text);""").execute()

      // variables(id, data_type, initial_value, is_evidence, is_query, mapping_relation, mapping_column)
      SQL("""drop table if exists variables CASCADE; 
        create table variables(id bigint primary key, data_type text,
        initial_value double precision, is_evidence boolean, is_query boolean,
        mapping_relation text, mapping_column text, mapping_id bigint);""").execute()
      
      // factor_variables(factor_id, variable_id, position, is_positive)
      SQL("""drop table if exists factor_variables; 
        create table factor_variables(factor_id bigint, variable_id bigint, 
        position int, is_positive boolean);""").execute()
      SQL("CREATE INDEX factor_idx ON factor_variables (factor_id);").execute()
      SQL("CREATE INDEX factor_variables_idx ON factor_variables (variable_id);").execute()
      SQL("analyze").execute()

      // inference_result(id, last_sample, probability)
      SQL("""drop table if exists inference_result CASCADE; 
        create table inference_result(id bigint primary key, last_sample boolean, 
        probability double precision);""").execute()

      // inference_result_weights(id, weight)
      SQL("""drop table if exists inference_result_weights CASCADE; 
        create table inference_result_weights(id bigint primary key, 
          weight double precision);""").execute()

      // A view for the mapped inference result.
      // The view is a join of the variables and inference result tables.
      SQL("""drop view if exists mapped_inference_result; 
      CREATE VIEW mapped_inference_result AS SELECT variables.*, inference_result.last_sample, inference_result.probability 
      FROM variables INNER JOIN inference_result ON variables.id = inference_result.id;
      """).execute()
    }

    def getLocalVariableIds(rowMap: Map[String, Any], factorVar: FactorFunctionVariable) : Array[Long] = {
      if (factorVar.isArray)
        // Postgres prefixes aggregated colimns with a dot
        rowMap(s".${factorVar.relation}.id").asInstanceOf[Array[java.lang.Long]].map(_.toLong)
      else
        Array(rowMap(s"${factorVar.relation}.id").asInstanceOf[Long])
    }

    def addFactor(factor: Factor) = { 
      factors += factor
    }

    def addVariable(variable: Variable) = {
      variables += variable
    }
    
    def addWeight(weight: Weight) = { 
      weights += weight
    }

    def dumpFactorGraph(serializer: Serializer) : Unit = {
      // Add all weights
      log.info(s"Dumping factor graph...")
      log.info("Serializing weights...")
      SQL(s"SELECT * from weights order by id asc")().iterator.foreach { row =>
        serializer.addWeight(
          row[Long]("id"),
          row[Boolean]("is_fixed"),
          row[Double]("initial_value"),
          row[String]("description")
        )
      }
      // Add all variables
      log.info("Serializing variables...")
      SQL(s"SELECT * from variables order by id asc")().iterator.foreach { row =>
        serializer.addVariable(
          row[Long]("id"),
          if (row[Boolean]("is_evidence")) row[Option[Double]]("initial_value") else None,
          row[String]("data_type")
        )
      }
      // Add all factors
      log.info("Serializing factors...")
      SQL(s"SELECT * from factors order by id asc")().iterator.foreach { row =>
        serializer.addFactor(
          row[Long]("id"),
          row[Long]("weight_id"),
          row[String]("factor_function")
        )
      }
      // Add all edges
      log.info("Serializing edges...")
      SQL(s"SELECT * from factor_variables")().iterator.foreach { row =>
        serializer.addEdge(
          row[Long]("variable_id"),
          row[Long]("factor_id"),
          row[Long]("position"),
          row[Boolean]("is_positive")
        )
      }
      log.info("Writing serialization result...")
      serializer.close()
    }

    def writebackInferenceResult(variableSchema: Map[String, String],
      variableOutputFile: String, weightsOutputFile: String) = {
      log.info("writing back inference result")

      // Copy the inference result back to the database
      PostgresDataStore.copyBatchData("COPY inference_result(id, last_sample, probability) FROM STDIN",
        Source.fromFile(variableOutputFile).reader())
      SQL("CREATE INDEX inference_result_idx ON inference_result using btree (probability);").execute()
      SQL("analyze").execute()

      // Each (relation, column) tuple is a variable in the plate model.
      // Find all (relation, column) combinations
      val relationsColumns = variableSchema.keys map (_.split('.')) map {
        case Array(relation, column) => (relation, column)
      } 

      // Generate a view for each (relation, column) combination.
      relationsColumns.foreach { case(relationName, columnName) => 
        val view_name = s"${relationName}_${columnName}_inference"
        log.info(s"creating view=${view_name}")
        SQL(s"""DROP VIEW IF EXISTS ${view_name}; CREATE VIEW ${view_name} AS
          SELECT ${relationName}.*, mir.last_sample, mir.probability FROM
          ${relationName} JOIN
            (SELECT mir.last_sample, mir.probability, mir.id, mir.mapping_id 
            FROM mapped_inference_result mir 
            WHERE mapping_relation = '${relationName}' AND mapping_column = '${columnName}'
            ORDER BY mir.probability DESC) 
          mir ON ${relationName}.id = mir.mapping_id""").execute()
      }


      // Copy the weight result back to the database
      PostgresDataStore.copyBatchData("COPY inference_result_weights(id, weight) FROM STDIN",
        Source.fromFile(weightsOutputFile).reader())
      SQL("""CREATE INDEX inference_result_weights_idx 
        ON inference_result_weights using btree (weight);""").execute()
      SQL("analyze").execute()
      
      // Create a view that maps weight descriptions to the weight values
      val mappedVeightsView = "inference_result_mapped_weights"
      log.info(s"creating view=${mappedVeightsView}")
      SQL(s"""create or replace view ${mappedVeightsView} AS
        SELECT weights.*, inference_result_weights.weight FROM
        weights JOIN inference_result_weights ON weights.id = inference_result_weights.id
        ORDER BY abs(weight) DESC""").execute()

      relationsColumns.foreach { case(relationName, columnName) => 
        val view_name = s"${relationName}_${columnName}_weights"
        log.info(s"creating view=${view_name}")
        SQL(s"""CREATE OR REPLACE VIEW ${view_name} AS
            WITH top_weights AS (SELECT id
            FROM inference_result_mapped_weights 
            ORDER BY abs(weight) DESC),
          relevant_variables AS (SELECT top_weights.id, is_evidence, initial_value FROM top_weights 
            INNER JOIN factors ON factors.weight_id = top_weights.id 
            INNER JOIN factor_variables ON factors.id = factor_variables.factor_id 
            INNER JOIN variables ON factor_variables.variable_id = variables.id
            WHERE variables.mapping_relation = '${relationName}' 
              AND variables.mapping_column = '${columnName}'),
          grouped_weights AS (SELECT relevant_variables.id,
            COUNT(CASE WHEN relevant_variables.is_evidence=true 
              AND relevant_variables.initial_value=1.0 THEN 1 END) AS true_count,
            COUNT(CASE WHEN relevant_variables.is_evidence=true 
              AND relevant_variables.initial_value=0.0 THEN 1 END) AS false_count,
            COUNT(CASE WHEN relevant_variables.is_evidence=false THEN 1 END) AS total_count
            FROM relevant_variables GROUP BY (id))
          SELECT description, weight, true_count, false_count, total_count
            FROM grouped_weights INNER JOIN inference_result_mapped_weights 
              ON grouped_weights.id = inference_result_mapped_weights.id
            ORDER BY abs(weight) DESC""").execute()
      }
      
    }

    def getCalibrationData(variable: String, buckets: List[Bucket]) : Map[Bucket, BucketData] = {
      val Array(relationName, columnName) = variable.split('.')
      val inference_view = s"${relationName}_${columnName}_inference"
{'fraction_non_alphanumeric': 0.10359743744182226, 'fraction_numerical': 0.0010403548157476865, 'mean_word_length': 3.147138964577657, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 16, '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': '3406346', 'n_tokens_mistral': 5103, 'n_tokens_neox': 4871, 'n_words': 1310}