Kontinuation commented on code in PR #521:
URL: https://github.com/apache/sedona-db/pull/521#discussion_r2697251934
##########
rust/sedona-spatial-join/src/utils/arrow_utils.rs:
##########
@@ -15,10 +15,149 @@
// specific language governing permissions and limitations
// under the License.
-use arrow::array::{Array, ArrayData, RecordBatch};
+use std::sync::Arc;
+
+use arrow::array::{Array, ArrayData, BinaryViewArray, ListArray, RecordBatch,
StringViewArray};
+use arrow_array::make_array;
use arrow_array::ArrayRef;
+use arrow_array::StructArray;
use arrow_schema::{ArrowError, DataType};
use datafusion_common::Result;
+use sedona_common::sedona_internal_err;
+
+/// Reconstruct `batch` to organize the payload buffers of each
`StringViewArray` and
+/// `BinaryViewArray` in sequential order by calling `gc()` on them.
+///
+/// Note this is a workaround until
<https://github.com/apache/arrow-rs/issues/7185> is
+/// available.
+///
+/// # Rationale
+///
+/// The `interleave` kernel does not reconstruct the inner buffers of view
arrays by default,
+/// leading to non-sequential payload locations. A single payload buffer might
be shared by
+/// multiple `RecordBatch`es or multiple rows in the same batch might
reference scattered
+/// locations in a large buffer.
+///
+/// When writing each batch to disk, the writer has to write all referenced
buffers. This
+/// causes extra disk reads and writes, and potentially execution failure
(e.g. No space left
+/// on device).
+///
+/// # Example
+///
+/// Before interleaving:
+/// batch1 -> buffer1 (large)
+/// batch2 -> buffer2 (large)
+///
+/// interleaved_batch -> buffer1 (sparse access)
+/// -> buffer2 (sparse access)
+///
+/// Then when spilling the interleaved batch, the writer has to write both
buffer1 and buffer2
+/// entirely, even if only a few bytes are used.
+pub(crate) fn compact_batch(batch: RecordBatch) -> Result<RecordBatch> {
+ let mut new_columns: Vec<Arc<dyn Array>> =
Vec::with_capacity(batch.num_columns());
+ let mut arr_mutated = false;
+
+ for array in batch.columns() {
+ let (new_array, mutated) = compact_array(Arc::clone(array))?;
+ new_columns.push(new_array);
+ arr_mutated |= mutated;
+ }
+
+ if arr_mutated {
+ Ok(RecordBatch::try_new(batch.schema(), new_columns)?)
+ } else {
+ Ok(batch)
+ }
+}
+
+/// Recursively compacts view arrays in `array` by calling `gc()` on them.
+/// Returns a tuple of the potentially new array and a boolean indicating
+/// whether any compaction was performed.
+pub(crate) fn compact_array(array: ArrayRef) -> Result<(ArrayRef, bool)> {
+ if let Some(view_array) = array.as_any().downcast_ref::<StringViewArray>()
{
+ return Ok((Arc::new(view_array.gc()), true));
+ }
+ if let Some(view_array) = array.as_any().downcast_ref::<BinaryViewArray>()
{
+ return Ok((Arc::new(view_array.gc()), true));
+ }
+
+ // Fast path for non-nested arrays
+ if !array.data_type().is_nested() {
+ return Ok((array, false));
+ }
+
+ // Avoid ArrayData -> ArrayRef roundtrips for commonly used data types,
+ // including StructArray and ListArray.
+
+ if let Some(struct_array) = array.as_any().downcast_ref::<StructArray>() {
+ let mut mutated = false;
+ let mut new_columns: Vec<ArrayRef> =
Vec::with_capacity(struct_array.num_columns());
+ for col in struct_array.columns() {
+ let (new_col, col_mutated) = compact_array(Arc::clone(col))?;
+ mutated |= col_mutated;
+ new_columns.push(new_col);
+ }
+
+ if !mutated {
+ return Ok((array, false));
+ }
+
+ let rebuilt = StructArray::new(
+ struct_array.fields().clone(),
+ new_columns,
+ struct_array.nulls().cloned(),
+ );
+ return Ok((Arc::new(rebuilt), true));
+ }
+
+ if let Some(list_array) = array.as_any().downcast_ref::<ListArray>() {
+ let (new_values, mutated) =
compact_array(list_array.values().clone())?;
+ if !mutated {
+ return Ok((array, false));
+ }
Review Comment:
I have added several test cases to cover the optimized case.
--
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]
