read ThriftFileMetaData instead of FileMetaData

Author: dantengsky

Cargo.lock

@@ -45,8 +45,8 @@ enable-histogram-metrics = [
 common-arrow = { path = "../../common/arrow" }
 common-ast = { path = "../ast" }
 common-base = { path = "../../common/base" }
+common-cache = { path = "../../common/cache" }
 common-catalog = { path = "../catalog" }
-common-cache = {path = "../../common/cache" }
 common-compress = { path = "../../common/compress" }
 common-config = { path = "../config" }
 common-exception = { path = "../../common/exception" }
@@ -161,13 +161,13 @@ num = "0.4.0"
 p256 = "0.13"
 pretty_assertions = "1.3.0"
 reqwest = { workspace = true }
+sysinfo = "0.28.3"
 temp-env = "0.3.0"
 tempfile = "3.4.0"
 toml = { version = "0.7.3", default-features = false }
 url = "2.3.1"
 walkdir = "2.3.2"
 wiremock = "0.5.14"
-sysinfo = "0.28.3"
 
 [build-dependencies]
 common-building = { path = "../../common/building" }

src/query/service/Cargo.toml

@@ -84,8 +84,7 @@ async fn test_index_meta_cache_size_bloom_meta() -> common_exception::Result<()>
 
     let cache_number = 300_000;
 
-    let meta: FileMetaData = FileMetaData::try_from_thrift(thrift_file_meta)?;
-    let bloom_index_meta = BloomIndexMeta::try_from(meta)?;
+    let bloom_index_meta = BloomIndexMeta::try_from(thrift_file_meta)?;
 
     let sys = System::new_all();
     let pid = get_current_pid().unwrap();

src/query/service/tests/it/storages/fuse/bloom_index_meta_size.rs

@@ -16,7 +16,7 @@ use std::collections::HashMap;
 use std::ops::Deref;
 use std::sync::Arc;
 
-use common_arrow::parquet::metadata::FileMetaData;
+use common_arrow::parquet::metadata::ThriftFileMetaData;
 use common_exception::ErrorCode;
 use common_exception::Result;
 use common_exception::Span;
@@ -60,12 +60,10 @@ pub struct BloomIndexMeta {
     pub columns: Vec<(String, SingleColumnMeta)>,
 }
 
-impl TryFrom<FileMetaData> for BloomIndexMeta {
+impl TryFrom<ThriftFileMetaData> for BloomIndexMeta {
     type Error = common_exception::ErrorCode;
 
-    fn try_from(meta: FileMetaData) -> std::result::Result<Self, Self::Error> {
-        // TODO from FileMeta directly or pass in Thrift
-        let mut meta = meta.into_thrift();
+    fn try_from(mut meta: ThriftFileMetaData) -> std::result::Result<Self, Self::Error> {
         let rg = meta.row_groups.remove(0);
         let mut col_metas = Vec::with_capacity(rg.columns.len());
         for x in &rg.columns {
@@ -92,7 +90,9 @@ impl TryFrom<FileMetaData> for BloomIndexMeta {
                     col_metas.push((column_name, res));
                 }
                 None => {
-                    panic!("")
+                    panic!(
+                        "expecting chunk meta data while converting ThriftFileMetaData to BloomIndexMeta"
+                    )
                 }
             }
         }

src/query/storages/common/index/src/bloom_index.rs

@@ -46,6 +46,7 @@ futures = "0.3.24"
 futures-util = "0.3.24"
 metrics = "0.20.1"
 opendal = { workspace = true }
+parquet-format-safe = "0.2"
 serde = { workspace = true }
 serde_json = { workspace = true }
 sha2 = "0.10.6"

src/query/storages/fuse/Cargo.toml

@@ -12,10 +12,16 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use common_arrow::parquet::read::read_metadata_async;
+use std::io::SeekFrom;
+
+use common_arrow::parquet::metadata::ThriftFileMetaData;
 use common_exception::ErrorCode;
 use common_exception::Result;
 use common_expression::TableSchemaRef;
+use futures::AsyncRead;
+use futures::AsyncSeek;
+use futures_util::AsyncReadExt;
+use futures_util::AsyncSeekExt;
 use opendal::Operator;
 use opendal::Reader;
 use storages_common_cache::InMemoryItemCacheReader;
@@ -31,6 +37,7 @@ use storages_common_table_meta::meta::TableSnapshotStatistics;
 use storages_common_table_meta::meta::TableSnapshotStatisticsVersion;
 
 use super::versioned_reader::VersionedReader;
+use crate::io::read::meta::meta_readers::thrift_file_meta_read::read_thrift_file_metadata;
 
 pub type TableSnapshotStatisticsReader =
     InMemoryItemCacheReader<TableSnapshotStatistics, LoaderWrapper<Operator>>;
@@ -107,12 +114,15 @@ impl Loader<SegmentInfo> for LoaderWrapper<(Operator, TableSchemaRef)> {
 impl Loader<BloomIndexMeta> for LoaderWrapper<Operator> {
     async fn load(&self, params: &LoadParams) -> Result<BloomIndexMeta> {
         let mut reader = bytes_reader(&self.0, params.location.as_str(), params.len_hint).await?;
-        let meta = read_metadata_async(&mut reader).await.map_err(|err| {
-            ErrorCode::Internal(format!(
-                "read file meta failed, {}, {:?}",
-                params.location, err
-            ))
-        })?;
+        // read the ThriftFileMetaData, omit unnecessary conversions
+        let meta = read_thrift_file_metadata(&mut reader)
+            .await
+            .map_err(|err| {
+                ErrorCode::StorageOther(format!(
+                    "read file meta failed, {}, {:?}",
+                    params.location, err
+                ))
+            })?;
 
         BloomIndexMeta::try_from(meta)
     }
@@ -126,3 +136,107 @@ async fn bytes_reader(op: &Operator, path: &str, len_hint: Option<u64>) -> Resul
     };
     Ok(reader)
 }
+
+mod thrift_file_meta_read {
+    use common_arrow::parquet::error::Error;
+    use parquet_format_safe::thrift::protocol::TCompactInputProtocol;
+
+    use super::*;
+
+    // the following code is copied from crate `parquet2`, with slight modification:
+    // return a ThriftFileMetaData instead of FileMetaData while reader parquet metadata,
+    // to avoid unnecessary conversions.
+    //
+    // It takes about 18s, to convert one million ThriftFileMetaData objects to parquet::metadata::FileMetaData,
+    // and the memory usage of FileMetaData is also larger.
+
+    const HEADER_SIZE: u64 = PARQUET_MAGIC.len() as u64;
+    const FOOTER_SIZE: u64 = 8;
+    const PARQUET_MAGIC: [u8; 4] = [b'P', b'A', b'R', b'1'];
+
+    /// The number of bytes read at the end of the parquet file on first read
+    const DEFAULT_FOOTER_READ_SIZE: u64 = 64 * 1024;
+
+    async fn stream_len(
+        seek: &mut (impl AsyncSeek + std::marker::Unpin),
+    ) -> std::result::Result<u64, std::io::Error> {
+        let old_pos = seek.seek(SeekFrom::Current(0)).await?;
+        let len = seek.seek(SeekFrom::End(0)).await?;
+
+        // Avoid seeking a third time when we were already at the end of the
+        // stream. The branch is usually way cheaper than a seek operation.
+        if old_pos != len {
+            seek.seek(SeekFrom::Start(old_pos)).await?;
+        }
+
+        Ok(len)
+    }
+
+    fn metadata_len(buffer: &[u8], len: usize) -> i32 {
+        i32::from_le_bytes(buffer[len - 8..len - 4].try_into().unwrap())
+    }
+
+    pub async fn read_thrift_file_metadata<R: AsyncRead + AsyncSeek + Send + std::marker::Unpin>(
+        reader: &mut R,
+    ) -> common_arrow::parquet::error::Result<ThriftFileMetaData> {
+        let file_size = stream_len(reader).await?;
+
+        if file_size < HEADER_SIZE + FOOTER_SIZE {
+            return Err(Error::OutOfSpec(
+                "A parquet file must containt a header and footer with at least 12 bytes".into(),
+            ));
+        }
+
+        // read and cache up to DEFAULT_FOOTER_READ_SIZE bytes from the end and process the footer
+        let default_end_len = std::cmp::min(DEFAULT_FOOTER_READ_SIZE, file_size) as usize;
+        reader
+            .seek(SeekFrom::End(-(default_end_len as i64)))
+            .await?;
+
+        let mut buffer = vec![];
+        buffer.try_reserve(default_end_len)?;
+        reader
+            .take(default_end_len as u64)
+            .read_to_end(&mut buffer)
+            .await?;
+
+        // check this is indeed a parquet file
+        if buffer[default_end_len - 4..] != PARQUET_MAGIC {
+            return Err(Error::OutOfSpec(
+                "Invalid Parquet file. Corrupt footer".into(),
+            ));
+        }
+
+        let metadata_len = metadata_len(&buffer, default_end_len);
+        let metadata_len: u64 = metadata_len.try_into()?;
+
+        let footer_len = FOOTER_SIZE + metadata_len;
+        if footer_len > file_size {
+            return Err(Error::OutOfSpec(
+                "The footer size must be smaller or equal to the file's size".into(),
+            ));
+        }
+
+        let reader = if (footer_len as usize) < buffer.len() {
+            // the whole metadata is in the bytes we already read
+            let remaining = buffer.len() - footer_len as usize;
+            &buffer[remaining..]
+        } else {
+            // the end of file read by default is not long enough, read again including the metadata.
+            reader.seek(SeekFrom::End(-(footer_len as i64))).await?;
+
+            buffer.clear();
+            buffer.try_reserve(footer_len as usize)?;
+            reader.take(footer_len).read_to_end(&mut buffer).await?;
+
+            &buffer
+        };
+
+        // a highly nested but sparse struct could result in many allocations
+        let max_size = reader.len() * 2 + 1024;
+
+        let mut prot = TCompactInputProtocol::new(reader, max_size);
+        let meta = ThriftFileMetaData::read_from_in_protocol(&mut prot)?;
+        Ok(meta)
+    }
+}