Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
179 changes: 178 additions & 1 deletion crates/integrations/datafusion/tests/blob_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
mod common;

use arrow_array::{Array, BinaryArray, Int32Array, RecordBatch, StringArray};
use common::{create_sql_context, create_test_env, exec};
use common::{assert_sql_error, create_sql_context, create_test_env, exec};
use paimon::catalog::Identifier;
use paimon::spec::{BlobDescriptor, BlobViewStruct};
use paimon::table::BranchManager;
Expand Down Expand Up @@ -577,6 +577,66 @@ async fn test_blob_resolve_descriptor_with_offset() {
assert_eq!(rows[0], (1, "Partial".into(), Some(b"PAYLOAD".to_vec())));
}

#[tokio::test]
async fn test_blob_resolve_unknown_length_descriptor() {
let (tmp, sql_context) = setup(BLOB_TABLE_DDL).await;

let source_data = b"HEADER_PAYLOAD_TRAILER";
let source_path = tmp.path().join("blob_unknown_length.bin");
std::fs::write(&source_path, source_data).unwrap();

let uri = format!("file://{}", source_path.display());
let full_hex = to_hex(&BlobDescriptor::new(uri.clone(), 0, -1).serialize());
let suffix_hex = to_hex(&BlobDescriptor::new(uri.clone(), 7, -1).serialize());
let eof_hex =
to_hex(&BlobDescriptor::new(uri.clone(), source_data.len() as i64, -1).serialize());
let past_eof_hex =
to_hex(&BlobDescriptor::new(uri, source_data.len() as i64 + 5, -1).serialize());

let sql = format!(
"INSERT INTO paimon.test_db.t (id, name, picture) VALUES \
(1, 'Full', X'{full_hex}'), \
(2, 'Suffix', X'{suffix_hex}'), \
(3, 'Eof', X'{eof_hex}'), \
(4, 'PastEof', X'{past_eof_hex}'), \
(5, 'Raw', X'524157'), \
(6, 'Null', NULL)"
);
exec(&sql_context, &sql).await;

let rows = query_id_name_picture(
&sql_context,
"SELECT id, name, picture FROM paimon.test_db.t ORDER BY id",
)
.await;
assert_eq!(
rows,
vec![
(1, "Full".into(), Some(source_data.to_vec())),
(2, "Suffix".into(), Some(b"PAYLOAD_TRAILER".to_vec())),
(3, "Eof".into(), Some(Vec::new())),
(4, "PastEof".into(), Some(Vec::new())),
(5, "Raw".into(), Some(b"RAW".to_vec())),
(6, "Null".into(), None),
]
);
}

#[tokio::test]
async fn test_blob_descriptor_short_read_returns_error() {
let (tmp, sql_context) = setup(BLOB_TABLE_DDL).await;

let source_path = tmp.path().join("blob_short_read.bin");
std::fs::write(&source_path, b"short").unwrap();
let uri = format!("file://{}", source_path.display());
let desc_hex = to_hex(&BlobDescriptor::new(uri, 0, 6).serialize());
let sql = format!(
"INSERT INTO paimon.test_db.t (id, name, picture) VALUES (1, 'Short', X'{desc_hex}')"
);

assert_sql_error(&sql_context, &sql, "Failed to read BlobDescriptor").await;
}

/// Blob files roll independently when `blob.target-file-size` is small.
#[tokio::test]
async fn test_blob_rolling() {
Expand Down Expand Up @@ -709,6 +769,123 @@ async fn test_blob_descriptor_field_resolve_descriptor_value() {
);
}

