refactor

Author: SkyFan2002

Cargo.lock

@@ -170,8 +170,7 @@ impl DataBlock {
 impl Column {
     pub fn into_arrow_rs(self) -> Arc<dyn arrow_array::Array> {
         let arrow2_array: Box<dyn databend_common_arrow::arrow::array::Array> = self.as_arrow();
-        let arrow_array: Arc<dyn arrow_array::Array> = arrow2_array.into();
-        arrow_array
+        arrow2_array.into()
     }
 }
 
@@ -200,7 +199,7 @@ fn arrow_field_from_arrow2_field(field: Arrow2Field) -> ArrowField {
         Arrow2DataType::Struct(f) => {
             ArrowDataType::Struct(f.into_iter().map(arrow_field_from_arrow2_field).collect())
         }
-        other => other.into(),
+        other => ArrowDataType::from(other),
     };
 
     ArrowField::new(field.name, data_type, field.is_nullable).with_metadata(metadata)

src/query/expression/src/converts/arrow/to.rs

@@ -12,6 +12,9 @@ doctest = false
 test = true
 
 [dependencies]
+arrow-array = { workspace = true }
+arrow-buffer = { workspace = true }
+arrow-schema = { workspace = true }
 databend-common-exception = { path = "../../../../common/exception" }
 databend-common-expression = { path = "../../../expression" }
 databend-storages-common-table-meta = { path = "../table_meta" }

src/query/storages/common/blocks/Cargo.toml

@@ -0,0 +1,90 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed 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 std::sync::Arc;
+
+use arrow_array::types::ByteArrayType;
+use arrow_array::types::LargeBinaryType;
+use arrow_array::types::LargeUtf8Type;
+use arrow_array::Array;
+use arrow_array::GenericByteArray;
+use arrow_schema::DataType as ArrowDataType;
+use databend_common_exception::Result;
+use databend_storages_common_table_meta::meta::ColumnStatistics;
+use parquet_rs::basic::Encoding;
+use parquet_rs::file::properties::WriterPropertiesBuilder;
+use parquet_rs::schema::types::ColumnPath;
+
+const ROWS_PER_DISTINCT_THRESHOLD: f64 = 10.0;
+const SAMPLE_ROWS: usize = 1000;
+const AVERAGE_PREFIX_LEN_THRESHOLD: f64 = 8.0;
+
+pub fn choose_byte_array_encoding(
+    mut props: WriterPropertiesBuilder,
+    stat: Option<&ColumnStatistics>,
+    array: Arc<dyn Array>,
+    column_name: &str,
+) -> Result<WriterPropertiesBuilder> {
+    if array.is_empty() {
+        return Ok(props);
+    }
+    let col_path = ColumnPath::new(vec![column_name.to_string()]);
+    let ndv = stat.as_ref().and_then(|s| s.distinct_of_values);
+    let num_rows = array.len();
+    if let Some(ndv) = ndv {
+        if num_rows as f64 / ndv as f64 > ROWS_PER_DISTINCT_THRESHOLD {
+            props = props.set_column_dictionary_enabled(col_path, true);
+            return Ok(props);
+        }
+    }
+    let data_type = array.data_type();
+    match data_type {
+        ArrowDataType::LargeBinary => {
+            if can_apply_delta_byte_array::<LargeBinaryType>(&array)? {
+                props = props.set_column_encoding(col_path, Encoding::DELTA_BYTE_ARRAY);
+                return Ok(props);
+            }
+        }
+        ArrowDataType::LargeUtf8 => {
+            if can_apply_delta_byte_array::<LargeUtf8Type>(&array)? {
+                props = props.set_column_encoding(col_path, Encoding::DELTA_BYTE_ARRAY);
+                return Ok(props);
+            }
+        }
+        _ => {}
+    };
+    props = props.set_column_encoding(col_path, Encoding::DELTA_LENGTH_BYTE_ARRAY);
+    Ok(props)
+}
+
+fn can_apply_delta_byte_array<T: ByteArrayType>(array: &dyn Array) -> Result<bool> {
+    let num_rows = array.len();
+    let array = array
+        .as_any()
+        .downcast_ref::<GenericByteArray<T>>()
+        .unwrap();
+    let mut sum_prefix_len = 0;
+    for i in 1..num_rows.min(SAMPLE_ROWS) {
+        let last: &[u8] = array.value(i - 1).as_ref();
+        let cur: &[u8] = array.value(i).as_ref();
+        let prefix_len = last
+            .iter()
+            .zip(cur.iter())
+            .take_while(|(a, b)| a == b)
+            .count();
+        sum_prefix_len += prefix_len;
+    }
+    let avg_prefix_len = sum_prefix_len as f64 / num_rows as f64;
+    Ok(avg_prefix_len > AVERAGE_PREFIX_LEN_THRESHOLD)
+}

src/query/storages/common/blocks/src/codec/byte_array.rs

@@ -0,0 +1,60 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed 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 databend_common_exception::Result;
+use databend_common_expression::DataBlock;
+use databend_common_expression::TableSchema;
+use databend_storages_common_table_meta::meta::StatisticsOfColumns;
+use parquet_rs::basic::Type as PhysicalType;
+use parquet_rs::file::properties::WriterPropertiesBuilder;
+use parquet_rs::schema::types::Type;
+use parquet_rs::schema::types::TypePtr;
+
+use super::byte_array::choose_byte_array_encoding;
+use super::int::choose_int_encoding;
+
+pub fn choose_codec(
+    mut props: WriterPropertiesBuilder,
+    block: &DataBlock,
+    parquet_fields: &[TypePtr],
+    table_schema: &TableSchema,
+    stat: &StatisticsOfColumns,
+) -> Result<WriterPropertiesBuilder> {
+    for ((parquet_field, table_field), entry) in parquet_fields
+        .iter()
+        .zip(table_schema.fields.iter())
+        .zip(block.columns())
+    {
+        let column = entry.to_column(block.num_rows());
+        let array = column.into_arrow_rs();
+        let stat = stat.get(&table_field.column_id);
+        let column_name = table_field.name.as_str();
+        match parquet_field.as_ref() {
+            Type::PrimitiveType { physical_type, .. } => match physical_type {
+                PhysicalType::BYTE_ARRAY | PhysicalType::FIXED_LEN_BYTE_ARRAY => {
+                    props = choose_byte_array_encoding(props, stat, array, column_name)?;
+                }
+                PhysicalType::INT32 | PhysicalType::INT64 | PhysicalType::INT96 => {
+                    props = choose_int_encoding(props, stat, array, column_name)?
+                }
+                _ => {}
+            },
+            Type::GroupType {
+                basic_info: _,
+                fields: _,
+            } => {} // TODO: handle nested fields
+        }
+    }
+    Ok(props)
+}

src/query/storages/common/blocks/src/codec/choose.rs

@@ -0,0 +1,85 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed 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 std::sync::Arc;
+
+use arrow_array::types::ArrowPrimitiveType;
+use arrow_array::types::Date32Type;
+use arrow_array::types::Decimal128Type;
+use arrow_array::types::Decimal256Type;
+use arrow_array::types::Int32Type;
+use arrow_array::types::Int64Type;
+use arrow_array::types::TimestampMicrosecondType;
+use arrow_array::types::UInt32Type;
+use arrow_array::types::UInt64Type;
+use arrow_array::Array;
+use arrow_array::ArrowNativeTypeOp;
+use arrow_array::PrimitiveArray;
+use arrow_buffer::ArrowNativeType;
+use arrow_schema::DataType as ArrowDataType;
+use arrow_schema::TimeUnit;
+use databend_common_exception::Result;
+use databend_storages_common_table_meta::meta::ColumnStatistics;
+use parquet_rs::basic::Encoding;
+use parquet_rs::file::properties::WriterPropertiesBuilder;
+use parquet_rs::schema::types::ColumnPath;
+
+const MAX_WIDTH_THRESHOLD: i64 = 3;
+
+pub fn choose_int_encoding(
+    mut props: WriterPropertiesBuilder,
+    _stat: Option<&ColumnStatistics>,
+    array: Arc<dyn Array>,
+    column_name: &str,
+) -> Result<WriterPropertiesBuilder> {
+    if array.is_empty() {
+        return Ok(props);
+    }
+    let col_path = ColumnPath::new(vec![column_name.to_string()]);
+    let data_type = array.data_type();
+    let apply_delta = match data_type {
+        ArrowDataType::Int32 => can_apply_delta_binary_pack::<Int32Type>(&array)?,
+        ArrowDataType::Int64 => can_apply_delta_binary_pack::<Int64Type>(&array)?,
+        ArrowDataType::UInt32 => can_apply_delta_binary_pack::<UInt32Type>(&array)?,
+        ArrowDataType::UInt64 => can_apply_delta_binary_pack::<UInt64Type>(&array)?,
+        ArrowDataType::Timestamp(TimeUnit::Microsecond, _) => {
+            can_apply_delta_binary_pack::<TimestampMicrosecondType>(&array)?
+        }
+        ArrowDataType::Date32 => can_apply_delta_binary_pack::<Date32Type>(&array)?,
+        ArrowDataType::Decimal128(_, _) => can_apply_delta_binary_pack::<Decimal128Type>(&array)?,
+        ArrowDataType::Decimal256(_, _) => can_apply_delta_binary_pack::<Decimal256Type>(&array)?,
+        _ => false,
+    };
+    if apply_delta {
+        props = props.set_column_encoding(col_path, Encoding::DELTA_BINARY_PACKED);
+    }
+    Ok(props)
+}
+
+fn can_apply_delta_binary_pack<T: ArrowPrimitiveType>(array: &dyn Array) -> Result<bool> {
+    let mut max_delta = T::Native::MIN_TOTAL_ORDER;
+    let mut min_delta = T::Native::MAX_TOTAL_ORDER;
+    let array = array.as_any().downcast_ref::<PrimitiveArray<T>>().unwrap();
+    for i in 1..array.len() {
+        let delta = array.value(i).sub_wrapping(array.value(i - 1));
+        if delta.is_gt(max_delta) {
+            max_delta = delta;
+        }
+        if delta.is_lt(min_delta) {
+            min_delta = delta;
+        }
+    }
+    let x = max_delta.sub_wrapping(min_delta).as_usize();
+    Ok(x <= (1 << MAX_WIDTH_THRESHOLD))
+}

src/query/storages/common/blocks/src/codec/int.rs

@@ -0,0 +1,19 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed 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.
+
+mod byte_array;
+mod choose;
+mod int;
+
+pub use choose::choose_codec;

src/query/storages/common/blocks/src/codec/mod.rs

@@ -14,5 +14,6 @@
 
 #![allow(clippy::uninlined_format_args)]
 
+mod codec;
 mod parquet_rs;
 pub use parquet_rs::blocks_to_parquet;