COW rewrite does not carry row lineage into replacement files on v3 tables #3401

#[cfg(test)]
mod repro_lineage {
    use std::collections::HashMap;
    use std::sync::{Arc, Mutex};

    use arrow_array::{ArrayRef, Int32Array, Int64Array, RecordBatch};
    use arrow_schema::{DataType, Field, Schema as ArrowSchema};
    use futures::TryStreamExt;
    use parquet::arrow::PARQUET_FIELD_ID_META_KEY;
    use tempfile::TempDir;

    use crate::cow_rewrite::{CowBatchRewrite, CowBatchRewriter, CowRewriteBuilder};
    use crate::io::LocalFsStorageFactory;
    use crate::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder};
    use crate::spec::{FormatVersion, NestedField, PrimitiveType, Schema, Type};
    use crate::table::Table;
    use crate::transaction::{ApplyTransactionAction, Transaction};
    use crate::{Catalog, CatalogBuilder, NamespaceIdent, Result, TableCreation};

    struct DropId2 {
        seen_columns: Mutex<Vec<String>>,
    }

    impl CowBatchRewriter for DropId2 {
        fn rewrite_batch(&self, batch: RecordBatch) -> Result<CowBatchRewrite> {
            *self.seen_columns.lock().unwrap() = batch
                .schema()
                .fields()
                .iter()
                .map(|f| f.name().clone())
                .collect();
            let ids = batch
                .column(0)
                .as_any()
                .downcast_ref::<Int32Array>()
                .unwrap();
            let keep: Vec<i32> = ids.values().iter().copied().filter(|v| *v != 2).collect();
            let changed = keep.len() != ids.len();
            let output = RecordBatch::try_new(batch.schema(), vec![
                Arc::new(Int32Array::from(keep)) as ArrayRef,
            ])
            .unwrap();
            Ok(CowBatchRewrite {
                output: Some(output),
                changed,
            })
        }
    }

    async fn read_lineage(table: &Table) -> Result<Vec<(i32, i64, i64)>> {
        let stream = table
            .scan()
            .select(["id", "_row_id", "_last_updated_sequence_number"])
            .build()?
            .to_arrow()
            .await?;
        let batches: Vec<RecordBatch> = stream.try_collect().await?;
        let mut rows = vec![];
        for b in batches {
            let id = b.column_by_name("id").unwrap().as_any().downcast_ref::<Int32Array>().unwrap();
            let rid = b.column_by_name("_row_id").unwrap().as_any().downcast_ref::<Int64Array>().unwrap();
            let lu = arrow_cast::cast(
                b.column_by_name("_last_updated_sequence_number").unwrap(),
                &DataType::Int64,
            )?;
            let lu = lu.as_any().downcast_ref::<Int64Array>().unwrap();
            for i in 0..b.num_rows() {
                rows.push((id.value(i), rid.value(i), lu.value(i)));
            }
        }
        rows.sort();
        Ok(rows)
    }

    #[tokio::test]
    async fn repro_cow_rewrite_v3_row_lineage() -> Result<()> {
        let temp_dir = TempDir::new().unwrap();
        let warehouse = format!("file://{}", temp_dir.path().join("warehouse").display());
        let catalog = MemoryCatalogBuilder::default()
            .with_storage_factory(Arc::new(LocalFsStorageFactory))
            .load("memory", HashMap::from([(MEMORY_CATALOG_WAREHOUSE.to_string(), warehouse)]))
            .await?;
        let namespace = NamespaceIdent::new("ns".to_string());
        catalog.create_namespace(&namespace, HashMap::new()).await?;

        let schema = Schema::builder()
            .with_fields(vec![
                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
            ])
            .build()?;
        let table = catalog
            .create_table(&namespace, TableCreation::builder()
                .name("v3_cow".to_string())
                .schema(schema)
                .format_version(FormatVersion::V3)
                .build())
            .await?;

        let arrow_schema = Arc::new(ArrowSchema::new(vec![
            Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
                PARQUET_FIELD_ID_META_KEY.to_string(),
                "1".to_string(),
            )])),
        ]));
        let batch = RecordBatch::try_new(arrow_schema, vec![
            Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef,
        ])?;
        let files = super::writer::write_replacement_batches(
            &table,
            table.metadata().current_schema().clone(),
            None,
            futures::stream::iter(vec![Ok(batch)]),
        )
        .await?;
        let tx = Transaction::new(&table);
        let tx = tx.fast_append().add_data_files(files).apply(tx)?;
        let table = tx.commit(&catalog).await?;

        let before = read_lineage(&table).await?;
        println!("BEFORE (id, _row_id, _last_updated_sequence_number): {before:?}");
        assert_eq!(before, vec![(1, 0, 1), (2, 1, 1), (3, 2, 1)]);

        let rewriter = Arc::new(DropId2 {
            seen_columns: Mutex::new(vec![]),
        });
        let result = CowRewriteBuilder::new(&table)
            .with_rewriter(rewriter.clone())
            .rewrite()
            .await?;
        println!("rewriter input columns: {:?}", rewriter.seen_columns.lock().unwrap());
        println!(
            "removed={} added={} added first_row_id={:?}",
            result.removed_data_files.len(),
            result.added_data_files.len(),
            result.added_data_files.iter().map(|f| f.first_row_id()).collect::<Vec<_>>()
        );
        assert_eq!(result.added_data_files.len(), 1);
        assert_eq!(result.added_data_files[0].first_row_id(), None);

        // Simulated commit: the replacement is appended while the source file is still live,
        // because no overwrite commit action exists yet.
        let tx = Transaction::new(&table);
        let tx = tx
            .fast_append()
            .add_data_files(result.added_data_files.clone())
            .apply(tx)?;
        let table = tx.commit(&catalog).await?;
        let after = read_lineage(&table).await?;
        println!("AFTER simulated commit: {after:?}");
        Ok(())
    }
}