Compare based on keys on hash diff

Author: lewazo

src/lib/src/api/local/compare.rs

@@ -46,7 +46,7 @@ pub fn compare_files(
     targets: Vec<String>,
     output: Option<PathBuf>,
 ) -> Result<CompareResult, OxenError> {
-    // Assert that the files exist in their respective commits and are tabular.
+    // Assert that the files exist in their respective commits.
     let file_1 = get_version_file(repo, &compare_entry_1)?;
     let file_2 = get_version_file(repo, &compare_entry_2)?;
 
@@ -414,7 +414,9 @@ fn compute_row_comparison(
     let mut dupes = CompareDupes { left: 0, right: 0 };
 
     let dataframes = match strategy {
-        CompareStrategy::Hash => hash_compare::compare(df_1, df_2, &schema_1, &schema_2)?,
+        CompareStrategy::Hash => {
+            hash_compare::compare(df_1, df_2, &schema_1, &schema_2, keys.to_owned())?
+        }
         CompareStrategy::Join => {
             // TODO: unsure if hash comparison or join is faster here - would guess join, could use some testing
             let (df_1, df_2) = hash_dfs(

src/lib/src/api/local/compare/hash_compare.rs

@@ -13,6 +13,7 @@ pub fn compare(
     head_df: &DataFrame,
     schema_1: &Schema,
     schema_2: &Schema,
+    keys: Vec<&str>,
 ) -> Result<(DataFrame, DataFrame, DataFrame, DataFrame), OxenError> {
     if schema_1.hash != schema_2.hash {
         return Err(OxenError::invalid_file_type(
@@ -21,7 +22,7 @@ pub fn compare(
     }
 
     // Compute row indices
-    let (added_indices, removed_indices) = compute_new_row_indices(base_df, head_df)?;
+    let (added_indices, removed_indices) = compute_new_row_indices(base_df, head_df, keys)?;
 
     // Take added from the current df
     let added_rows = if !added_indices.is_empty() {
@@ -52,10 +53,11 @@ pub fn compare(
 fn compute_new_row_indices(
     base_df: &DataFrame,
     head_df: &DataFrame,
+    keys: Vec<&str>,
 ) -> Result<(Vec<u32>, Vec<u32>), OxenError> {
     // Hash the rows
-    let base_df = tabular::df_hash_rows(base_df.clone())?;
-    let head_df = tabular::df_hash_rows(head_df.clone())?;
+    let base_df = tabular::df_hash_rows(base_df.clone(), Some(keys.clone()))?;
+    let head_df = tabular::df_hash_rows(head_df.clone(), Some(keys.clone()))?;
 
     log::debug!("diff_current got current hashes base_df {:?}", base_df);
     log::debug!("diff_current got current hashes head_df {:?}", head_df);

src/lib/src/api/local/diff.rs

@@ -3,21 +3,19 @@ use serde::{Deserialize, Serialize};
 use crate::core::df::tabular;
 use crate::core::index::object_db_reader::ObjectDBReader;
 use crate::core::index::CommitDirEntryReader;
+use crate::core::index::CommitEntryReader;
 use crate::error::OxenError;
 use crate::model::diff::diff_entry_status::DiffEntryStatus;
-use crate::model::diff::generic_diff::GenericDiff;
 use crate::model::{Commit, CommitEntry, DataFrameDiff, DiffEntry, LocalRepository, Schema};
 use crate::opts::DFOpts;
 use crate::view::compare::AddRemoveModifyCounts;
 use crate::view::Pagination;
 use crate::{constants, util};
 
-use crate::core::index::CommitEntryReader;
-use colored::Colorize;
-use difference::{Changeset, Difference};
 use polars::export::ahash::HashMap;
 use polars::prelude::DataFrame;
 use polars::prelude::IntoLazy;
+
 use std::collections::HashSet;
 use std::path::{Path, PathBuf};
 use std::str::FromStr;
@@ -79,8 +77,8 @@ pub fn get_version_file_from_commit_id(
 
 pub fn count_added_rows(base_df: DataFrame, head_df: DataFrame) -> Result<usize, OxenError> {
     // Hash the rows
-    let base_df = tabular::df_hash_rows(base_df)?;
-    let head_df = tabular::df_hash_rows(head_df)?;
+    let base_df = tabular::df_hash_rows(base_df, None)?;
+    let head_df = tabular::df_hash_rows(head_df, None)?;
 
     // log::debug!("count_added_rows got base_df {}", base_df);
     // log::debug!("count_added_rows got head_df {}", head_df);
@@ -111,8 +109,8 @@ pub fn count_added_rows(base_df: DataFrame, head_df: DataFrame) -> Result<usize,
 
 pub fn count_removed_rows(base_df: DataFrame, head_df: DataFrame) -> Result<usize, OxenError> {
     // Hash the rows
-    let base_df = tabular::df_hash_rows(base_df)?;
-    let head_df = tabular::df_hash_rows(head_df)?;
+    let base_df = tabular::df_hash_rows(base_df, None)?;
+    let head_df = tabular::df_hash_rows(head_df, None)?;
 
     // log::debug!("count_removed_rows got base_df {}", base_df);
     // log::debug!("count_removed_rows got head_df {}", head_df);
@@ -149,8 +147,8 @@ pub fn compute_new_row_indices(
     head_df: &DataFrame,
 ) -> Result<(Vec<u32>, Vec<u32>), OxenError> {
     // Hash the rows
-    let base_df = tabular::df_hash_rows(base_df.clone())?;
-    let head_df = tabular::df_hash_rows(head_df.clone())?;
+    let base_df = tabular::df_hash_rows(base_df.clone(), None)?;
+    let head_df = tabular::df_hash_rows(head_df.clone(), None)?;
 
     log::debug!("diff_current got current hashes base_df {:?}", base_df);
     log::debug!("diff_current got current hashes head_df {:?}", head_df);

src/lib/src/core/df/tabular.rs

@@ -700,13 +700,23 @@ pub fn any_val_to_bytes(value: &AnyValue) -> Vec<u8> {
     }
 }
 
-pub fn df_hash_rows(df: DataFrame) -> Result<DataFrame, OxenError> {
+pub fn df_hash_rows(df: DataFrame, keys: Option<Vec<&str>>) -> Result<DataFrame, OxenError> {
     let num_rows = df.height() as i64;
 
     let mut col_names = vec![];
     let schema = df.schema();
-    for field in schema.iter_fields() {
-        col_names.push(col(field.name()));
+
+    match keys {
+        Some(keys) => {
+            for key in keys {
+                col_names.push(col(key));
+            }
+        }
+        None => {
+            for field in schema.iter_fields() {
+                col_names.push(col(field.name()));
+            }
+        }
     }
     // println!("Hashing: {:?}", col_names);
     // println!("{:?}", df);