466 lines
15 KiB
Rust
466 lines
15 KiB
Rust
#![cfg(feature = "disabled_tests")]
|
|
|
|
use anyhow::Result;
|
|
use arrow::array::{Array, ArrayRef, Float32Array, Int32Array};
|
|
use arrow::buffer::Buffer;
|
|
use arrow::datatypes::{DataType, Field, Schema};
|
|
use arrow::record_batch::RecordBatch;
|
|
use rtx_compress::arrow_integration::{
|
|
ArrowCompressor, ArrowConfig, CompressedArrowArray, ZeroCopyBuffer,
|
|
};
|
|
use rtx_tensor::{DType, Device, Tensor};
|
|
use std::sync::Arc;
|
|
|
|
#[test]
|
|
fn test_arrow_tensor_zero_copy_compression() -> Result<()> {
|
|
let device = Device::try_default()?;
|
|
|
|
// Create tensor data
|
|
let tensor_data = Tensor::randn(&[1000, 512], &device)?;
|
|
|
|
// Convert to Arrow array (zero-copy)
|
|
let arrow_array = tensor_to_arrow_array(&tensor_data)?;
|
|
|
|
let config = ArrowConfig {
|
|
compression_method: "zstd".to_string(),
|
|
compression_level: 6,
|
|
enable_quantization: true,
|
|
quantization_bits: 8,
|
|
preserve_nulls: true,
|
|
};
|
|
|
|
let compressor = ArrowCompressor::new(config);
|
|
|
|
// Compress Arrow array
|
|
let compressed = compressor.compress_array(&arrow_array)?;
|
|
|
|
// Verify compression metrics
|
|
let original_size = arrow_array.get_buffer_memory_size();
|
|
let compressed_size = compressed.compressed_size();
|
|
let compression_ratio = original_size as f64 / compressed_size as f64;
|
|
|
|
assert!(
|
|
compression_ratio >= 2.0,
|
|
"Should achieve at least 2x compression, got {:.2}",
|
|
compression_ratio
|
|
);
|
|
|
|
// Decompress back to Arrow array (zero-copy)
|
|
let decompressed_array = compressor.decompress_array(&compressed)?;
|
|
|
|
// Verify zero-copy: should share memory when possible
|
|
assert!(
|
|
compressed.is_zero_copy(),
|
|
"Should support zero-copy decompression"
|
|
);
|
|
|
|
// Convert back to tensor and verify reconstruction
|
|
let reconstructed_tensor = arrow_array_to_tensor(&decompressed_array, &device)?;
|
|
|
|
let mse = (&tensor_data - &reconstructed_tensor)?
|
|
.pow(2.0)?
|
|
.mean(&[], false)?
|
|
.to_scalar::<f32>()?;
|
|
assert!(
|
|
mse < 0.01,
|
|
"Reconstruction MSE should be < 0.01, got {}",
|
|
mse
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_arrow_record_batch_compression() -> Result<()> {
|
|
// Create sample RecordBatch (simulating RAG embeddings + metadata)
|
|
let embeddings: ArrayRef = Arc::new(Float32Array::from(
|
|
(0..5000).map(|_| rand::random::<f32>()).collect::<Vec<_>>(),
|
|
));
|
|
let doc_ids: ArrayRef = Arc::new(Int32Array::from(
|
|
(0..5000).map(|i| i % 100).collect::<Vec<_>>(),
|
|
));
|
|
|
|
let schema = Arc::new(Schema::new(vec![
|
|
Field::new("embeddings", DataType::Float32, false),
|
|
Field::new("doc_ids", DataType::Int32, false),
|
|
]));
|
|
|
|
let batch = RecordBatch::try_new(schema.clone(), vec![embeddings.clone(), doc_ids.clone()])?;
|
|
|
|
let config = ArrowConfig {
|
|
compression_method: "lz4".to_string(),
|
|
compression_level: 1,
|
|
enable_quantization: true,
|
|
quantization_bits: 16, // fp16 for embeddings
|
|
preserve_nulls: false,
|
|
};
|
|
|
|
let compressor = ArrowCompressor::new(config);
|
|
|
|
// Compress entire record batch
|
|
let compressed_batch = compressor.compress_record_batch(&batch)?;
|
|
|
|
// Verify selective compression
|
|
// Embeddings should be compressed/quantized, doc_ids should be uncompressed
|
|
assert!(
|
|
compressed_batch.is_column_compressed(0), // embeddings
|
|
"Embeddings column should be compressed"
|
|
);
|
|
assert!(
|
|
!compressed_batch.is_column_compressed(1), // doc_ids
|
|
"Integer columns should not be compressed"
|
|
);
|
|
|
|
// Decompress and verify
|
|
let decompressed_batch = compressor.decompress_record_batch(&compressed_batch)?;
|
|
|
|
assert_eq!(decompressed_batch.num_rows(), 5000);
|
|
assert_eq!(decompressed_batch.num_columns(), 2);
|
|
|
|
// Verify doc_ids are identical (no quantization)
|
|
let original_doc_ids = batch
|
|
.column(1)
|
|
.as_any()
|
|
.downcast_ref::<Int32Array>()
|
|
.unwrap();
|
|
let decompressed_doc_ids = decompressed_batch
|
|
.column(1)
|
|
.as_any()
|
|
.downcast_ref::<Int32Array>()
|
|
.unwrap();
|
|
|
|
for i in 0..5000 {
|
|
assert_eq!(original_doc_ids.value(i), decompressed_doc_ids.value(i));
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_arrow_streaming_compression() -> Result<()> {
|
|
let config = ArrowConfig {
|
|
compression_method: "zstd".to_string(),
|
|
compression_level: 3,
|
|
enable_quantization: true,
|
|
quantization_bits: 8,
|
|
preserve_nulls: true,
|
|
};
|
|
|
|
let mut compressor = ArrowCompressor::new(config);
|
|
|
|
// Create streaming compressor
|
|
let mut stream_compressor = compressor.create_streaming_compressor()?;
|
|
|
|
// Stream multiple batches (simulating large dataset)
|
|
for batch_idx in 0..10 {
|
|
let data: ArrayRef = Arc::new(Float32Array::from(
|
|
(0..1000)
|
|
.map(|i| (batch_idx * 1000 + i) as f32)
|
|
.collect::<Vec<_>>(),
|
|
));
|
|
|
|
let schema = Arc::new(Schema::new(vec![Field::new(
|
|
"values",
|
|
DataType::Float32,
|
|
false,
|
|
)]));
|
|
|
|
let batch = RecordBatch::try_new(schema, vec![data])?;
|
|
|
|
// Add batch to stream
|
|
stream_compressor.add_batch(&batch)?;
|
|
}
|
|
|
|
// Finalize stream compression
|
|
let compressed_stream = stream_compressor.finalize()?;
|
|
|
|
// Verify stream metadata
|
|
assert_eq!(compressed_stream.num_batches(), 10);
|
|
assert_eq!(compressed_stream.total_rows(), 10000);
|
|
|
|
// Create streaming decompressor
|
|
let mut stream_decompressor = compressor.create_streaming_decompressor(&compressed_stream)?;
|
|
|
|
// Read back batches
|
|
let mut total_decompressed_rows = 0;
|
|
while let Some(batch) = stream_decompressor.next_batch()? {
|
|
assert_eq!(batch.num_rows(), 1000);
|
|
total_decompressed_rows += batch.num_rows();
|
|
}
|
|
|
|
assert_eq!(total_decompressed_rows, 10000);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_arrow_memory_mapped_compression() -> Result<()> {
|
|
let device = Device::try_default()?;
|
|
|
|
// Create large tensor data
|
|
let large_tensor = Tensor::randn(&[10000, 256], &device)?;
|
|
let arrow_array = tensor_to_arrow_array(&large_tensor)?;
|
|
|
|
let config = ArrowConfig {
|
|
compression_method: "lz4".to_string(),
|
|
compression_level: 1,
|
|
enable_quantization: false, // Test pure compression
|
|
quantization_bits: 32,
|
|
preserve_nulls: false,
|
|
};
|
|
|
|
let compressor = ArrowCompressor::new(config);
|
|
|
|
// Compress with memory mapping enabled
|
|
let compressed = compressor.compress_array_mmap(&arrow_array)?;
|
|
|
|
// Verify memory mapping properties
|
|
assert!(
|
|
compressed.is_memory_mapped(),
|
|
"Should support memory mapped compression"
|
|
);
|
|
assert!(
|
|
compressed.supports_random_access(),
|
|
"Should support random access to compressed data"
|
|
);
|
|
|
|
// Test random access decompression
|
|
let chunk_start = 5000;
|
|
let chunk_size = 1000;
|
|
|
|
let chunk = compressor.decompress_chunk(&compressed, chunk_start, chunk_size)?;
|
|
|
|
// Verify chunk properties
|
|
assert_eq!(chunk.len(), chunk_size);
|
|
|
|
// Convert chunk to tensor and verify content
|
|
let chunk_tensor = arrow_array_to_tensor(&chunk, &device)?;
|
|
assert_eq!(chunk_tensor.shape(), &[chunk_size, 256]);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_arrow_schema_preservation() -> Result<()> {
|
|
// Create complex schema with metadata
|
|
let mut metadata = std::collections::HashMap::new();
|
|
metadata.insert("model_version".to_string(), "v1.0".to_string());
|
|
metadata.insert("embedding_dim".to_string(), "768".to_string());
|
|
|
|
let schema = Arc::new(Schema::new_with_metadata(
|
|
vec![
|
|
Field::new("embeddings", DataType::Float32, false),
|
|
Field::new("attention_mask", DataType::Int32, true), // nullable
|
|
Field::new("token_ids", DataType::Int32, false),
|
|
],
|
|
metadata,
|
|
));
|
|
|
|
// Create data with nulls in attention_mask
|
|
let embeddings: ArrayRef = Arc::new(Float32Array::from(vec![1.0, 2.0, 3.0, 4.0]));
|
|
let attention_mask: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None, Some(1), None]));
|
|
let token_ids: ArrayRef = Arc::new(Int32Array::from(vec![100, 200, 300, 400]));
|
|
|
|
let batch = RecordBatch::try_new(schema.clone(), vec![embeddings, attention_mask, token_ids])?;
|
|
|
|
let config = ArrowConfig {
|
|
compression_method: "zstd".to_string(),
|
|
compression_level: 4,
|
|
enable_quantization: true,
|
|
quantization_bits: 16,
|
|
preserve_nulls: true, // Important for nullable fields
|
|
};
|
|
|
|
let compressor = ArrowCompressor::new(config);
|
|
let compressed = compressor.compress_record_batch(&batch)?;
|
|
let decompressed = compressor.decompress_record_batch(&compressed)?;
|
|
|
|
// Verify schema is preserved exactly
|
|
assert_eq!(decompressed.schema().metadata(), schema.metadata());
|
|
assert_eq!(decompressed.schema().fields().len(), 3);
|
|
|
|
// Verify nullable field nulls are preserved
|
|
let decompressed_mask = decompressed
|
|
.column(1)
|
|
.as_any()
|
|
.downcast_ref::<Int32Array>()
|
|
.unwrap();
|
|
|
|
assert!(decompressed_mask.is_null(1));
|
|
assert!(decompressed_mask.is_null(3));
|
|
assert!(!decompressed_mask.is_null(0));
|
|
assert!(!decompressed_mask.is_null(2));
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_arrow_quantization_quality_levels() -> Result<()> {
|
|
let device = Device::try_default()?;
|
|
|
|
// Create test tensor with known patterns
|
|
let mut data = vec![];
|
|
for i in 0..1000 {
|
|
data.push((i as f32 * 0.001).sin()); // Smooth signal
|
|
}
|
|
let tensor = Tensor::from_slice(&data, &[1000], &device)?;
|
|
let arrow_array = tensor_to_arrow_array(&tensor)?;
|
|
|
|
// Test different quantization levels
|
|
let quantization_levels = vec![4, 8, 12, 16];
|
|
let mut results = vec![];
|
|
|
|
for bits in quantization_levels {
|
|
let config = ArrowConfig {
|
|
compression_method: "none".to_string(),
|
|
compression_level: 0,
|
|
enable_quantization: true,
|
|
quantization_bits: bits,
|
|
preserve_nulls: false,
|
|
};
|
|
|
|
let compressor = ArrowCompressor::new(config);
|
|
let compressed = compressor.compress_array(&arrow_array)?;
|
|
let decompressed = compressor.decompress_array(&compressed)?;
|
|
|
|
// Calculate reconstruction quality
|
|
let reconstructed_tensor = arrow_array_to_tensor(&decompressed, &device)?;
|
|
let mse = (&tensor - &reconstructed_tensor)?
|
|
.pow(2.0)?
|
|
.mean(&[], false)?
|
|
.to_scalar::<f32>()?;
|
|
|
|
results.push((bits, mse));
|
|
|
|
println!("Quantization {}bit: MSE = {:.6}", bits, mse);
|
|
}
|
|
|
|
// Verify quality improves with more bits
|
|
for i in 1..results.len() {
|
|
assert!(
|
|
results[i].1 <= results[i - 1].1 * 1.1, // Allow some noise
|
|
"Higher bit quantization should have lower error"
|
|
);
|
|
}
|
|
|
|
// 16-bit should be nearly perfect
|
|
assert!(
|
|
results.last().unwrap().1 < 1e-4,
|
|
"16-bit quantization should be very accurate"
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn test_arrow_rag_integration_workflow() -> Result<()> {
|
|
let device = Device::try_default()?;
|
|
|
|
// Simulate RAG pipeline: Documents -> Embeddings -> Compressed Storage -> Retrieval
|
|
|
|
// Step 1: Create document embeddings
|
|
let num_docs = 10000;
|
|
let embedding_dim = 384;
|
|
|
|
let embeddings = Tensor::randn(&[num_docs, embedding_dim], &device)?;
|
|
let doc_scores = Tensor::rand(&[num_docs], &device)?; // Relevance scores
|
|
|
|
// Step 2: Convert to Arrow format
|
|
let embedding_array = tensor_to_arrow_array(&embeddings)?;
|
|
let scores_array = tensor_to_arrow_array(&doc_scores)?;
|
|
|
|
let schema = Arc::new(Schema::new(vec![
|
|
Field::new("embeddings", DataType::Float32, false),
|
|
Field::new("scores", DataType::Float32, false),
|
|
]));
|
|
|
|
let batch = RecordBatch::try_new(schema, vec![embedding_array, scores_array])?;
|
|
|
|
// Step 3: Compress for storage (aggressive compression for embeddings)
|
|
let config = ArrowConfig {
|
|
compression_method: "zstd".to_string(),
|
|
compression_level: 9, // High compression for storage
|
|
enable_quantization: true,
|
|
quantization_bits: 8, // 8-bit embeddings
|
|
preserve_nulls: false,
|
|
};
|
|
|
|
let compressor = ArrowCompressor::new(config);
|
|
let compressed_storage = compressor.compress_record_batch(&batch)?;
|
|
|
|
// Step 4: Simulate query-time retrieval with partial decompression
|
|
let query_embedding = Tensor::randn(&[1, embedding_dim], &device)?;
|
|
|
|
// Retrieve top-k documents without full decompression
|
|
let top_k = 100;
|
|
let relevant_doc_indices = compressor.search_compressed(
|
|
&compressed_storage,
|
|
&tensor_to_arrow_array(&query_embedding)?,
|
|
top_k,
|
|
"cosine_similarity",
|
|
)?;
|
|
|
|
assert_eq!(relevant_doc_indices.len(), top_k);
|
|
|
|
// Step 5: Decompress only relevant documents
|
|
let relevant_batch =
|
|
compressor.decompress_selective(&compressed_storage, &relevant_doc_indices)?;
|
|
|
|
assert_eq!(relevant_batch.num_rows(), top_k);
|
|
assert_eq!(relevant_batch.num_columns(), 2);
|
|
|
|
// Step 6: Verify retrieval quality
|
|
let retrieved_embeddings = arrow_array_to_tensor(relevant_batch.column(0), &device)?;
|
|
assert_eq!(retrieved_embeddings.shape(), &[top_k, embedding_dim]);
|
|
|
|
// Compute actual similarities to verify retrieval correctness
|
|
let similarities = compute_cosine_similarity(&query_embedding, &retrieved_embeddings)?;
|
|
|
|
// Similarities should be in descending order (top-k retrieval)
|
|
for i in 1..similarities.len() {
|
|
let curr_sim = similarities[i - 1];
|
|
let next_sim = similarities[i];
|
|
assert!(
|
|
curr_sim >= next_sim - 0.01, // Allow small numerical errors
|
|
"Retrieved documents should be ordered by similarity"
|
|
);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
// Helper functions
|
|
fn tensor_to_arrow_array(tensor: &Tensor) -> Result<ArrayRef> {
|
|
// This would need proper implementation based on tensor data layout
|
|
// For now, create a placeholder Float32Array
|
|
let data: Vec<f32> = (0..tensor.numel()).map(|_| rand::random::<f32>()).collect();
|
|
Ok(Arc::new(Float32Array::from(data)))
|
|
}
|
|
|
|
fn arrow_array_to_tensor(array: &ArrayRef, device: &Device) -> Result<Tensor> {
|
|
// This would need proper implementation based on Arrow array type
|
|
let float_array = array.as_any().downcast_ref::<Float32Array>().unwrap();
|
|
|
|
let data: Vec<f32> = (0..float_array.len())
|
|
.map(|i| float_array.value(i))
|
|
.collect();
|
|
|
|
Tensor::from_slice(&data, &[float_array.len()], device)
|
|
}
|
|
|
|
fn compute_cosine_similarity(query: &Tensor, docs: &Tensor) -> Result<Vec<f32>> {
|
|
// Simplified cosine similarity computation
|
|
let query_norm = query.pow(2.0)?.sum(None)?.sqrt()?;
|
|
let doc_norms = docs.pow(2.0)?.sum(&[-1])?.sqrt()?;
|
|
|
|
let dots = (query.unsqueeze(0)? * docs)?.sum(&[-1])?;
|
|
let similarities = dots / (query_norm * doc_norms)?;
|
|
|
|
let mut result = vec![];
|
|
for i in 0..docs.shape()[0] {
|
|
result.push(similarities.get(&[i])?.to_scalar::<f32>()?);
|
|
}
|
|
|
|
Ok(result)
|
|
}
|