feat: avoid building dictionary for massive unique column values

analytic_engine/Cargo.toml

@@ -50,6 +50,7 @@ datafusion = { workspace = true }
 future_ext = { workspace = true }
 futures = { workspace = true }
 generic_error = { workspace = true }
+hash_ext = { workspace = true }
 hex = { workspace = true }
 hyperloglog = { workspace = true }
 id_allocator = { workspace = true }

analytic_engine/src/sst/factory.rs

@@ -25,6 +25,7 @@ use snafu::{ResultExt, Snafu};
 use table_engine::predicate::PredicateRef;
 use trace_metric::MetricsCollector;
 
+use super::parquet::writer::WriteOptions;
 use crate::{
     sst::{
         file::Level,
@@ -200,11 +201,16 @@ impl Factory for FactoryImpl {
         store_picker: &'a ObjectStorePickerRef,
         level: Level,
     ) -> Result<Box<dyn SstWriter + Send + 'a>> {
+        let write_options = WriteOptions {
+            num_rows_per_row_group: options.num_rows_per_row_group,
+            max_buffer_size: options.max_buffer_size,
+            compression: options.compression.into(),
+            sst_level: level,
+        };
         Ok(Box::new(ParquetSstWriter::new(
             path,
-            level,
+            write_options,
             store_picker,
-            options,
         )))
     }
 }

analytic_engine/src/sst/parquet/encoding.rs

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::convert::TryFrom;
+use std::{collections::HashMap, convert::TryFrom};
 
 use arrow::{compute, record_batch::RecordBatch as ArrowRecordBatch};
 use async_trait::async_trait;
@@ -26,6 +26,7 @@ use parquet::{
     arrow::AsyncArrowWriter,
     basic::Compression,
     file::{metadata::KeyValue, properties::WriterProperties},
+    schema::types::ColumnPath,
 };
 use prost::{bytes, Message};
 use snafu::{ensure, Backtrace, OptionExt, ResultExt, Snafu};
@@ -237,25 +238,40 @@ struct ColumnarRecordEncoder<W> {
     arrow_schema: ArrowSchemaRef,
 }
 
+#[derive(Debug, Clone)]
+pub struct ColumnEncoding {
+    pub enable_dict: bool,
+}
+
+#[derive(Debug, Clone)]
+pub struct EncodeOptions {
+    pub num_rows_per_row_group: usize,
+    pub max_buffer_size: usize,
+    pub compression: Compression,
+    pub column_encodings: HashMap<String, ColumnEncoding>,
+}
+
 impl<W: AsyncWrite + Send + Unpin> ColumnarRecordEncoder<W> {
-    fn try_new(
-        sink: W,
-        schema: &Schema,
-        num_rows_per_row_group: usize,
-        max_buffer_size: usize,
-        compression: Compression,
-    ) -> Result<Self> {
+    fn try_new(sink: W, schema: &Schema, options: &EncodeOptions) -> Result<Self> {
         let arrow_schema = schema.to_arrow_schema_ref();
 
-        let write_props = WriterProperties::builder()
-            .set_max_row_group_size(num_rows_per_row_group)
-            .set_compression(compression)
-            .build();
+        let write_props = {
+            let mut builder = WriterProperties::builder()
+                .set_max_row_group_size(options.num_rows_per_row_group)
+                .set_compression(options.compression);
+
+            for (col_name, encoding) in &options.column_encodings {
+                let col_path = ColumnPath::new(vec![col_name.to_string()]);
+                builder = builder.set_column_dictionary_enabled(col_path, encoding.enable_dict);
+            }
+
+            builder.build()
+        };
 
         let arrow_writer = AsyncArrowWriter::try_new(
             sink,
             arrow_schema.clone(),
-            max_buffer_size,
+            options.max_buffer_size,
             Some(write_props),
         )
         .box_err()
@@ -326,18 +342,10 @@ impl ParquetEncoder {
     pub fn try_new<W: AsyncWrite + Unpin + Send + 'static>(
         sink: W,
         schema: &Schema,
-        num_rows_per_row_group: usize,
-        max_buffer_size: usize,
-        compression: Compression,
+        options: &EncodeOptions,
     ) -> Result<Self> {
         Ok(ParquetEncoder {
-            record_encoder: Box::new(ColumnarRecordEncoder::try_new(
-                sink,
-                schema,
-                num_rows_per_row_group,
-                max_buffer_size,
-                compression,
-            )?),
+            record_encoder: Box::new(ColumnarRecordEncoder::try_new(sink, schema, options)?),
         })
     }
 

analytic_engine/src/sst/parquet/meta_data.rs

@@ -353,14 +353,14 @@ pub struct ParquetMetaData {
 
 pub type ParquetMetaDataRef = Arc<ParquetMetaData>;
 
-impl From<MetaData> for ParquetMetaData {
-    fn from(meta: MetaData) -> Self {
+impl From<&MetaData> for ParquetMetaData {
+    fn from(meta: &MetaData) -> Self {
         Self {
-            min_key: meta.min_key,
-            max_key: meta.max_key,
+            min_key: meta.min_key.clone(),
+            max_key: meta.max_key.clone(),
             time_range: meta.time_range,
             max_sequence: meta.max_sequence,
-            schema: meta.schema,
+            schema: meta.schema.clone(),
             parquet_filter: None,
             column_values: None,
         }