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/instance/engine.rs

@@ -26,10 +26,14 @@ use table_engine::{
 };
 use wal::manager::WalLocation;
 
-use super::open::{TableContext, TablesOfShardContext};
 use crate::{
     engine::build_space_id,
-    instance::{close::Closer, drop::Dropper, open::OpenTablesOfShardResult, Instance},
+    instance::{
+        close::Closer,
+        drop::Dropper,
+        open::{OpenTablesOfShardResult, TableContext, TablesOfShardContext},
+        Instance,
+    },
     space::{MemSizeOptions, Space, SpaceAndTable, SpaceContext, SpaceId, SpaceRef},
 };
 

analytic_engine/src/instance/open.rs

@@ -26,13 +26,13 @@ use snafu::ResultExt;
 use table_engine::{engine::TableDef, table::TableId};
 use wal::manager::WalManagerRef;
 
-use super::{engine::OpenTablesOfShard, flush_compaction::Flusher};
 use crate::{
     compaction::scheduler::SchedulerImpl,
     context::OpenContext,
     engine,
     instance::{
-        engine::{OpenManifest, ReadMetaUpdate, Result},
+        engine::{OpenManifest, OpenTablesOfShard, ReadMetaUpdate, Result},
+        flush_compaction::Flusher,
         mem_collector::MemUsageCollector,
         wal_replayer::{ReplayMode, WalReplayer},
         Instance, SpaceStore,

analytic_engine/src/instance/serial_executor.rs

@@ -30,9 +30,8 @@ use tokio::sync::{
     watch::{self, Receiver, Sender},
 };
 
-use super::flush_compaction::{BackgroundFlushFailed, TableFlushOptions};
 use crate::{
-    instance::flush_compaction::{Other, Result},
+    instance::flush_compaction::{BackgroundFlushFailed, Other, Result, TableFlushOptions},
     table::data::TableData,
 };
 

analytic_engine/src/sst/factory.rs

@@ -32,7 +32,10 @@ use crate::{
         header::HeaderParser,
         meta_data::cache::MetaCacheRef,
         metrics::MaybeTableLevelMetrics as SstMaybeTableLevelMetrics,
-        parquet::{writer::ParquetSstWriter, AsyncParquetReader, ThreadedReader},
+        parquet::{
+            writer::{ParquetSstWriter, WriteOptions},
+            AsyncParquetReader, ThreadedReader,
+        },
         reader::SstReader,
         writer::SstWriter,
     },
@@ -200,11 +203,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/meta_data/cache.rs

@@ -175,7 +175,7 @@ mod tests {
     use parquet::{arrow::ArrowWriter, file::footer};
     use parquet_ext::ParquetMetaData;
 
-    use super::MetaData;
+    use super::*;
     use crate::{
         sst::parquet::{
             encoding::{self, META_PATH_KEY, META_VERSION_KEY},

analytic_engine/src/sst/meta_data/metadata_reader.rs

@@ -20,19 +20,18 @@ use object_store::{ObjectStoreRef, Path};
 use parquet::{data_type::AsBytes, file::metadata::KeyValue};
 use snafu::{ensure, OptionExt, ResultExt};
 
-use super::UnknownMetaVersion;
 use crate::sst::{
     meta_data::{
         DecodeCustomMetaData, FetchAndDecodeSstMeta, FetchFromStore, KvMetaDataNotFound,
-        KvMetaPathEmpty,
+        KvMetaPathEmpty, UnknownMetaVersion,
     },
     parquet::{
         encoding::{self, decode_sst_meta_data_from_bytes, META_VERSION_CURRENT, META_VERSION_V1},
         meta_data::{ParquetMetaData, ParquetMetaDataRef},
     },
 };
 
-define_result!(super::Error);
+define_result!(crate::sst::meta_data::Error);
 
 #[async_trait]
 pub trait CustomMetadataReader {

analytic_engine/src/sst/parquet/async_reader.rs

@@ -57,7 +57,6 @@ use tokio::sync::{
 };
 use trace_metric::{MetricsCollector, TraceMetricWhenDrop};
 
-use super::meta_data::ColumnValueSet;
 use crate::{
     prefetchable_stream::{NoopPrefetcher, PrefetchableStream},
     sst::{
@@ -68,7 +67,9 @@ use crate::{
         },
         metrics::MaybeTableLevelMetrics,
         parquet::{
-            encoding::ParquetDecoder, meta_data::ParquetFilter, row_group_pruner::RowGroupPruner,
+            encoding::ParquetDecoder,
+            meta_data::{ColumnValueSet, ParquetFilter},
+            row_group_pruner::RowGroupPruner,
         },
         reader::{error::*, Result, SstReader},
     },

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,
         }

analytic_engine/src/sst/parquet/row_group_pruner.rs

@@ -36,9 +36,8 @@ use parquet_ext::prune::{
 use snafu::ensure;
 use trace_metric::{MetricsCollector, TraceMetricWhenDrop};
 
-use super::meta_data::ColumnValueSet;
 use crate::sst::{
-    parquet::meta_data::ParquetFilter,
+    parquet::meta_data::{ColumnValueSet, ParquetFilter},
     reader::error::{OtherNoCause, Result},
 };