#[tokio::test]
async fn test_blob_descriptor_field_resolve_unknown_length_descriptor() {
let (tmp, sql_context) = setup(
"CREATE TABLE paimon.test_db.t (\
id INT, \
name STRING, \
picture BLOB \
) WITH (\
'data-evolution.enabled' = 'true', \
'row-tracking.enabled' = 'true', \
'blob-descriptor-field' = 'picture'\
)",
)
.await;

let source_data = b"HEADER_PAYLOAD_TRAILER";
let source_path = tmp.path().join("descriptor_unknown_length.bin");
std::fs::write(&source_path, source_data).unwrap();

let uri = format!("file://{}", source_path.display());
let bounded_hex = to_hex(&BlobDescriptor::new(uri.clone(), 0, 6).serialize());
let suffix_hex = to_hex(&BlobDescriptor::new(uri.clone(), 7, -1).serialize());
let eof_hex =
to_hex(&BlobDescriptor::new(uri.clone(), source_data.len() as i64, -1).serialize());
let past_eof_hex =
to_hex(&BlobDescriptor::new(uri, source_data.len() as i64 + 5, -1).serialize());
let sql = format!(
"INSERT INTO paimon.test_db.t (id, name, picture) VALUES \
(1, 'Bounded', X'{bounded_hex}'), \
(2, 'Suffix', X'{suffix_hex}'), \
(3, 'Eof', X'{eof_hex}'), \
(4, 'PastEof', X'{past_eof_hex}'), \
(5, 'Raw', X'524157'), \
(6, 'Null', NULL)"
);
exec(&sql_context, &sql).await;

let rows = query_id_name_picture(
&sql_context,
"SELECT id, name, picture FROM paimon.test_db.t ORDER BY id",
)
.await;
assert_eq!(
rows,
vec![
(1, "Bounded".into(), Some(b"HEADER".to_vec())),
(2, "Suffix".into(), Some(b"PAYLOAD_TRAILER".to_vec())),
(3, "Eof".into(), Some(Vec::new())),
(4, "PastEof".into(), Some(Vec::new())),
(5, "Raw".into(), Some(b"RAW".to_vec())),
(6, "Null".into(), None),
]
);
}

#[tokio::test]
async fn test_blob_descriptor_field_rejects_invalid_length() {
let (tmp, sql_context) = setup(
"CREATE TABLE paimon.test_db.t (\
id INT, \
name STRING, \
picture BLOB \
) WITH (\
'data-evolution.enabled' = 'true', \
'row-tracking.enabled' = 'true', \
'blob-descriptor-field' = 'picture'\
)",
)
.await;

let uri = format!("file://{}", tmp.path().join("unused.bin").display());
let desc_hex = to_hex(&BlobDescriptor::new(uri, 0, -2).serialize());
let sql = format!(
"INSERT INTO paimon.test_db.t (id, name, picture) VALUES (1, 'Invalid', X'{desc_hex}')"
);
exec(&sql_context, &sql).await;

assert_sql_error(
&sql_context,
"SELECT id, name, picture FROM paimon.test_db.t",
"length must be -1 or non-negative",
)
.await;
}

#[tokio::test]
async fn test_blob_descriptor_field_short_read_returns_error() {
let (tmp, sql_context) = setup(
"CREATE TABLE paimon.test_db.t (\
id INT, \
name STRING, \
picture BLOB \
) WITH (\
'data-evolution.enabled' = 'true', \
'row-tracking.enabled' = 'true', \
'blob-descriptor-field' = 'picture'\
)",
)
.await;

let source_path = tmp.path().join("descriptor_short_read.bin");
std::fs::write(&source_path, b"short").unwrap();
let uri = format!("file://{}", source_path.display());
let desc_hex = to_hex(&BlobDescriptor::new(uri, 0, 6).serialize());
let sql = format!(
"INSERT INTO paimon.test_db.t (id, name, picture) VALUES (1, 'Short', X'{desc_hex}')"
);
exec(&sql_context, &sql).await;

assert_sql_error(
&sql_context,
"SELECT id, name, picture FROM paimon.test_db.t",
"Failed to read BlobDescriptor",
)
.await;
}

