kosiew commented on code in PR #26787:
URL: https://github.com/apache/datafusion/pull/26787#discussion_r4563203929
##########
datafusion/core/src/datasource/file_format/arrow.rs:
##########
@@ -96,4 +97,111 @@ mod tests {
Ok(())
}
+
+ /// Registers table `t`, whose batches are slices of one batch, so that the
+ /// views of each batch reference the string data of all the batches.
+ /// Returns the size of the string data.
+ fn register_sliced_views(ctx: &SessionContext) -> Result<usize> {
+ let num_rows = 40_000;
+ let strings: StringViewArray = (0..num_rows)
+ .map(|i| Some(format!("a string longer than twelve bytes {i:05}")))
+ .collect();
+ let data_size = strings.data_buffers().iter().map(|b| b.len()).sum();
+ let partitions = Int32Array::from_iter_values((0..num_rows).map(|i| i
% 4));
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("s", DataType::Utf8View, true),
+ Field::new("p", DataType::Int32, false),
+ ]));
+ let batch = RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(strings), Arc::new(partitions)],
+ )?;
+ let batches = (0..8).map(|i| batch.slice(i * 5000, 5000)).collect();
+ ctx.register_table("t", Arc::new(MemTable::try_new(schema,
vec![batches])?))?;
+ Ok(data_size)
+ }
+
+ /// Returns the total size of the files under `path`
+ fn total_file_size(path: &std::path::Path) -> std::io::Result<u64> {
+ let mut size = 0;
+ for entry in std::fs::read_dir(path)? {
+ let entry = entry?;
+ let metadata = entry.metadata()?;
+ size += if metadata.is_dir() {
+ total_file_size(&entry.path())?
+ } else {
+ metadata.len()
+ };
+ }
+ Ok(size)
+ }
+
+ #[tokio::test]
+ async fn test_write_arrow_compacts_view_buffers() -> Result<()> {
+ let ctx = SessionContext::new();
+ let data_size = register_sliced_views(&ctx)?;
+
+ let tmp_dir = tempfile::TempDir::new()?;
+ let path = format!("{}/views.arrow", tmp_dir.path().to_string_lossy());
+ // Uncompressed, so that the file size is the size of the written
buffers
+ ctx.sql(&format!(
+ "COPY t TO '{path}' STORED AS ARROW OPTIONS ('format.compression'
'uncompressed')"
+ ))
+ .await?
+ .collect()
+ .await?;
+
+ // Without compaction, each of the 8 batches holds all the string data
+ let file_size = std::fs::metadata(&path)?.len() as usize;
+ assert!(
+ file_size < 2 * data_size,
+ "wrote {file_size} bytes for {data_size} bytes of strings"
+ );
+
+ let rows = ctx
+ .read_arrow(&path, ArrowReadOptions::default())
+ .await?
+ .count()
+ .await?;
+ assert_eq!(rows, 40_000);
+
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn test_write_partitioned_arrow_compacts_view_buffers() ->
Result<()> {
+ let ctx = SessionContext::new();
+ let data_size = register_sliced_views(&ctx)?;
+
+ let tmp_dir = tempfile::TempDir::new()?;
Review Comment:
Could we strengthen both read-back checks by comparing decoded values as
multisets, without assuming row order? For partitioned output, compare each p=N
directory against its expected rows and check the physical schema without the
partition column. Size and row count alone cannot detect incorrect values or
partition routing. This is optional test hardening.
##########
datafusion/core/src/datasource/file_format/arrow.rs:
##########
@@ -96,4 +97,111 @@ mod tests {
Ok(())
}
+
+ /// Registers table `t`, whose batches are slices of one batch, so that the
+ /// views of each batch reference the string data of all the batches.
+ /// Returns the size of the string data.
+ fn register_sliced_views(ctx: &SessionContext) -> Result<usize> {
Review Comment:
Could we add a small COPY roundtrip test with sliced BinaryView values and
NULLs? Use out-of-line values with unique backing buffers exceeding 10 KiB and
slices that reference only part of that data, so compaction runs. Please
compare decoded values, not just row counts. This is optional integration
coverage since existing helper tests already cover BinaryView.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
