You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@sedona.apache.org by "Kontinuation (via GitHub)" <gi...@apache.org> on 2026/01/14 15:29:07 UTC
[PR] chore(rust/sedona-spatial-join): More accurate batch in-memory size estimation [sedona-db]
Kontinuation opened a new pull request, #515:
URL: https://github.com/apache/sedona-db/pull/515
This patch improves the accuracy of memory usage estimation by implementing our own functions for estimating the in-memory sizes of record batches and arrow arrays.
The rationale is similar to https://github.com/apache/datafusion/pull/13377. If we don't roll our own memory usage estimation function but call `RecordBatch::get_array_memory_size` instead, we'll get insanely inaccurate numbers for spilled batches read using `arrow::ipc::reader::StreamReader`.
Future work: use the memory pool API of arrow-rs for more accurate memory usage accounting. See https://github.com/apache/arrow-rs/issues/8137.
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] chore(rust/sedona-spatial-join): More accurate batch in-memory size estimation [sedona-db]
Posted by "Kontinuation (via GitHub)" <gi...@apache.org>.
Kontinuation commented on code in PR #515:
URL: https://github.com/apache/sedona-db/pull/515#discussion_r2694869758
##########
rust/sedona-spatial-join/src/utils/arrow_utils.rs:
##########
@@ -0,0 +1,166 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::{Array, ArrayData, RecordBatch};
+use arrow_array::ArrayRef;
+use arrow_schema::{ArrowError, DataType};
+use datafusion_common::Result;
+
+/// Estimate the in-memory size of a given RecordBatch. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_record_batch_memory_size(batch: &RecordBatch) -> Result<usize> {
+ let mut total_size = 0;
+
+ for array in batch.columns() {
+ let array_data = array.to_data();
+ total_size += count_array_data_memory_size(&array_data)?;
+ }
+
+ Ok(total_size)
+}
+
+/// Estimate the in-memory size of a given Arrow array. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_array_memory_size(array: &ArrayRef) -> Result<usize> {
+ let array_data = array.to_data();
+ let size = count_array_data_memory_size(&array_data)?;
+ Ok(size)
+}
+
+/// The maximum number of bytes that can be stored inline in a byte view.
+///
+/// See [`ByteView`] and [`GenericByteViewArray`] for more information on the
+/// layout of the views.
+///
+/// [`GenericByteViewArray`]: https://docs.rs/arrow/latest/arrow/array/struct.GenericByteViewArray.html
+pub const MAX_INLINE_VIEW_LEN: u32 = 12;
+
+/// Count the memory usage of `array_data` and its children recursively.
+fn count_array_data_memory_size(array_data: &ArrayData) -> core::result::Result<usize, ArrowError> {
+ Ok(get_binary_view_value_size(array_data)? + array_data.get_slice_memory_size()?)
+}
+
+fn get_binary_view_value_size(array_data: &ArrayData) -> Result<usize, ArrowError> {
+ let mut result: usize = 0;
+ let array_data_type = array_data.data_type();
+
+ if matches!(array_data_type, DataType::BinaryView | DataType::Utf8View) {
+ // The views buffer contains length view structures with the following layout:
+ // https://arrow.apache.org/docs/format/Columnar.html#variable-size-binary-view-layout
+ //
+ // * Short strings, length <= 12
+ // | Bytes 0-3 | Bytes 4-15 |
+ // |------------|---------------------------------------|
+ // | length | data (padded with 0) |
+ //
+ // * Long strings, length > 12
+ // | Bytes 0-3 | Bytes 4-7 | Bytes 8-11 | Bytes 12-15 |
+ // |------------|------------|------------|-------------|
+ // | length | prefix | buf. index | offset |
+ let views = &array_data.buffer::<u128>(0)[..array_data.len()];
+ result = views
+ .iter()
+ .map(|v| {
+ let len = *v as u32;
+ if len > MAX_INLINE_VIEW_LEN {
+ len as usize
+ } else {
+ 0
+ }
+ })
+ .sum();
+ }
+
+ for child in array_data.child_data() {
+ result += get_binary_view_value_size(child)?;
+ }
+ Ok(result)
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow_array::{BinaryViewArray, StringViewArray, StructArray};
+ use arrow_schema::{DataType, Field};
+ use std::sync::Arc;
+
+ #[test]
+ fn test_string_view_array_memory_size() {
+ let array = StringViewArray::from(vec![
+ "short", // Inline
+ "Long string that is definitely longer than 12 bytes", // 51 bytes
+ ]);
+ let array_ref: ArrayRef = Arc::new(array);
+ let size = get_array_memory_size(&array_ref).unwrap();
+ // Views: 2 * 16 = 32 bytes
+ // Data: 51 bytes
+ // Total: 83 bytes
+ assert_eq!(size, 83);
+ }
+
+ #[test]
+ fn test_binary_view_array_memory_size() {
+ let array = BinaryViewArray::from(vec![
+ "short".as_bytes(),
+ "Long string that is definitely longer than 12 bytes".as_bytes(),
+ ]);
+ let array_ref: ArrayRef = Arc::new(array);
+ let size = get_array_memory_size(&array_ref).unwrap();
+ assert_eq!(size, 83);
+ }
+
+ #[test]
+ fn test_struct_array_with_view_memory_size() {
+ let string_view_array = StringViewArray::from(vec![
+ "short",
+ "Long string that is definitely longer than 12 bytes",
+ ]);
+ let boolean_array = arrow_array::BooleanArray::from(vec![true, false]);
+
+ let struct_array = StructArray::from(vec![
+ (
+ Arc::new(Field::new("a", DataType::Utf8View, false)),
+ Arc::new(string_view_array) as ArrayRef,
+ ),
+ (
+ Arc::new(Field::new("b", DataType::Boolean, false)),
+ Arc::new(boolean_array) as ArrayRef,
+ ),
+ ]);
+
+ let array_ref: ArrayRef = Arc::new(struct_array);
+ let size = get_array_memory_size(&array_ref).unwrap();
+ // 83 (StringView) + 1 (Boolean values) = 84
+ assert_eq!(size, 84);
Review Comment:
I'd like to leave `get_array_data_memory_size` as is and log a ticket.
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] chore(rust/sedona-spatial-join): More accurate batch in-memory size estimation [sedona-db]
Posted by "Kontinuation (via GitHub)" <gi...@apache.org>.
Kontinuation commented on PR #515:
URL: https://github.com/apache/sedona-db/pull/515#issuecomment-3752922147
I have added the requested tests and also renamed the private function `count_array_data_memory_size` to `get_array_data_memory_size` since the original name looked awful.
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] chore(rust/sedona-spatial-join): More accurate batch in-memory size estimation [sedona-db]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #515:
URL: https://github.com/apache/sedona-db/pull/515#discussion_r2691655827
##########
rust/sedona-spatial-join/src/utils/arrow_utils.rs:
##########
@@ -0,0 +1,166 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::{Array, ArrayData, RecordBatch};
+use arrow_array::ArrayRef;
+use arrow_schema::{ArrowError, DataType};
+use datafusion_common::Result;
+
+/// Estimate the in-memory size of a given RecordBatch. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_record_batch_memory_size(batch: &RecordBatch) -> Result<usize> {
+ let mut total_size = 0;
+
+ for array in batch.columns() {
+ let array_data = array.to_data();
+ total_size += count_array_data_memory_size(&array_data)?;
+ }
+
+ Ok(total_size)
+}
+
+/// Estimate the in-memory size of a given Arrow array. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_array_memory_size(array: &ArrayRef) -> Result<usize> {
+ let array_data = array.to_data();
+ let size = count_array_data_memory_size(&array_data)?;
+ Ok(size)
+}
Review Comment:
Is https://docs.rs/arrow/latest/arrow/array/struct.ArrayData.html#method.get_slice_memory_size helpful here?
##########
rust/sedona-spatial-join/src/utils/arrow_utils.rs:
##########
@@ -0,0 +1,166 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::{Array, ArrayData, RecordBatch};
+use arrow_array::ArrayRef;
+use arrow_schema::{ArrowError, DataType};
+use datafusion_common::Result;
+
+/// Estimate the in-memory size of a given RecordBatch. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_record_batch_memory_size(batch: &RecordBatch) -> Result<usize> {
+ let mut total_size = 0;
+
+ for array in batch.columns() {
+ let array_data = array.to_data();
+ total_size += count_array_data_memory_size(&array_data)?;
+ }
+
+ Ok(total_size)
+}
+
+/// Estimate the in-memory size of a given Arrow array. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_array_memory_size(array: &ArrayRef) -> Result<usize> {
+ let array_data = array.to_data();
+ let size = count_array_data_memory_size(&array_data)?;
+ Ok(size)
+}
+
+/// The maximum number of bytes that can be stored inline in a byte view.
+///
+/// See [`ByteView`] and [`GenericByteViewArray`] for more information on the
+/// layout of the views.
+///
+/// [`GenericByteViewArray`]: https://docs.rs/arrow/latest/arrow/array/struct.GenericByteViewArray.html
+pub const MAX_INLINE_VIEW_LEN: u32 = 12;
+
+/// Count the memory usage of `array_data` and its children recursively.
+fn count_array_data_memory_size(array_data: &ArrayData) -> core::result::Result<usize, ArrowError> {
+ Ok(get_binary_view_value_size(array_data)? + array_data.get_slice_memory_size()?)
+}
+
+fn get_binary_view_value_size(array_data: &ArrayData) -> Result<usize, ArrowError> {
+ let mut result: usize = 0;
+ let array_data_type = array_data.data_type();
+
+ if matches!(array_data_type, DataType::BinaryView | DataType::Utf8View) {
Review Comment:
This seems like it is missing a branch for non BinaryView types. I imagine we aren't trying to ignore the memory occupied by non-view types?
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] chore(rust/sedona-spatial-join): More accurate batch in-memory size estimation [sedona-db]
Posted by "Kontinuation (via GitHub)" <gi...@apache.org>.
Kontinuation commented on code in PR #515:
URL: https://github.com/apache/sedona-db/pull/515#discussion_r2694753611
##########
rust/sedona-spatial-join/src/utils/arrow_utils.rs:
##########
@@ -0,0 +1,166 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::{Array, ArrayData, RecordBatch};
+use arrow_array::ArrayRef;
+use arrow_schema::{ArrowError, DataType};
+use datafusion_common::Result;
+
+/// Estimate the in-memory size of a given RecordBatch. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_record_batch_memory_size(batch: &RecordBatch) -> Result<usize> {
+ let mut total_size = 0;
+
+ for array in batch.columns() {
+ let array_data = array.to_data();
+ total_size += count_array_data_memory_size(&array_data)?;
+ }
+
+ Ok(total_size)
+}
+
+/// Estimate the in-memory size of a given Arrow array. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_array_memory_size(array: &ArrayRef) -> Result<usize> {
+ let array_data = array.to_data();
+ let size = count_array_data_memory_size(&array_data)?;
+ Ok(size)
+}
+
+/// The maximum number of bytes that can be stored inline in a byte view.
+///
+/// See [`ByteView`] and [`GenericByteViewArray`] for more information on the
+/// layout of the views.
+///
+/// [`GenericByteViewArray`]: https://docs.rs/arrow/latest/arrow/array/struct.GenericByteViewArray.html
+pub const MAX_INLINE_VIEW_LEN: u32 = 12;
+
+/// Count the memory usage of `array_data` and its children recursively.
+fn count_array_data_memory_size(array_data: &ArrayData) -> core::result::Result<usize, ArrowError> {
+ Ok(get_binary_view_value_size(array_data)? + array_data.get_slice_memory_size()?)
+}
+
+fn get_binary_view_value_size(array_data: &ArrayData) -> Result<usize, ArrowError> {
+ let mut result: usize = 0;
+ let array_data_type = array_data.data_type();
+
+ if matches!(array_data_type, DataType::BinaryView | DataType::Utf8View) {
+ // The views buffer contains length view structures with the following layout:
+ // https://arrow.apache.org/docs/format/Columnar.html#variable-size-binary-view-layout
+ //
+ // * Short strings, length <= 12
+ // | Bytes 0-3 | Bytes 4-15 |
+ // |------------|---------------------------------------|
+ // | length | data (padded with 0) |
+ //
+ // * Long strings, length > 12
+ // | Bytes 0-3 | Bytes 4-7 | Bytes 8-11 | Bytes 12-15 |
+ // |------------|------------|------------|-------------|
+ // | length | prefix | buf. index | offset |
+ let views = &array_data.buffer::<u128>(0)[..array_data.len()];
+ result = views
+ .iter()
+ .map(|v| {
+ let len = *v as u32;
+ if len > MAX_INLINE_VIEW_LEN {
+ len as usize
+ } else {
+ 0
+ }
+ })
+ .sum();
+ }
+
+ for child in array_data.child_data() {
+ result += get_binary_view_value_size(child)?;
+ }
+ Ok(result)
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow_array::{BinaryViewArray, StringViewArray, StructArray};
+ use arrow_schema::{DataType, Field};
+ use std::sync::Arc;
+
+ #[test]
+ fn test_string_view_array_memory_size() {
+ let array = StringViewArray::from(vec![
+ "short", // Inline
+ "Long string that is definitely longer than 12 bytes", // 51 bytes
+ ]);
+ let array_ref: ArrayRef = Arc::new(array);
+ let size = get_array_memory_size(&array_ref).unwrap();
+ // Views: 2 * 16 = 32 bytes
+ // Data: 51 bytes
+ // Total: 83 bytes
+ assert_eq!(size, 83);
+ }
+
+ #[test]
+ fn test_binary_view_array_memory_size() {
+ let array = BinaryViewArray::from(vec![
+ "short".as_bytes(),
+ "Long string that is definitely longer than 12 bytes".as_bytes(),
+ ]);
+ let array_ref: ArrayRef = Arc::new(array);
+ let size = get_array_memory_size(&array_ref).unwrap();
+ assert_eq!(size, 83);
+ }
+
+ #[test]
+ fn test_struct_array_with_view_memory_size() {
+ let string_view_array = StringViewArray::from(vec![
+ "short",
+ "Long string that is definitely longer than 12 bytes",
+ ]);
+ let boolean_array = arrow_array::BooleanArray::from(vec![true, false]);
+
+ let struct_array = StructArray::from(vec![
+ (
+ Arc::new(Field::new("a", DataType::Utf8View, false)),
+ Arc::new(string_view_array) as ArrayRef,
+ ),
+ (
+ Arc::new(Field::new("b", DataType::Boolean, false)),
+ Arc::new(boolean_array) as ArrayRef,
+ ),
+ ]);
+
+ let array_ref: ArrayRef = Arc::new(struct_array);
+ let size = get_array_memory_size(&array_ref).unwrap();
+ // 83 (StringView) + 1 (Boolean values) = 84
+ assert_eq!(size, 84);
+ }
+
+ #[test]
+ fn test_sliced_view_array_memory_size() {
+ let array = StringViewArray::from(vec![
+ "short",
+ "Long string that is definitely longer than 12 bytes",
+ "Another long string to make buffer larger",
+ ]);
+ let sliced = array.slice(0, 2);
Review Comment:
`array.slice(2, 1)` is invalid since the array contains only 2 rows. I have tested `array.slice(1, 1)` and `array.slice(2, 0)` instead.
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] chore(rust/sedona-spatial-join): More accurate batch in-memory size estimation [sedona-db]
Posted by "Kontinuation (via GitHub)" <gi...@apache.org>.
Kontinuation commented on code in PR #515:
URL: https://github.com/apache/sedona-db/pull/515#discussion_r2692580163
##########
rust/sedona-spatial-join/src/utils/arrow_utils.rs:
##########
@@ -0,0 +1,166 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::{Array, ArrayData, RecordBatch};
+use arrow_array::ArrayRef;
+use arrow_schema::{ArrowError, DataType};
+use datafusion_common::Result;
+
+/// Estimate the in-memory size of a given RecordBatch. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_record_batch_memory_size(batch: &RecordBatch) -> Result<usize> {
+ let mut total_size = 0;
+
+ for array in batch.columns() {
+ let array_data = array.to_data();
+ total_size += count_array_data_memory_size(&array_data)?;
+ }
+
+ Ok(total_size)
+}
+
+/// Estimate the in-memory size of a given Arrow array. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_array_memory_size(array: &ArrayRef) -> Result<usize> {
+ let array_data = array.to_data();
+ let size = count_array_data_memory_size(&array_data)?;
+ Ok(size)
+}
Review Comment:
It is already using `get_slice_memory_size` under the hood. However, `get_slice_memory_size` does not take buffers of binary/utf8 views into account, `count_array_data_memory_size` is meant to fix that.
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] chore(rust/sedona-spatial-join): More accurate batch in-memory size estimation [sedona-db]
Posted by "Kontinuation (via GitHub)" <gi...@apache.org>.
Kontinuation commented on code in PR #515:
URL: https://github.com/apache/sedona-db/pull/515#discussion_r2692967502
##########
rust/sedona-spatial-join/src/utils/arrow_utils.rs:
##########
@@ -0,0 +1,166 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::{Array, ArrayData, RecordBatch};
+use arrow_array::ArrayRef;
+use arrow_schema::{ArrowError, DataType};
+use datafusion_common::Result;
+
+/// Estimate the in-memory size of a given RecordBatch. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_record_batch_memory_size(batch: &RecordBatch) -> Result<usize> {
+ let mut total_size = 0;
+
+ for array in batch.columns() {
+ let array_data = array.to_data();
+ total_size += count_array_data_memory_size(&array_data)?;
+ }
+
+ Ok(total_size)
+}
+
+/// Estimate the in-memory size of a given Arrow array. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_array_memory_size(array: &ArrayRef) -> Result<usize> {
+ let array_data = array.to_data();
+ let size = count_array_data_memory_size(&array_data)?;
+ Ok(size)
+}
+
+/// The maximum number of bytes that can be stored inline in a byte view.
+///
+/// See [`ByteView`] and [`GenericByteViewArray`] for more information on the
+/// layout of the views.
+///
+/// [`GenericByteViewArray`]: https://docs.rs/arrow/latest/arrow/array/struct.GenericByteViewArray.html
+pub const MAX_INLINE_VIEW_LEN: u32 = 12;
+
+/// Count the memory usage of `array_data` and its children recursively.
+fn count_array_data_memory_size(array_data: &ArrayData) -> core::result::Result<usize, ArrowError> {
+ Ok(get_binary_view_value_size(array_data)? + array_data.get_slice_memory_size()?)
+}
+
+fn get_binary_view_value_size(array_data: &ArrayData) -> Result<usize, ArrowError> {
+ let mut result: usize = 0;
+ let array_data_type = array_data.data_type();
+
+ if matches!(array_data_type, DataType::BinaryView | DataType::Utf8View) {
+ // The views buffer contains length view structures with the following layout:
+ // https://arrow.apache.org/docs/format/Columnar.html#variable-size-binary-view-layout
+ //
+ // * Short strings, length <= 12
+ // | Bytes 0-3 | Bytes 4-15 |
+ // |------------|---------------------------------------|
+ // | length | data (padded with 0) |
+ //
+ // * Long strings, length > 12
+ // | Bytes 0-3 | Bytes 4-7 | Bytes 8-11 | Bytes 12-15 |
+ // |------------|------------|------------|-------------|
+ // | length | prefix | buf. index | offset |
+ let views = &array_data.buffer::<u128>(0)[..array_data.len()];
Review Comment:
This does not look correct to me. `ArrayData::buffer` has already taken offset into consideration: https://github.com/apache/arrow-rs/blob/57.2.0/arrow-data/src/data.rs#L580-L588
I'll add the test cases for buffers not starting at 0 anyway.
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] chore(rust/sedona-spatial-join): More accurate batch in-memory size estimation [sedona-db]
Posted by "Kontinuation (via GitHub)" <gi...@apache.org>.
Kontinuation merged PR #515:
URL: https://github.com/apache/sedona-db/pull/515
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] chore(rust/sedona-spatial-join): More accurate batch in-memory size estimation [sedona-db]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #515:
URL: https://github.com/apache/sedona-db/pull/515#discussion_r2694697886
##########
rust/sedona-spatial-join/src/utils/arrow_utils.rs:
##########
@@ -0,0 +1,166 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::{Array, ArrayData, RecordBatch};
+use arrow_array::ArrayRef;
+use arrow_schema::{ArrowError, DataType};
+use datafusion_common::Result;
+
+/// Estimate the in-memory size of a given RecordBatch. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_record_batch_memory_size(batch: &RecordBatch) -> Result<usize> {
+ let mut total_size = 0;
+
+ for array in batch.columns() {
+ let array_data = array.to_data();
+ total_size += count_array_data_memory_size(&array_data)?;
+ }
+
+ Ok(total_size)
+}
+
+/// Estimate the in-memory size of a given Arrow array. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_array_memory_size(array: &ArrayRef) -> Result<usize> {
+ let array_data = array.to_data();
+ let size = count_array_data_memory_size(&array_data)?;
+ Ok(size)
+}
+
+/// The maximum number of bytes that can be stored inline in a byte view.
+///
+/// See [`ByteView`] and [`GenericByteViewArray`] for more information on the
+/// layout of the views.
+///
+/// [`GenericByteViewArray`]: https://docs.rs/arrow/latest/arrow/array/struct.GenericByteViewArray.html
+pub const MAX_INLINE_VIEW_LEN: u32 = 12;
+
+/// Count the memory usage of `array_data` and its children recursively.
+fn count_array_data_memory_size(array_data: &ArrayData) -> core::result::Result<usize, ArrowError> {
+ Ok(get_binary_view_value_size(array_data)? + array_data.get_slice_memory_size()?)
+}
+
+fn get_binary_view_value_size(array_data: &ArrayData) -> Result<usize, ArrowError> {
+ let mut result: usize = 0;
+ let array_data_type = array_data.data_type();
+
+ if matches!(array_data_type, DataType::BinaryView | DataType::Utf8View) {
+ // The views buffer contains length view structures with the following layout:
+ // https://arrow.apache.org/docs/format/Columnar.html#variable-size-binary-view-layout
+ //
+ // * Short strings, length <= 12
+ // | Bytes 0-3 | Bytes 4-15 |
+ // |------------|---------------------------------------|
+ // | length | data (padded with 0) |
+ //
+ // * Long strings, length > 12
+ // | Bytes 0-3 | Bytes 4-7 | Bytes 8-11 | Bytes 12-15 |
+ // |------------|------------|------------|-------------|
+ // | length | prefix | buf. index | offset |
+ let views = &array_data.buffer::<u128>(0)[..array_data.len()];
+ result = views
+ .iter()
+ .map(|v| {
+ let len = *v as u32;
+ if len > MAX_INLINE_VIEW_LEN {
+ len as usize
+ } else {
+ 0
+ }
+ })
+ .sum();
+ }
+
+ for child in array_data.child_data() {
+ result += get_binary_view_value_size(child)?;
+ }
+ Ok(result)
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow_array::{BinaryViewArray, StringViewArray, StructArray};
+ use arrow_schema::{DataType, Field};
+ use std::sync::Arc;
+
+ #[test]
+ fn test_string_view_array_memory_size() {
+ let array = StringViewArray::from(vec![
+ "short", // Inline
+ "Long string that is definitely longer than 12 bytes", // 51 bytes
+ ]);
+ let array_ref: ArrayRef = Arc::new(array);
+ let size = get_array_memory_size(&array_ref).unwrap();
+ // Views: 2 * 16 = 32 bytes
+ // Data: 51 bytes
+ // Total: 83 bytes
+ assert_eq!(size, 83);
+ }
+
+ #[test]
+ fn test_binary_view_array_memory_size() {
+ let array = BinaryViewArray::from(vec![
+ "short".as_bytes(),
+ "Long string that is definitely longer than 12 bytes".as_bytes(),
+ ]);
+ let array_ref: ArrayRef = Arc::new(array);
+ let size = get_array_memory_size(&array_ref).unwrap();
+ assert_eq!(size, 83);
+ }
+
+ #[test]
+ fn test_struct_array_with_view_memory_size() {
+ let string_view_array = StringViewArray::from(vec![
+ "short",
+ "Long string that is definitely longer than 12 bytes",
+ ]);
+ let boolean_array = arrow_array::BooleanArray::from(vec![true, false]);
+
+ let struct_array = StructArray::from(vec![
+ (
+ Arc::new(Field::new("a", DataType::Utf8View, false)),
+ Arc::new(string_view_array) as ArrayRef,
+ ),
+ (
+ Arc::new(Field::new("b", DataType::Boolean, false)),
+ Arc::new(boolean_array) as ArrayRef,
+ ),
+ ]);
+
+ let array_ref: ArrayRef = Arc::new(struct_array);
+ let size = get_array_memory_size(&array_ref).unwrap();
+ // 83 (StringView) + 1 (Boolean values) = 84
+ assert_eq!(size, 84);
Review Comment:
What do you think about this test case? (Also OK to punt with a comment in the appropriate location that we arenot calculating correct memory usage of sliced lists/structs?)
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] chore(rust/sedona-spatial-join): More accurate batch in-memory size estimation [sedona-db]
Posted by "Kontinuation (via GitHub)" <gi...@apache.org>.
Kontinuation commented on code in PR #515:
URL: https://github.com/apache/sedona-db/pull/515#discussion_r2694896405
##########
rust/sedona-spatial-join/src/utils/arrow_utils.rs:
##########
@@ -0,0 +1,166 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::{Array, ArrayData, RecordBatch};
+use arrow_array::ArrayRef;
+use arrow_schema::{ArrowError, DataType};
+use datafusion_common::Result;
+
+/// Estimate the in-memory size of a given RecordBatch. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_record_batch_memory_size(batch: &RecordBatch) -> Result<usize> {
+ let mut total_size = 0;
+
+ for array in batch.columns() {
+ let array_data = array.to_data();
+ total_size += count_array_data_memory_size(&array_data)?;
+ }
+
+ Ok(total_size)
+}
+
+/// Estimate the in-memory size of a given Arrow array. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_array_memory_size(array: &ArrayRef) -> Result<usize> {
+ let array_data = array.to_data();
+ let size = count_array_data_memory_size(&array_data)?;
+ Ok(size)
+}
+
+/// The maximum number of bytes that can be stored inline in a byte view.
+///
+/// See [`ByteView`] and [`GenericByteViewArray`] for more information on the
+/// layout of the views.
+///
+/// [`GenericByteViewArray`]: https://docs.rs/arrow/latest/arrow/array/struct.GenericByteViewArray.html
+pub const MAX_INLINE_VIEW_LEN: u32 = 12;
+
+/// Count the memory usage of `array_data` and its children recursively.
+fn count_array_data_memory_size(array_data: &ArrayData) -> core::result::Result<usize, ArrowError> {
+ Ok(get_binary_view_value_size(array_data)? + array_data.get_slice_memory_size()?)
+}
+
+fn get_binary_view_value_size(array_data: &ArrayData) -> Result<usize, ArrowError> {
+ let mut result: usize = 0;
+ let array_data_type = array_data.data_type();
+
+ if matches!(array_data_type, DataType::BinaryView | DataType::Utf8View) {
+ // The views buffer contains length view structures with the following layout:
+ // https://arrow.apache.org/docs/format/Columnar.html#variable-size-binary-view-layout
+ //
+ // * Short strings, length <= 12
+ // | Bytes 0-3 | Bytes 4-15 |
+ // |------------|---------------------------------------|
+ // | length | data (padded with 0) |
+ //
+ // * Long strings, length > 12
+ // | Bytes 0-3 | Bytes 4-7 | Bytes 8-11 | Bytes 12-15 |
+ // |------------|------------|------------|-------------|
+ // | length | prefix | buf. index | offset |
+ let views = &array_data.buffer::<u128>(0)[..array_data.len()];
+ result = views
+ .iter()
+ .map(|v| {
+ let len = *v as u32;
+ if len > MAX_INLINE_VIEW_LEN {
+ len as usize
+ } else {
+ 0
+ }
+ })
+ .sum();
+ }
+
+ for child in array_data.child_data() {
+ result += get_binary_view_value_size(child)?;
+ }
+ Ok(result)
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow_array::{BinaryViewArray, StringViewArray, StructArray};
+ use arrow_schema::{DataType, Field};
+ use std::sync::Arc;
+
+ #[test]
+ fn test_string_view_array_memory_size() {
+ let array = StringViewArray::from(vec![
+ "short", // Inline
+ "Long string that is definitely longer than 12 bytes", // 51 bytes
+ ]);
+ let array_ref: ArrayRef = Arc::new(array);
+ let size = get_array_memory_size(&array_ref).unwrap();
+ // Views: 2 * 16 = 32 bytes
+ // Data: 51 bytes
+ // Total: 83 bytes
+ assert_eq!(size, 83);
+ }
+
+ #[test]
+ fn test_binary_view_array_memory_size() {
+ let array = BinaryViewArray::from(vec![
+ "short".as_bytes(),
+ "Long string that is definitely longer than 12 bytes".as_bytes(),
+ ]);
+ let array_ref: ArrayRef = Arc::new(array);
+ let size = get_array_memory_size(&array_ref).unwrap();
+ assert_eq!(size, 83);
+ }
+
+ #[test]
+ fn test_struct_array_with_view_memory_size() {
+ let string_view_array = StringViewArray::from(vec![
+ "short",
+ "Long string that is definitely longer than 12 bytes",
+ ]);
+ let boolean_array = arrow_array::BooleanArray::from(vec![true, false]);
+
+ let struct_array = StructArray::from(vec![
+ (
+ Arc::new(Field::new("a", DataType::Utf8View, false)),
+ Arc::new(string_view_array) as ArrayRef,
+ ),
+ (
+ Arc::new(Field::new("b", DataType::Boolean, false)),
+ Arc::new(boolean_array) as ArrayRef,
+ ),
+ ]);
+
+ let array_ref: ArrayRef = Arc::new(struct_array);
+ let size = get_array_memory_size(&array_ref).unwrap();
+ // 83 (StringView) + 1 (Boolean values) = 84
+ assert_eq!(size, 84);
Review Comment:
And also added a XFAIL test case for list array slices.
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] chore(rust/sedona-spatial-join): More accurate batch in-memory size estimation [sedona-db]
Posted by "Kontinuation (via GitHub)" <gi...@apache.org>.
Kontinuation commented on code in PR #515:
URL: https://github.com/apache/sedona-db/pull/515#discussion_r2692581206
##########
rust/sedona-spatial-join/src/utils/arrow_utils.rs:
##########
@@ -0,0 +1,166 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::{Array, ArrayData, RecordBatch};
+use arrow_array::ArrayRef;
+use arrow_schema::{ArrowError, DataType};
+use datafusion_common::Result;
+
+/// Estimate the in-memory size of a given RecordBatch. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_record_batch_memory_size(batch: &RecordBatch) -> Result<usize> {
+ let mut total_size = 0;
+
+ for array in batch.columns() {
+ let array_data = array.to_data();
+ total_size += count_array_data_memory_size(&array_data)?;
+ }
+
+ Ok(total_size)
+}
+
+/// Estimate the in-memory size of a given Arrow array. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_array_memory_size(array: &ArrayRef) -> Result<usize> {
+ let array_data = array.to_data();
+ let size = count_array_data_memory_size(&array_data)?;
+ Ok(size)
+}
+
+/// The maximum number of bytes that can be stored inline in a byte view.
+///
+/// See [`ByteView`] and [`GenericByteViewArray`] for more information on the
+/// layout of the views.
+///
+/// [`GenericByteViewArray`]: https://docs.rs/arrow/latest/arrow/array/struct.GenericByteViewArray.html
+pub const MAX_INLINE_VIEW_LEN: u32 = 12;
+
+/// Count the memory usage of `array_data` and its children recursively.
+fn count_array_data_memory_size(array_data: &ArrayData) -> core::result::Result<usize, ArrowError> {
+ Ok(get_binary_view_value_size(array_data)? + array_data.get_slice_memory_size()?)
+}
+
+fn get_binary_view_value_size(array_data: &ArrayData) -> Result<usize, ArrowError> {
+ let mut result: usize = 0;
+ let array_data_type = array_data.data_type();
+
+ if matches!(array_data_type, DataType::BinaryView | DataType::Utf8View) {
Review Comment:
Non-view types were already taken care by `get_slice_memory_size`.
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] chore(rust/sedona-spatial-join): More accurate batch in-memory size estimation [sedona-db]
Posted by "Kontinuation (via GitHub)" <gi...@apache.org>.
Kontinuation commented on code in PR #515:
URL: https://github.com/apache/sedona-db/pull/515#discussion_r2694750679
##########
rust/sedona-spatial-join/src/utils/arrow_utils.rs:
##########
@@ -0,0 +1,166 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::{Array, ArrayData, RecordBatch};
+use arrow_array::ArrayRef;
+use arrow_schema::{ArrowError, DataType};
+use datafusion_common::Result;
+
+/// Estimate the in-memory size of a given RecordBatch. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_record_batch_memory_size(batch: &RecordBatch) -> Result<usize> {
+ let mut total_size = 0;
+
+ for array in batch.columns() {
+ let array_data = array.to_data();
+ total_size += count_array_data_memory_size(&array_data)?;
+ }
+
+ Ok(total_size)
+}
+
+/// Estimate the in-memory size of a given Arrow array. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_array_memory_size(array: &ArrayRef) -> Result<usize> {
+ let array_data = array.to_data();
+ let size = count_array_data_memory_size(&array_data)?;
+ Ok(size)
+}
+
+/// The maximum number of bytes that can be stored inline in a byte view.
+///
+/// See [`ByteView`] and [`GenericByteViewArray`] for more information on the
+/// layout of the views.
+///
+/// [`GenericByteViewArray`]: https://docs.rs/arrow/latest/arrow/array/struct.GenericByteViewArray.html
+pub const MAX_INLINE_VIEW_LEN: u32 = 12;
+
+/// Count the memory usage of `array_data` and its children recursively.
+fn count_array_data_memory_size(array_data: &ArrayData) -> core::result::Result<usize, ArrowError> {
+ Ok(get_binary_view_value_size(array_data)? + array_data.get_slice_memory_size()?)
+}
+
+fn get_binary_view_value_size(array_data: &ArrayData) -> Result<usize, ArrowError> {
+ let mut result: usize = 0;
+ let array_data_type = array_data.data_type();
+
+ if matches!(array_data_type, DataType::BinaryView | DataType::Utf8View) {
+ // The views buffer contains length view structures with the following layout:
+ // https://arrow.apache.org/docs/format/Columnar.html#variable-size-binary-view-layout
+ //
+ // * Short strings, length <= 12
+ // | Bytes 0-3 | Bytes 4-15 |
+ // |------------|---------------------------------------|
+ // | length | data (padded with 0) |
+ //
+ // * Long strings, length > 12
+ // | Bytes 0-3 | Bytes 4-7 | Bytes 8-11 | Bytes 12-15 |
+ // |------------|------------|------------|-------------|
+ // | length | prefix | buf. index | offset |
+ let views = &array_data.buffer::<u128>(0)[..array_data.len()];
+ result = views
+ .iter()
+ .map(|v| {
+ let len = *v as u32;
+ if len > MAX_INLINE_VIEW_LEN {
+ len as usize
+ } else {
+ 0
+ }
+ })
+ .sum();
+ }
+
+ for child in array_data.child_data() {
+ result += get_binary_view_value_size(child)?;
+ }
+ Ok(result)
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow_array::{BinaryViewArray, StringViewArray, StructArray};
+ use arrow_schema::{DataType, Field};
+ use std::sync::Arc;
+
+ #[test]
+ fn test_string_view_array_memory_size() {
+ let array = StringViewArray::from(vec![
+ "short", // Inline
+ "Long string that is definitely longer than 12 bytes", // 51 bytes
+ ]);
+ let array_ref: ArrayRef = Arc::new(array);
+ let size = get_array_memory_size(&array_ref).unwrap();
+ // Views: 2 * 16 = 32 bytes
+ // Data: 51 bytes
+ // Total: 83 bytes
+ assert_eq!(size, 83);
+ }
+
+ #[test]
+ fn test_binary_view_array_memory_size() {
+ let array = BinaryViewArray::from(vec![
+ "short".as_bytes(),
+ "Long string that is definitely longer than 12 bytes".as_bytes(),
+ ]);
+ let array_ref: ArrayRef = Arc::new(array);
+ let size = get_array_memory_size(&array_ref).unwrap();
+ assert_eq!(size, 83);
+ }
+
+ #[test]
+ fn test_struct_array_with_view_memory_size() {
+ let string_view_array = StringViewArray::from(vec![
+ "short",
+ "Long string that is definitely longer than 12 bytes",
+ ]);
+ let boolean_array = arrow_array::BooleanArray::from(vec![true, false]);
+
+ let struct_array = StructArray::from(vec![
+ (
+ Arc::new(Field::new("a", DataType::Utf8View, false)),
+ Arc::new(string_view_array) as ArrayRef,
+ ),
+ (
+ Arc::new(Field::new("b", DataType::Boolean, false)),
+ Arc::new(boolean_array) as ArrayRef,
+ ),
+ ]);
+
+ let array_ref: ArrayRef = Arc::new(struct_array);
+ let size = get_array_memory_size(&array_ref).unwrap();
+ // 83 (StringView) + 1 (Boolean values) = 84
+ assert_eq!(size, 84);
Review Comment:
I have added `array_ref.slice(0, 1)` test below and it gave correct result.
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] chore(rust/sedona-spatial-join): More accurate batch in-memory size estimation [sedona-db]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #515:
URL: https://github.com/apache/sedona-db/pull/515#discussion_r2692907297
##########
rust/sedona-spatial-join/src/utils/arrow_utils.rs:
##########
@@ -0,0 +1,166 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::{Array, ArrayData, RecordBatch};
+use arrow_array::ArrayRef;
+use arrow_schema::{ArrowError, DataType};
+use datafusion_common::Result;
+
+/// Estimate the in-memory size of a given RecordBatch. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_record_batch_memory_size(batch: &RecordBatch) -> Result<usize> {
+ let mut total_size = 0;
+
+ for array in batch.columns() {
+ let array_data = array.to_data();
+ total_size += count_array_data_memory_size(&array_data)?;
+ }
+
+ Ok(total_size)
+}
+
+/// Estimate the in-memory size of a given Arrow array. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_array_memory_size(array: &ArrayRef) -> Result<usize> {
+ let array_data = array.to_data();
+ let size = count_array_data_memory_size(&array_data)?;
+ Ok(size)
+}
+
+/// The maximum number of bytes that can be stored inline in a byte view.
+///
+/// See [`ByteView`] and [`GenericByteViewArray`] for more information on the
+/// layout of the views.
+///
+/// [`GenericByteViewArray`]: https://docs.rs/arrow/latest/arrow/array/struct.GenericByteViewArray.html
+pub const MAX_INLINE_VIEW_LEN: u32 = 12;
+
+/// Count the memory usage of `array_data` and its children recursively.
+fn count_array_data_memory_size(array_data: &ArrayData) -> core::result::Result<usize, ArrowError> {
+ Ok(get_binary_view_value_size(array_data)? + array_data.get_slice_memory_size()?)
+}
+
+fn get_binary_view_value_size(array_data: &ArrayData) -> Result<usize, ArrowError> {
+ let mut result: usize = 0;
+ let array_data_type = array_data.data_type();
+
+ if matches!(array_data_type, DataType::BinaryView | DataType::Utf8View) {
+ // The views buffer contains length view structures with the following layout:
+ // https://arrow.apache.org/docs/format/Columnar.html#variable-size-binary-view-layout
+ //
+ // * Short strings, length <= 12
+ // | Bytes 0-3 | Bytes 4-15 |
+ // |------------|---------------------------------------|
+ // | length | data (padded with 0) |
+ //
+ // * Long strings, length > 12
+ // | Bytes 0-3 | Bytes 4-7 | Bytes 8-11 | Bytes 12-15 |
+ // |------------|------------|------------|-------------|
+ // | length | prefix | buf. index | offset |
+ let views = &array_data.buffer::<u128>(0)[..array_data.len()];
Review Comment:
Should this take into account the offset?
```suggestion
let views = &array_data.buffer::<u128>(0)[array_data.offset()..array_data.len()];
```
##########
rust/sedona-spatial-join/src/utils/arrow_utils.rs:
##########
@@ -0,0 +1,166 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::{Array, ArrayData, RecordBatch};
+use arrow_array::ArrayRef;
+use arrow_schema::{ArrowError, DataType};
+use datafusion_common::Result;
+
+/// Estimate the in-memory size of a given RecordBatch. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_record_batch_memory_size(batch: &RecordBatch) -> Result<usize> {
+ let mut total_size = 0;
+
+ for array in batch.columns() {
+ let array_data = array.to_data();
+ total_size += count_array_data_memory_size(&array_data)?;
+ }
+
+ Ok(total_size)
+}
+
+/// Estimate the in-memory size of a given Arrow array. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
Review Comment:
```suggestion
/// Estimate the in-memory size of a given Arrow array. This function estimates the
/// size as if the underlying buffers were copied to somewhere else and not shared,
/// including the sizes of each BinaryView item (which is otherwise not counted by
/// `array_data.get_slice_memory_size()`).
```
##########
rust/sedona-spatial-join/src/utils/arrow_utils.rs:
##########
@@ -0,0 +1,166 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::{Array, ArrayData, RecordBatch};
+use arrow_array::ArrayRef;
+use arrow_schema::{ArrowError, DataType};
+use datafusion_common::Result;
+
+/// Estimate the in-memory size of a given RecordBatch. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_record_batch_memory_size(batch: &RecordBatch) -> Result<usize> {
+ let mut total_size = 0;
+
+ for array in batch.columns() {
+ let array_data = array.to_data();
+ total_size += count_array_data_memory_size(&array_data)?;
+ }
+
+ Ok(total_size)
+}
+
+/// Estimate the in-memory size of a given Arrow array. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_array_memory_size(array: &ArrayRef) -> Result<usize> {
+ let array_data = array.to_data();
+ let size = count_array_data_memory_size(&array_data)?;
+ Ok(size)
+}
+
+/// The maximum number of bytes that can be stored inline in a byte view.
+///
+/// See [`ByteView`] and [`GenericByteViewArray`] for more information on the
+/// layout of the views.
+///
+/// [`GenericByteViewArray`]: https://docs.rs/arrow/latest/arrow/array/struct.GenericByteViewArray.html
+pub const MAX_INLINE_VIEW_LEN: u32 = 12;
+
+/// Count the memory usage of `array_data` and its children recursively.
+fn count_array_data_memory_size(array_data: &ArrayData) -> core::result::Result<usize, ArrowError> {
+ Ok(get_binary_view_value_size(array_data)? + array_data.get_slice_memory_size()?)
+}
+
+fn get_binary_view_value_size(array_data: &ArrayData) -> Result<usize, ArrowError> {
+ let mut result: usize = 0;
+ let array_data_type = array_data.data_type();
+
+ if matches!(array_data_type, DataType::BinaryView | DataType::Utf8View) {
Review Comment:
Got it! I suggested a comment so that future me also doesn't think that 🙂
##########
rust/sedona-spatial-join/src/utils/arrow_utils.rs:
##########
@@ -0,0 +1,166 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::{Array, ArrayData, RecordBatch};
+use arrow_array::ArrayRef;
+use arrow_schema::{ArrowError, DataType};
+use datafusion_common::Result;
+
+/// Estimate the in-memory size of a given RecordBatch. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_record_batch_memory_size(batch: &RecordBatch) -> Result<usize> {
+ let mut total_size = 0;
+
+ for array in batch.columns() {
+ let array_data = array.to_data();
+ total_size += count_array_data_memory_size(&array_data)?;
+ }
+
+ Ok(total_size)
+}
+
+/// Estimate the in-memory size of a given Arrow array. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_array_memory_size(array: &ArrayRef) -> Result<usize> {
+ let array_data = array.to_data();
+ let size = count_array_data_memory_size(&array_data)?;
+ Ok(size)
+}
+
+/// The maximum number of bytes that can be stored inline in a byte view.
+///
+/// See [`ByteView`] and [`GenericByteViewArray`] for more information on the
+/// layout of the views.
+///
+/// [`GenericByteViewArray`]: https://docs.rs/arrow/latest/arrow/array/struct.GenericByteViewArray.html
+pub const MAX_INLINE_VIEW_LEN: u32 = 12;
+
+/// Count the memory usage of `array_data` and its children recursively.
+fn count_array_data_memory_size(array_data: &ArrayData) -> core::result::Result<usize, ArrowError> {
+ Ok(get_binary_view_value_size(array_data)? + array_data.get_slice_memory_size()?)
+}
+
+fn get_binary_view_value_size(array_data: &ArrayData) -> Result<usize, ArrowError> {
+ let mut result: usize = 0;
+ let array_data_type = array_data.data_type();
+
+ if matches!(array_data_type, DataType::BinaryView | DataType::Utf8View) {
+ // The views buffer contains length view structures with the following layout:
+ // https://arrow.apache.org/docs/format/Columnar.html#variable-size-binary-view-layout
+ //
+ // * Short strings, length <= 12
+ // | Bytes 0-3 | Bytes 4-15 |
+ // |------------|---------------------------------------|
+ // | length | data (padded with 0) |
+ //
+ // * Long strings, length > 12
+ // | Bytes 0-3 | Bytes 4-7 | Bytes 8-11 | Bytes 12-15 |
+ // |------------|------------|------------|-------------|
+ // | length | prefix | buf. index | offset |
+ let views = &array_data.buffer::<u128>(0)[..array_data.len()];
+ result = views
+ .iter()
+ .map(|v| {
+ let len = *v as u32;
+ if len > MAX_INLINE_VIEW_LEN {
+ len as usize
+ } else {
+ 0
+ }
+ })
+ .sum();
+ }
+
+ for child in array_data.child_data() {
+ result += get_binary_view_value_size(child)?;
+ }
+ Ok(result)
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow_array::{BinaryViewArray, StringViewArray, StructArray};
+ use arrow_schema::{DataType, Field};
+ use std::sync::Arc;
+
+ #[test]
+ fn test_string_view_array_memory_size() {
+ let array = StringViewArray::from(vec![
+ "short", // Inline
+ "Long string that is definitely longer than 12 bytes", // 51 bytes
+ ]);
+ let array_ref: ArrayRef = Arc::new(array);
+ let size = get_array_memory_size(&array_ref).unwrap();
+ // Views: 2 * 16 = 32 bytes
+ // Data: 51 bytes
+ // Total: 83 bytes
+ assert_eq!(size, 83);
+ }
+
+ #[test]
+ fn test_binary_view_array_memory_size() {
+ let array = BinaryViewArray::from(vec![
+ "short".as_bytes(),
+ "Long string that is definitely longer than 12 bytes".as_bytes(),
+ ]);
+ let array_ref: ArrayRef = Arc::new(array);
+ let size = get_array_memory_size(&array_ref).unwrap();
+ assert_eq!(size, 83);
+ }
+
+ #[test]
+ fn test_struct_array_with_view_memory_size() {
+ let string_view_array = StringViewArray::from(vec![
+ "short",
+ "Long string that is definitely longer than 12 bytes",
+ ]);
+ let boolean_array = arrow_array::BooleanArray::from(vec![true, false]);
+
+ let struct_array = StructArray::from(vec![
+ (
+ Arc::new(Field::new("a", DataType::Utf8View, false)),
+ Arc::new(string_view_array) as ArrayRef,
+ ),
+ (
+ Arc::new(Field::new("b", DataType::Boolean, false)),
+ Arc::new(boolean_array) as ArrayRef,
+ ),
+ ]);
+
+ let array_ref: ArrayRef = Arc::new(struct_array);
+ let size = get_array_memory_size(&array_ref).unwrap();
+ // 83 (StringView) + 1 (Boolean values) = 84
+ assert_eq!(size, 84);
+ }
+
+ #[test]
+ fn test_sliced_view_array_memory_size() {
+ let array = StringViewArray::from(vec![
+ "short",
+ "Long string that is definitely longer than 12 bytes",
+ "Another long string to make buffer larger",
+ ]);
+ let sliced = array.slice(0, 2);
Review Comment:
We can also test `array.slice(2, 1)` to test the non-zero offset case
##########
rust/sedona-spatial-join/src/utils/arrow_utils.rs:
##########
@@ -0,0 +1,166 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::{Array, ArrayData, RecordBatch};
+use arrow_array::ArrayRef;
+use arrow_schema::{ArrowError, DataType};
+use datafusion_common::Result;
+
+/// Estimate the in-memory size of a given RecordBatch. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_record_batch_memory_size(batch: &RecordBatch) -> Result<usize> {
+ let mut total_size = 0;
+
+ for array in batch.columns() {
+ let array_data = array.to_data();
+ total_size += count_array_data_memory_size(&array_data)?;
+ }
+
+ Ok(total_size)
+}
+
+/// Estimate the in-memory size of a given Arrow array. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_array_memory_size(array: &ArrayRef) -> Result<usize> {
+ let array_data = array.to_data();
+ let size = count_array_data_memory_size(&array_data)?;
+ Ok(size)
+}
+
+/// The maximum number of bytes that can be stored inline in a byte view.
+///
+/// See [`ByteView`] and [`GenericByteViewArray`] for more information on the
+/// layout of the views.
+///
+/// [`GenericByteViewArray`]: https://docs.rs/arrow/latest/arrow/array/struct.GenericByteViewArray.html
+pub const MAX_INLINE_VIEW_LEN: u32 = 12;
+
+/// Count the memory usage of `array_data` and its children recursively.
+fn count_array_data_memory_size(array_data: &ArrayData) -> core::result::Result<usize, ArrowError> {
+ Ok(get_binary_view_value_size(array_data)? + array_data.get_slice_memory_size()?)
+}
+
+fn get_binary_view_value_size(array_data: &ArrayData) -> Result<usize, ArrowError> {
+ let mut result: usize = 0;
+ let array_data_type = array_data.data_type();
+
+ if matches!(array_data_type, DataType::BinaryView | DataType::Utf8View) {
+ // The views buffer contains length view structures with the following layout:
+ // https://arrow.apache.org/docs/format/Columnar.html#variable-size-binary-view-layout
+ //
+ // * Short strings, length <= 12
+ // | Bytes 0-3 | Bytes 4-15 |
+ // |------------|---------------------------------------|
+ // | length | data (padded with 0) |
+ //
+ // * Long strings, length > 12
+ // | Bytes 0-3 | Bytes 4-7 | Bytes 8-11 | Bytes 12-15 |
+ // |------------|------------|------------|-------------|
+ // | length | prefix | buf. index | offset |
+ let views = &array_data.buffer::<u128>(0)[..array_data.len()];
+ result = views
+ .iter()
+ .map(|v| {
+ let len = *v as u32;
+ if len > MAX_INLINE_VIEW_LEN {
+ len as usize
+ } else {
+ 0
+ }
+ })
+ .sum();
+ }
+
+ for child in array_data.child_data() {
+ result += get_binary_view_value_size(child)?;
+ }
Review Comment:
Possibly worth a comment since it wasn't clear to me initially what was going on:
```suggestion
// If this was not a BinaryView/Utf8View array, count the bytes of any BinaryView/Utf8View
// children, taking into account the slice of this array that applies to the child.
for child in array_data.child_data() {
result += get_binary_view_value_size(child)?;
}
```
...or change the comment to acknowledge that we aren't taking into account the slice of the child we are dealing with, which is super annoying math and is very tricky to get right with little payoff (probably not that many sliced list arrays are coming through this function).
##########
rust/sedona-spatial-join/src/utils/arrow_utils.rs:
##########
@@ -0,0 +1,166 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::{Array, ArrayData, RecordBatch};
+use arrow_array::ArrayRef;
+use arrow_schema::{ArrowError, DataType};
+use datafusion_common::Result;
+
+/// Estimate the in-memory size of a given RecordBatch. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_record_batch_memory_size(batch: &RecordBatch) -> Result<usize> {
+ let mut total_size = 0;
+
+ for array in batch.columns() {
+ let array_data = array.to_data();
+ total_size += count_array_data_memory_size(&array_data)?;
+ }
+
+ Ok(total_size)
+}
+
+/// Estimate the in-memory size of a given Arrow array. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_array_memory_size(array: &ArrayRef) -> Result<usize> {
+ let array_data = array.to_data();
+ let size = count_array_data_memory_size(&array_data)?;
+ Ok(size)
+}
+
+/// The maximum number of bytes that can be stored inline in a byte view.
+///
+/// See [`ByteView`] and [`GenericByteViewArray`] for more information on the
+/// layout of the views.
+///
+/// [`GenericByteViewArray`]: https://docs.rs/arrow/latest/arrow/array/struct.GenericByteViewArray.html
+pub const MAX_INLINE_VIEW_LEN: u32 = 12;
+
+/// Count the memory usage of `array_data` and its children recursively.
+fn count_array_data_memory_size(array_data: &ArrayData) -> core::result::Result<usize, ArrowError> {
+ Ok(get_binary_view_value_size(array_data)? + array_data.get_slice_memory_size()?)
+}
+
+fn get_binary_view_value_size(array_data: &ArrayData) -> Result<usize, ArrowError> {
+ let mut result: usize = 0;
+ let array_data_type = array_data.data_type();
+
+ if matches!(array_data_type, DataType::BinaryView | DataType::Utf8View) {
+ // The views buffer contains length view structures with the following layout:
+ // https://arrow.apache.org/docs/format/Columnar.html#variable-size-binary-view-layout
+ //
+ // * Short strings, length <= 12
+ // | Bytes 0-3 | Bytes 4-15 |
+ // |------------|---------------------------------------|
+ // | length | data (padded with 0) |
+ //
+ // * Long strings, length > 12
+ // | Bytes 0-3 | Bytes 4-7 | Bytes 8-11 | Bytes 12-15 |
+ // |------------|------------|------------|-------------|
+ // | length | prefix | buf. index | offset |
+ let views = &array_data.buffer::<u128>(0)[..array_data.len()];
+ result = views
+ .iter()
+ .map(|v| {
+ let len = *v as u32;
+ if len > MAX_INLINE_VIEW_LEN {
+ len as usize
+ } else {
+ 0
+ }
+ })
+ .sum();
+ }
+
+ for child in array_data.child_data() {
+ result += get_binary_view_value_size(child)?;
+ }
+ Ok(result)
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow_array::{BinaryViewArray, StringViewArray, StructArray};
+ use arrow_schema::{DataType, Field};
+ use std::sync::Arc;
+
+ #[test]
+ fn test_string_view_array_memory_size() {
+ let array = StringViewArray::from(vec![
+ "short", // Inline
+ "Long string that is definitely longer than 12 bytes", // 51 bytes
+ ]);
+ let array_ref: ArrayRef = Arc::new(array);
+ let size = get_array_memory_size(&array_ref).unwrap();
+ // Views: 2 * 16 = 32 bytes
+ // Data: 51 bytes
+ // Total: 83 bytes
+ assert_eq!(size, 83);
+ }
+
+ #[test]
+ fn test_binary_view_array_memory_size() {
+ let array = BinaryViewArray::from(vec![
+ "short".as_bytes(),
+ "Long string that is definitely longer than 12 bytes".as_bytes(),
+ ]);
+ let array_ref: ArrayRef = Arc::new(array);
+ let size = get_array_memory_size(&array_ref).unwrap();
+ assert_eq!(size, 83);
+ }
+
+ #[test]
+ fn test_struct_array_with_view_memory_size() {
+ let string_view_array = StringViewArray::from(vec![
+ "short",
+ "Long string that is definitely longer than 12 bytes",
+ ]);
+ let boolean_array = arrow_array::BooleanArray::from(vec![true, false]);
+
+ let struct_array = StructArray::from(vec![
+ (
+ Arc::new(Field::new("a", DataType::Utf8View, false)),
+ Arc::new(string_view_array) as ArrayRef,
+ ),
+ (
+ Arc::new(Field::new("b", DataType::Boolean, false)),
+ Arc::new(boolean_array) as ArrayRef,
+ ),
+ ]);
+
+ let array_ref: ArrayRef = Arc::new(struct_array);
+ let size = get_array_memory_size(&array_ref).unwrap();
+ // 83 (StringView) + 1 (Boolean values) = 84
+ assert_eq!(size, 84);
Review Comment:
I think `array_ref.slice(0, 1)` will give the wrong answer with the current code (the slice is not projected into the children)
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] chore(rust/sedona-spatial-join): More accurate batch in-memory size estimation [sedona-db]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #515:
URL: https://github.com/apache/sedona-db/pull/515#discussion_r2694835218
##########
rust/sedona-spatial-join/src/utils/arrow_utils.rs:
##########
@@ -0,0 +1,166 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::{Array, ArrayData, RecordBatch};
+use arrow_array::ArrayRef;
+use arrow_schema::{ArrowError, DataType};
+use datafusion_common::Result;
+
+/// Estimate the in-memory size of a given RecordBatch. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_record_batch_memory_size(batch: &RecordBatch) -> Result<usize> {
+ let mut total_size = 0;
+
+ for array in batch.columns() {
+ let array_data = array.to_data();
+ total_size += count_array_data_memory_size(&array_data)?;
+ }
+
+ Ok(total_size)
+}
+
+/// Estimate the in-memory size of a given Arrow array. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_array_memory_size(array: &ArrayRef) -> Result<usize> {
+ let array_data = array.to_data();
+ let size = count_array_data_memory_size(&array_data)?;
+ Ok(size)
+}
+
+/// The maximum number of bytes that can be stored inline in a byte view.
+///
+/// See [`ByteView`] and [`GenericByteViewArray`] for more information on the
+/// layout of the views.
+///
+/// [`GenericByteViewArray`]: https://docs.rs/arrow/latest/arrow/array/struct.GenericByteViewArray.html
+pub const MAX_INLINE_VIEW_LEN: u32 = 12;
+
+/// Count the memory usage of `array_data` and its children recursively.
+fn count_array_data_memory_size(array_data: &ArrayData) -> core::result::Result<usize, ArrowError> {
+ Ok(get_binary_view_value_size(array_data)? + array_data.get_slice_memory_size()?)
+}
+
+fn get_binary_view_value_size(array_data: &ArrayData) -> Result<usize, ArrowError> {
+ let mut result: usize = 0;
+ let array_data_type = array_data.data_type();
+
+ if matches!(array_data_type, DataType::BinaryView | DataType::Utf8View) {
+ // The views buffer contains length view structures with the following layout:
+ // https://arrow.apache.org/docs/format/Columnar.html#variable-size-binary-view-layout
+ //
+ // * Short strings, length <= 12
+ // | Bytes 0-3 | Bytes 4-15 |
+ // |------------|---------------------------------------|
+ // | length | data (padded with 0) |
+ //
+ // * Long strings, length > 12
+ // | Bytes 0-3 | Bytes 4-7 | Bytes 8-11 | Bytes 12-15 |
+ // |------------|------------|------------|-------------|
+ // | length | prefix | buf. index | offset |
+ let views = &array_data.buffer::<u128>(0)[..array_data.len()];
+ result = views
+ .iter()
+ .map(|v| {
+ let len = *v as u32;
+ if len > MAX_INLINE_VIEW_LEN {
+ len as usize
+ } else {
+ 0
+ }
+ })
+ .sum();
+ }
+
+ for child in array_data.child_data() {
+ result += get_binary_view_value_size(child)?;
+ }
+ Ok(result)
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow_array::{BinaryViewArray, StringViewArray, StructArray};
+ use arrow_schema::{DataType, Field};
+ use std::sync::Arc;
+
+ #[test]
+ fn test_string_view_array_memory_size() {
+ let array = StringViewArray::from(vec![
+ "short", // Inline
+ "Long string that is definitely longer than 12 bytes", // 51 bytes
+ ]);
+ let array_ref: ArrayRef = Arc::new(array);
+ let size = get_array_memory_size(&array_ref).unwrap();
+ // Views: 2 * 16 = 32 bytes
+ // Data: 51 bytes
+ // Total: 83 bytes
+ assert_eq!(size, 83);
+ }
+
+ #[test]
+ fn test_binary_view_array_memory_size() {
+ let array = BinaryViewArray::from(vec![
+ "short".as_bytes(),
+ "Long string that is definitely longer than 12 bytes".as_bytes(),
+ ]);
+ let array_ref: ArrayRef = Arc::new(array);
+ let size = get_array_memory_size(&array_ref).unwrap();
+ assert_eq!(size, 83);
+ }
+
+ #[test]
+ fn test_struct_array_with_view_memory_size() {
+ let string_view_array = StringViewArray::from(vec![
+ "short",
+ "Long string that is definitely longer than 12 bytes",
+ ]);
+ let boolean_array = arrow_array::BooleanArray::from(vec![true, false]);
+
+ let struct_array = StructArray::from(vec![
+ (
+ Arc::new(Field::new("a", DataType::Utf8View, false)),
+ Arc::new(string_view_array) as ArrayRef,
+ ),
+ (
+ Arc::new(Field::new("b", DataType::Boolean, false)),
+ Arc::new(boolean_array) as ArrayRef,
+ ),
+ ]);
+
+ let array_ref: ArrayRef = Arc::new(struct_array);
+ let size = get_array_memory_size(&array_ref).unwrap();
+ // 83 (StringView) + 1 (Boolean values) = 84
+ assert_eq!(size, 84);
Review Comment:
Ah, arrow-rs propagates the slice for a StructArray on `slice()`: https://docs.rs/arrow-array/57.2.0/src/arrow_array/array/struct_array.rs.html#337
It doesn't do this for list arrays: https://docs.rs/arrow-array/57.2.0/src/arrow_array/array/list_array.rs.html#396
...and possibly not for some other types like dictionaries, list views, or run-end-encoded. I think it's OK to punt on all of that but perhaps open an issue and link to it here in case this comes up.
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] chore(rust/sedona-spatial-join): More accurate batch in-memory size estimation [sedona-db]
Posted by "Kontinuation (via GitHub)" <gi...@apache.org>.
Kontinuation commented on code in PR #515:
URL: https://github.com/apache/sedona-db/pull/515#discussion_r2694869758
##########
rust/sedona-spatial-join/src/utils/arrow_utils.rs:
##########
@@ -0,0 +1,166 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::{Array, ArrayData, RecordBatch};
+use arrow_array::ArrayRef;
+use arrow_schema::{ArrowError, DataType};
+use datafusion_common::Result;
+
+/// Estimate the in-memory size of a given RecordBatch. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_record_batch_memory_size(batch: &RecordBatch) -> Result<usize> {
+ let mut total_size = 0;
+
+ for array in batch.columns() {
+ let array_data = array.to_data();
+ total_size += count_array_data_memory_size(&array_data)?;
+ }
+
+ Ok(total_size)
+}
+
+/// Estimate the in-memory size of a given Arrow array. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_array_memory_size(array: &ArrayRef) -> Result<usize> {
+ let array_data = array.to_data();
+ let size = count_array_data_memory_size(&array_data)?;
+ Ok(size)
+}
+
+/// The maximum number of bytes that can be stored inline in a byte view.
+///
+/// See [`ByteView`] and [`GenericByteViewArray`] for more information on the
+/// layout of the views.
+///
+/// [`GenericByteViewArray`]: https://docs.rs/arrow/latest/arrow/array/struct.GenericByteViewArray.html
+pub const MAX_INLINE_VIEW_LEN: u32 = 12;
+
+/// Count the memory usage of `array_data` and its children recursively.
+fn count_array_data_memory_size(array_data: &ArrayData) -> core::result::Result<usize, ArrowError> {
+ Ok(get_binary_view_value_size(array_data)? + array_data.get_slice_memory_size()?)
+}
+
+fn get_binary_view_value_size(array_data: &ArrayData) -> Result<usize, ArrowError> {
+ let mut result: usize = 0;
+ let array_data_type = array_data.data_type();
+
+ if matches!(array_data_type, DataType::BinaryView | DataType::Utf8View) {
+ // The views buffer contains length view structures with the following layout:
+ // https://arrow.apache.org/docs/format/Columnar.html#variable-size-binary-view-layout
+ //
+ // * Short strings, length <= 12
+ // | Bytes 0-3 | Bytes 4-15 |
+ // |------------|---------------------------------------|
+ // | length | data (padded with 0) |
+ //
+ // * Long strings, length > 12
+ // | Bytes 0-3 | Bytes 4-7 | Bytes 8-11 | Bytes 12-15 |
+ // |------------|------------|------------|-------------|
+ // | length | prefix | buf. index | offset |
+ let views = &array_data.buffer::<u128>(0)[..array_data.len()];
+ result = views
+ .iter()
+ .map(|v| {
+ let len = *v as u32;
+ if len > MAX_INLINE_VIEW_LEN {
+ len as usize
+ } else {
+ 0
+ }
+ })
+ .sum();
+ }
+
+ for child in array_data.child_data() {
+ result += get_binary_view_value_size(child)?;
+ }
+ Ok(result)
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow_array::{BinaryViewArray, StringViewArray, StructArray};
+ use arrow_schema::{DataType, Field};
+ use std::sync::Arc;
+
+ #[test]
+ fn test_string_view_array_memory_size() {
+ let array = StringViewArray::from(vec![
+ "short", // Inline
+ "Long string that is definitely longer than 12 bytes", // 51 bytes
+ ]);
+ let array_ref: ArrayRef = Arc::new(array);
+ let size = get_array_memory_size(&array_ref).unwrap();
+ // Views: 2 * 16 = 32 bytes
+ // Data: 51 bytes
+ // Total: 83 bytes
+ assert_eq!(size, 83);
+ }
+
+ #[test]
+ fn test_binary_view_array_memory_size() {
+ let array = BinaryViewArray::from(vec![
+ "short".as_bytes(),
+ "Long string that is definitely longer than 12 bytes".as_bytes(),
+ ]);
+ let array_ref: ArrayRef = Arc::new(array);
+ let size = get_array_memory_size(&array_ref).unwrap();
+ assert_eq!(size, 83);
+ }
+
+ #[test]
+ fn test_struct_array_with_view_memory_size() {
+ let string_view_array = StringViewArray::from(vec![
+ "short",
+ "Long string that is definitely longer than 12 bytes",
+ ]);
+ let boolean_array = arrow_array::BooleanArray::from(vec![true, false]);
+
+ let struct_array = StructArray::from(vec![
+ (
+ Arc::new(Field::new("a", DataType::Utf8View, false)),
+ Arc::new(string_view_array) as ArrayRef,
+ ),
+ (
+ Arc::new(Field::new("b", DataType::Boolean, false)),
+ Arc::new(boolean_array) as ArrayRef,
+ ),
+ ]);
+
+ let array_ref: ArrayRef = Arc::new(struct_array);
+ let size = get_array_memory_size(&array_ref).unwrap();
+ // 83 (StringView) + 1 (Boolean values) = 84
+ assert_eq!(size, 84);
Review Comment:
I'd like to leave `get_array_data_memory_size` as is and log a ticket. https://github.com/apache/sedona-db/issues/519
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] chore(rust/sedona-spatial-join): More accurate batch in-memory size estimation [sedona-db]
Posted by "Kontinuation (via GitHub)" <gi...@apache.org>.
Kontinuation commented on code in PR #515:
URL: https://github.com/apache/sedona-db/pull/515#discussion_r2694761640
##########
rust/sedona-spatial-join/src/utils/arrow_utils.rs:
##########
@@ -0,0 +1,166 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::array::{Array, ArrayData, RecordBatch};
+use arrow_array::ArrayRef;
+use arrow_schema::{ArrowError, DataType};
+use datafusion_common::Result;
+
+/// Estimate the in-memory size of a given RecordBatch. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_record_batch_memory_size(batch: &RecordBatch) -> Result<usize> {
+ let mut total_size = 0;
+
+ for array in batch.columns() {
+ let array_data = array.to_data();
+ total_size += count_array_data_memory_size(&array_data)?;
+ }
+
+ Ok(total_size)
+}
+
+/// Estimate the in-memory size of a given Arrow array. This function estimates the
+/// size as if the underlying buffers were copied to somewhere else and not shared.
+pub(crate) fn get_array_memory_size(array: &ArrayRef) -> Result<usize> {
+ let array_data = array.to_data();
+ let size = count_array_data_memory_size(&array_data)?;
+ Ok(size)
+}
+
+/// The maximum number of bytes that can be stored inline in a byte view.
+///
+/// See [`ByteView`] and [`GenericByteViewArray`] for more information on the
+/// layout of the views.
+///
+/// [`GenericByteViewArray`]: https://docs.rs/arrow/latest/arrow/array/struct.GenericByteViewArray.html
+pub const MAX_INLINE_VIEW_LEN: u32 = 12;
+
+/// Count the memory usage of `array_data` and its children recursively.
+fn count_array_data_memory_size(array_data: &ArrayData) -> core::result::Result<usize, ArrowError> {
+ Ok(get_binary_view_value_size(array_data)? + array_data.get_slice_memory_size()?)
+}
+
+fn get_binary_view_value_size(array_data: &ArrayData) -> Result<usize, ArrowError> {
+ let mut result: usize = 0;
+ let array_data_type = array_data.data_type();
+
+ if matches!(array_data_type, DataType::BinaryView | DataType::Utf8View) {
Review Comment:
Comment was added to `get_array_data_memory_size`.
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org