#[tokio::test]
async fn test_blob_descriptor_filter_before_resolve_skips_filtered_bad_descriptor() {
let (tmp, sql_context) = setup(
Expand Down
116 changes: 102 additions & 14 deletions crates/paimon/src/arrow/format/blob.rs
Original file line number Diff line number Diff line change
Expand Up @@ -617,6 +617,23 @@ impl BlobFormatWriter {

const BLOB_WRITE_BUFFER_SIZE: u64 = 8 * 1024 * 1024; // 8 MB

fn checked_blob_entry_length(payload_len: u64) -> crate::Result<i64> {
let entry_length = payload_len
.checked_add(BLOB_ENTRY_OVERHEAD)
.ok_or_else(|| Error::DataInvalid {
message: format!(
"Blob entry length overflows u64: payload_length={payload_len}, overhead={BLOB_ENTRY_OVERHEAD}"
),
source: None,
})?;
i64::try_from(entry_length).map_err(|e| Error::DataInvalid {
message: format!(
"Blob entry length exceeds i64: payload_length={payload_len}, entry_length={entry_length}"
),
source: Some(Box::new(e)),
})
}

#[async_trait]
impl FormatFileWriter for BlobFormatWriter {
async fn write(&mut self, batch: &RecordBatch) -> crate::Result<()> {
Expand All @@ -643,9 +660,7 @@ impl FormatFileWriter for BlobFormatWriter {

if BlobDescriptor::is_blob_descriptor(value) {
let desc = BlobDescriptor::deserialize(value)?;
let payload_len = desc.length() as u64;
let entry_length = (payload_len + BLOB_ENTRY_OVERHEAD) as i64;
self.lengths.push(entry_length);
let range = desc.range_spec()?;

let file_io = self.file_io.as_ref().ok_or_else(|| Error::DataInvalid {
message:
Expand All @@ -654,7 +669,47 @@ impl FormatFileWriter for BlobFormatWriter {
source: None,
})?;
let input = file_io.new_input(desc.uri())?;
let reader = input.reader().await?;
let offset = range.offset();
let payload_len = match range.length() {
Some(length) => length,
None => input
.metadata()
.await
.map_err(|e| Error::UnexpectedError {
message: format!(
"Failed to read metadata for BlobDescriptor '{}': {e}",
desc.uri()
),
source: Some(Box::new(e)),
})?
.size
.saturating_sub(offset),
};
let end = offset
.checked_add(payload_len)
.ok_or_else(|| Error::DataInvalid {
message: format!(
"BlobDescriptor range overflows u64: offset={offset}, length={payload_len}"
),
source: None,
})?;
let entry_length = checked_blob_entry_length(payload_len)?;
let entry_length_u64 = entry_length as u64;
let bytes_written = self
.bytes_written
.checked_add(entry_length_u64)
.ok_or_else(|| Error::DataInvalid {
message: format!(
"Blob file size overflows u64: current_size={}, entry_length={entry_length_u64}",
self.bytes_written
),
source: None,
})?;
let reader = if payload_len == 0 {
None
} else {
Some(input.reader().await?)
};

let mut hasher = crc32fast::Hasher::new();

Expand All @@ -664,15 +719,34 @@ impl FormatFileWriter for BlobFormatWriter {
.await?;

// Stream payload in chunks to avoid loading entire blob into memory
let start = desc.offset() as u64;
let end = start + payload_len;
let mut pos = start;
while pos < end {
let chunk_end = (pos + BLOB_WRITE_BUFFER_SIZE).min(end);
let chunk = reader.read(pos..chunk_end).await?;
hasher.update(&chunk);
self.writer.write(chunk).await?;
pos = chunk_end;
if let Some(reader) = reader.as_ref() {
let mut pos = offset;
while pos < end {
let chunk_end = pos.saturating_add(BLOB_WRITE_BUFFER_SIZE).min(end);
let chunk = reader.read(pos..chunk_end).await.map_err(|e| {
Error::UnexpectedError {
message: format!(
"Failed to read BlobDescriptor '{}' range {pos}..{chunk_end}: {e}",
desc.uri()
),
source: Some(Box::new(e)),
}
})?;
let actual_len = chunk.len() as u64;
let expected_len = chunk_end - pos;
if actual_len != expected_len {
return Err(Error::DataInvalid {
message: format!(
"Failed to read BlobDescriptor '{}': short read for range {pos}..{chunk_end}, expected={expected_len} bytes, actual={actual_len} bytes",
desc.uri()
),
source: None,
});
}
hasher.update(&chunk);
self.writer.write(chunk).await?;
pos = chunk_end;
}
}

let entry_length_bytes = entry_length.to_le_bytes();
Expand All @@ -685,7 +759,8 @@ impl FormatFileWriter for BlobFormatWriter {
.write(Bytes::copy_from_slice(&hasher.finalize().to_le_bytes()))
.await?;

self.bytes_written += entry_length as u64;
self.lengths.push(entry_length);
self.bytes_written = bytes_written;
} else {
let entry_length = (value.len() + BLOB_ENTRY_OVERHEAD as usize) as i64;
self.lengths.push(entry_length);
Expand Down Expand Up @@ -1050,6 +1125,19 @@ mod tests {
assert_eq!(decoded, values);
}

#[test]
fn test_checked_blob_entry_length() {
assert_eq!(
checked_blob_entry_length(0).unwrap(),
BLOB_ENTRY_OVERHEAD as i64
);

let max_payload = i64::MAX as u64 - BLOB_ENTRY_OVERHEAD;
assert_eq!(checked_blob_entry_length(max_payload).unwrap(), i64::MAX);
assert!(checked_blob_entry_length(max_payload + 1).is_err());
assert!(checked_blob_entry_length(u64::MAX).is_err());
}

fn basic_blob_rows() -> [Option<&'static [u8]>; 4] {
[
Some(&b"hello"[..]),
Expand Down
Loading
Loading