Skip to content

Commit 1d486d6

Browse files
affected rows count for insert (#18)
1 parent a9da1c2 commit 1d486d6

File tree

1 file changed

+12
-2
lines changed

1 file changed

+12
-2
lines changed

datafusion_iceberg/src/table.rs

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -602,14 +602,16 @@ impl DataSink for IcebergDataSink {
602602
)
603603
.await?;
604604

605+
let count = metadata_files.iter().map(|x|x.record_count()).fold(0, |acc, x| acc+ x);
606+
605607
table
606608
.new_transaction(self.0.branch.as_deref())
607609
.append(metadata_files)
608610
.commit()
609611
.await
610612
.map_err(DataFusionIcebergError::from)?;
611613

612-
Ok(0)
614+
Ok(count as u64)
613615
}
614616
fn metrics(&self) -> Option<MetricsSet> {
615617
None
@@ -953,7 +955,7 @@ mod tests {
953955

954956
ctx.register_table("orders", table.clone()).unwrap();
955957

956-
ctx.sql(
958+
let res = ctx.sql(
957959
"INSERT INTO orders (id, customer_id, product_id, date, amount) VALUES
958960
(1, 1, 1, '2020-01-01', 1),
959961
(2, 2, 1, '2020-01-01', 1),
@@ -968,6 +970,14 @@ mod tests {
968970
.await
969971
.expect("Failed to insert values into table");
970972

973+
let count = res[0]
974+
.column(0)
975+
.as_any()
976+
.downcast_ref::<Int64Array>()
977+
.unwrap()
978+
.values()[0];
979+
assert_eq!(count, 6);
980+
971981
let batches = ctx
972982
.sql("select product_id, sum(amount) from orders where customer_id = 1 group by product_id;")
973983
.await

0 commit comments

Comments
 (0)