Merge branch 'databendlabs:main' into add-glue-catalog

src/common/arrow/src/arrow/array/binview/mod.rs

@@ -166,6 +166,18 @@ impl<T: ViewType + ?Sized> BinaryViewArrayGeneric<T> {
         total_bytes_len: usize,
         total_buffer_len: usize,
     ) -> Self {
+        #[cfg(debug_assertions)]
+        {
+            if total_bytes_len != UNKNOWN_LEN as usize {
+                let total = views.iter().map(|v| v.length as usize).sum::<usize>();
+                assert_eq!(total, total_bytes_len);
+            }
+
+            if total_buffer_len != UNKNOWN_LEN as usize {
+                let total = buffers.iter().map(|v| v.len()).sum::<usize>();
+                assert_eq!(total, total_buffer_len);
+            }
+        }
         // # Safety
         // The caller must ensure
         // - the data is valid utf8 (if required)

src/common/arrow/src/arrow/array/binview/mutable.rs

@@ -155,12 +155,11 @@ impl<T: ViewType + ?Sized> MutableBinaryViewArray<T> {
     #[inline]
     pub(crate) unsafe fn push_view_unchecked(&mut self, v: View, buffers: &[Buffer<u8>]) {
         let len = v.length;
-        self.total_bytes_len += len as usize;
         if len <= 12 {
+            self.total_bytes_len += len as usize;
             debug_assert!(self.views.capacity() > self.views.len());
             self.views.push(v)
         } else {
-            self.total_buffer_len += len as usize;
             let data = buffers.get_unchecked(v.buffer_idx as usize);
             let offset = v.offset as usize;
             let bytes = data.get_unchecked(offset..offset + len as usize);
@@ -263,12 +262,19 @@ impl<T: ViewType + ?Sized> MutableBinaryViewArray<T> {
         // Push and pop to get the properly encoded value.
         // For long string this leads to a dictionary encoding,
         // as we push the string only once in the buffers
+
+        let old_bytes_len = self.total_bytes_len;
+
         let view_value = value
             .map(|v| {
                 self.push_value_ignore_validity(v);
                 self.views.pop().unwrap()
             })
             .unwrap_or_default();
+
+        self.total_bytes_len +=
+            (self.total_bytes_len - old_bytes_len) * additional.saturating_sub(1);
+
         self.views
             .extend(std::iter::repeat(view_value).take(additional));
     }

src/common/arrow/src/arrow/array/growable/binview.rs

@@ -74,7 +74,6 @@ impl<'a, T: ViewType + ?Sized> GrowableBinaryViewArray<'a, T> {
         capacity: usize,
     ) -> Self {
         let data_type = arrays[0].data_type().clone();
-
         // if any of the arrays has nulls, insertions from any array requires setting bits
         // as there is at least one array with nulls.
         if !use_validity & arrays.iter().any(|array| array.null_count() > 0) {
@@ -91,11 +90,8 @@ impl<'a, T: ViewType + ?Sized> GrowableBinaryViewArray<'a, T> {
                     .map(|buf| BufferKey { inner: buf })
             })
             .collect::<ArrowIndexSet<_>>();
-        let total_buffer_len = arrays
-            .iter()
-            .map(|arr| arr.data_buffers().len())
-            .sum::<usize>();
 
+        let total_buffer_len = buffers.iter().map(|v| v.inner.len()).sum();
         Self {
             arrays,
             data_type,

src/common/arrow/tests/it/arrow/array/binview/mutable_values.rs

@@ -29,3 +29,17 @@ fn extend_from_iter() {
             .as_box()
     )
 }
+
+#[test]
+fn extend_from_repeats() {
+    let mut b = MutableBinaryViewArray::<str>::new();
+    b.extend_constant(4, Some("databend"));
+
+    let a = b.clone();
+    b.extend_trusted_len_values(a.values_iter());
+
+    assert_eq!(
+        b.as_box(),
+        MutableBinaryViewArray::<str>::from_values_iter(vec!["databend"; 8].into_iter()).as_box()
+    )
+}

src/common/base/Cargo.toml

@@ -26,6 +26,7 @@ databend-common-exception = { workspace = true }
 async-backtrace = { workspace = true }
 async-trait = { workspace = true }
 borsh = { workspace = true }
+bytes = { workspace = true }
 bytesize = { workspace = true }
 chrono = { workspace = true }
 ctrlc = { workspace = true }

src/common/base/src/base/dma.rs

@@ -28,6 +28,7 @@ use std::path::Path;
 use std::ptr;
 use std::ptr::NonNull;
 
+use bytes::Bytes;
 use rustix::fs::OFlags;
 use tokio::fs::File;
 use tokio::io::AsyncSeekExt;
@@ -116,10 +117,6 @@ impl DmaAllocator {
             Layout::from_size_align(layout.size(), self.0.as_usize()).unwrap()
         }
     }
-
-    fn real_cap(&self, cap: usize) -> usize {
-        self.0.align_up(cap)
-    }
 }
 
 unsafe impl Allocator for DmaAllocator {
@@ -131,6 +128,10 @@ unsafe impl Allocator for DmaAllocator {
         Global {}.allocate_zeroed(self.real_layout(layout))
     }
 
+    unsafe fn deallocate(&self, ptr: std::ptr::NonNull<u8>, layout: Layout) {
+        Global {}.deallocate(ptr, self.real_layout(layout))
+    }
+
     unsafe fn grow(
         &self,
         ptr: NonNull<u8>,
@@ -157,20 +158,38 @@ unsafe impl Allocator for DmaAllocator {
         )
     }
 
-    unsafe fn deallocate(&self, ptr: std::ptr::NonNull<u8>, layout: Layout) {
-        Global {}.deallocate(ptr, self.real_layout(layout))
+    unsafe fn shrink(
+        &self,
+        ptr: NonNull<u8>,
+        old_layout: Layout,
+        new_layout: Layout,
+    ) -> Result<NonNull<[u8]>, AllocError> {
+        Global {}.shrink(
+            ptr,
+            self.real_layout(old_layout),
+            self.real_layout(new_layout),
+        )
     }
 }
 
 type DmaBuffer = Vec<u8, DmaAllocator>;
 
-pub fn dma_buffer_as_vec(mut buf: DmaBuffer) -> Vec<u8> {
-    let ptr = buf.as_mut_ptr();
-    let len = buf.len();
-    let cap = buf.allocator().real_cap(buf.capacity());
-    std::mem::forget(buf);
-
-    unsafe { Vec::from_raw_parts(ptr, len, cap) }
+pub fn dma_buffer_to_bytes(buf: DmaBuffer) -> Bytes {
+    if buf.is_empty() {
+        return Bytes::new();
+    }
+    let (ptr, len, cap, alloc) = buf.into_raw_parts_with_alloc();
+    // Memory fitting
+    let old_layout = Layout::from_size_align(cap, alloc.0.as_usize()).unwrap();
+    let new_layout = Layout::from_size_align(len, std::mem::align_of::<u8>()).unwrap();
+    let data = unsafe {
+        let p = Global {}
+            .shrink(NonNull::new(ptr).unwrap(), old_layout, new_layout)
+            .unwrap();
+        let cap = p.len();
+        Vec::from_raw_parts(p.cast().as_mut(), len, cap)
+    };
+    Bytes::from(data)
 }
 
 /// A `DmaFile` is similar to a `File`, but it is opened with the `O_DIRECT` file in order to
@@ -697,4 +716,28 @@ mod tests {
 
         let _ = std::fs::remove_file(filename);
     }
+
+    #[test]
+    fn test_dma_buffer_to_bytes() {
+        let want = (0..10_u8).collect::<Vec<_>>();
+        let alloc = DmaAllocator::new(Alignment::new(4096).unwrap());
+        let mut buf = DmaBuffer::with_capacity_in(3000, alloc);
+        buf.extend_from_slice(&want);
+
+        println!("{:?} {}", buf.as_ptr(), buf.capacity());
+        buf.shrink_to_fit();
+        println!("{:?} {}", buf.as_ptr(), buf.capacity());
+        buf.reserve(3000 - buf.capacity());
+        println!("{:?} {}", buf.as_ptr(), buf.capacity());
+
+        // let slice = buf.into_boxed_slice();
+        // println!("{:?}", slice.as_ptr());
+
+        let got = dma_buffer_to_bytes(buf);
+        println!("{:?}", got.as_ptr());
+        assert_eq!(&want, &got);
+
+        let buf = got.to_vec();
+        println!("{:?} {}", buf.as_ptr(), buf.capacity());
+    }
 }

src/common/base/src/base/mod.rs

@@ -28,7 +28,7 @@ mod take_mut;
 mod uniq_id;
 mod watch_notify;
 
-pub use dma::dma_buffer_as_vec;
+pub use dma::dma_buffer_to_bytes;
 pub use dma::dma_read_file;
 pub use dma::dma_read_file_range;
 pub use dma::dma_write_file_vectored;

src/common/base/src/lib.rs

@@ -25,6 +25,7 @@
 #![feature(slice_swap_unchecked)]
 #![feature(variant_count)]
 #![feature(ptr_alignment_type)]
+#![feature(vec_into_raw_parts)]
 
 pub mod base;
 pub mod containers;
@@ -35,7 +36,6 @@ pub mod http_client;
 pub mod mem_allocator;
 pub mod rangemap;
 pub mod runtime;
-pub mod slice_ext;
 pub mod vec_ext;
 pub mod version;
 

src/common/storage/Cargo.toml

@@ -34,6 +34,7 @@ regex = { workspace = true }
 reqwest = { workspace = true }
 serde = { workspace = true }
 thiserror = { workspace = true }
+url = { workspace = true }
 
 [dev-dependencies]
 

src/common/storage/src/http_client.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::HashMap;
 use std::future;
 use std::mem;
 use std::str::FromStr;
@@ -26,6 +27,7 @@ use opendal::raw::parse_content_length;
 use opendal::raw::HttpBody;
 use opendal::raw::HttpFetch;
 use opendal::Buffer;
+use url::Url;
 
 pub struct StorageHttpClient {
     client: reqwest::Client,
@@ -46,13 +48,14 @@ impl HttpFetch for StorageHttpClient {
         let uri = req.uri().clone();
         let is_head = req.method() == http::Method::HEAD;
 
-        let host = uri.host().unwrap_or_default();
+        let url = Url::parse(uri.to_string().as_str()).expect("input request url must be valid");
+        let host = url.host_str().unwrap_or_default();
         let method = match req.method() {
             &http::Method::GET => {
-                if uri.path() == "/" {
-                    "LIST"
-                } else {
-                    "GET"
+                let query: HashMap<_, _> = url.query_pairs().collect();
+                match query.get("list-type") {
+                    Some(_) => "LIST",
+                    None => "GET",
                 }
             }
             m => m.as_str(),

src/query/expression/src/kernels/filter.rs

@@ -358,7 +358,7 @@ impl<'a> FilterVisitor<'a> {
                         new_views,
                         values.data.data_buffers().clone(),
                         None,
-                        Some(values.data.total_buffer_len()),
+                        None,
                     )
                 };
                 StringColumn::new(new_col)

src/query/expression/src/kernels/take.rs

@@ -19,7 +19,6 @@ use databend_common_arrow::arrow::array::Array;
 use databend_common_arrow::arrow::array::Utf8ViewArray;
 use databend_common_arrow::arrow::bitmap::Bitmap;
 use databend_common_arrow::arrow::buffer::Buffer;
-use databend_common_base::slice_ext::GetSaferUnchecked;
 use databend_common_exception::Result;
 use string::StringColumnBuilder;
 
@@ -258,7 +257,7 @@ where I: databend_common_arrow::arrow::types::Index
         let result: Vec<T> = self
             .indices
             .iter()
-            .map(|index| unsafe { *col.get_unchecked_release(index.to_usize()) })
+            .map(|index| unsafe { *col.get_unchecked(index.to_usize()) })
             .collect();
         result.into()
     }
@@ -292,7 +291,7 @@ where I: databend_common_arrow::arrow::types::Index
                     new_views,
                     col.data.data_buffers().clone(),
                     None,
-                    Some(col.data.total_buffer_len()),
+                    None,
                 )
             };
             StringColumn::new(new_col)

src/query/expression/src/kernels/take_compact.rs

@@ -238,7 +238,7 @@ impl<'a> TakeCompactVisitor<'a> {
                 new_views,
                 col.data.data_buffers().clone(),
                 None,
-                Some(col.data.total_buffer_len()),
+                None,
             )
         };
         StringColumn::new(new_col)