Fix column encoding so it matches the underlying madness (#58)

Author: coracuity

src/commands/transform.rs

@@ -47,6 +47,8 @@ pub async fn run(args: TransformCommand) -> Result<()> {
         parquet_statistics,
         parquet_writer_version,
         parquet_no_dictionary,
+        parquet_column_dictionary,
+        parquet_column_no_dictionary,
         parquet_encoding,
         parquet_column_encoding,
         parquet_sorted_metadata,
@@ -200,6 +202,8 @@ pub async fn run(args: TransformCommand) -> Result<()> {
         parquet_statistics,
         parquet_writer_version,
         parquet_no_dictionary,
+        parquet_column_dictionary,
+        parquet_column_no_dictionary,
         parquet_encoding,
         parquet_column_encoding,
         bloom_filter,
@@ -338,6 +342,8 @@ fn create_sink_factory(
     parquet_statistics: Option<ParquetStatistics>,
     parquet_writer_version: Option<ParquetWriterVersion>,
     parquet_no_dictionary: bool,
+    parquet_column_dictionary: Vec<String>,
+    parquet_column_no_dictionary: Vec<String>,
     parquet_encoding: Option<ParquetEncoding>,
     parquet_column_encoding: Vec<ColumnEncodingConfig>,
     parquet_bloom_filter: BloomFilterConfig,
@@ -380,6 +386,8 @@ fn create_sink_factory(
                     options = options.with_no_dictionary(true);
                 }
                 options = options
+                    .with_column_dictionary(parquet_column_dictionary.clone())
+                    .with_column_no_dictionary(parquet_column_no_dictionary.clone())
                     .with_encoding(parquet_encoding)
                     .with_column_encodings(parquet_column_encoding.clone())
                     .with_bloom_filters(parquet_bloom_filter.clone());

src/lib.rs

@@ -792,31 +792,63 @@ pub struct TransformCommand {
     #[arg(long, value_enum)]
     pub parquet_writer_version: Option<ParquetWriterVersion>,
 
-    /// Disable dictionary encoding for Parquet columns.
-    #[arg(long)]
+    /// Disable dictionary encoding globally for all Parquet columns.
+    ///
+    /// Dictionary encoding builds a dictionary of unique values and stores references to it,
+    /// which is very effective for low-cardinality columns (few unique values). When disabled,
+    /// columns use their data page encoding directly.
+    ///
+    /// Default: dictionary encoding is enabled.
+    ///
+    /// Use --parquet-column-dictionary or --parquet-column-no-dictionary to override
+    /// this setting for specific columns.
+    #[arg(long, verbatim_doc_comment)]
     pub parquet_no_dictionary: bool,
 
-    /// Encoding for Parquet column data pages.
+    /// Enable dictionary encoding for specific columns. Can be specified multiple times.
     ///
-    /// If dictionary encoding is disabled (via --parquet-no-dictionary), this is the primary
-    /// encoding for all columns. If dictionary encoding is enabled (the default), this is the
-    /// fallback encoding when dictionary encoding falls back (e.g., dictionary size exceeded).
+    /// Overrides --parquet-no-dictionary for the named columns, enabling dictionary encoding
+    /// even when it's globally disabled.
     ///
-    /// Options: plain, rle, delta-binary-packed, delta-length-byte-array, delta-byte-array, byte-stream-split
-    #[arg(long, value_enum)]
+    /// Useful when most columns have high cardinality (dictionary disabled globally) but
+    /// a few columns have low cardinality and would benefit from dictionary encoding.
+    #[arg(long, value_name = "COLUMN", verbatim_doc_comment)]
+    pub parquet_column_dictionary: Vec<String>,
+
+    /// Disable dictionary encoding for specific columns. Can be specified multiple times.
+    ///
+    /// Overrides the default (dictionary enabled) for the named columns.
+    ///
+    /// Useful for high-cardinality columns like UUIDs or timestamps where dictionary
+    /// encoding adds overhead without compression benefit.
+    #[arg(long, value_name = "COLUMN", verbatim_doc_comment)]
+    pub parquet_column_no_dictionary: Vec<String>,
+
+    /// Data page encoding for Parquet columns.
+    ///
+    /// This encoding is used for column data pages. Its role depends on dictionary encoding:
+    /// - Dictionary enabled (default): this is the fallback encoding, used when the dictionary
+    ///   becomes too large or is inefficient for the data.
+    /// - Dictionary disabled: this is the primary encoding for all data.
+    ///
+    /// If not specified, the writer automatically selects an encoding based on column type
+    /// and writer version. With Parquet v2: integers use delta-binary-packed, strings use
+    /// delta-byte-array, booleans use rle. With Parquet v1: everything uses plain.
+    ///
+    /// Options: plain, rle, delta-binary-packed, delta-length-byte-array, delta-byte-array,
+    /// byte-stream-split
+    #[arg(long, value_enum, verbatim_doc_comment)]
     pub parquet_encoding: Option<ParquetEncoding>,
 
-    /// Set encoding for specific columns. Can be specified multiple times.
-    /// Overrides --parquet-encoding for the named column.
+    /// Set data page encoding for specific columns. Can be specified multiple times.
     ///
-    /// If dictionary encoding is disabled for this column, this is the primary encoding.
-    /// If dictionary encoding is enabled (the default), this is the fallback encoding
-    /// when dictionary encoding falls back (e.g., dictionary size exceeded).
+    /// Overrides --parquet-encoding for the named column. See --parquet-encoding for
+    /// how this interacts with dictionary encoding.
     ///
     /// Format: COLUMN=ENCODING
     ///
-    /// Encoding options: plain, rle, delta-binary-packed, delta-length-byte-array,
-    /// delta-byte-array, byte-stream-split
+    /// Options: plain, rle, delta-binary-packed, delta-length-byte-array, delta-byte-array,
+    /// byte-stream-split
     ///
     /// Examples:
     ///   --parquet-column-encoding id=delta-binary-packed

src/sinks/parquet/mod.rs

@@ -32,6 +32,8 @@ pub struct ParquetSinkOptions {
     bloom_filters: BloomFilterConfig,
     statistics: ParquetStatistics,
     no_dictionary: bool,
+    column_dictionary: Vec<String>,
+    column_no_dictionary: Vec<String>,
     writer_version: ParquetWriterVersion,
     encoding: Option<ParquetEncoding>,
     column_encodings: Vec<ColumnEncodingConfig>,
@@ -56,6 +58,8 @@ impl ParquetSinkOptions {
             bloom_filters: BloomFilterConfig::default(),
             statistics: ParquetStatistics::default(),
             no_dictionary: false,
+            column_dictionary: Vec::new(),
+            column_no_dictionary: Vec::new(),
             writer_version: ParquetWriterVersion::default(),
             encoding: None,
             column_encodings: Vec::new(),
@@ -104,6 +108,16 @@ impl ParquetSinkOptions {
         self
     }
 
+    pub fn with_column_dictionary(mut self, column_dictionary: Vec<String>) -> Self {
+        self.column_dictionary = column_dictionary;
+        self
+    }
+
+    pub fn with_column_no_dictionary(mut self, column_no_dictionary: Vec<String>) -> Self {
+        self.column_no_dictionary = column_no_dictionary;
+        self
+    }
+
     pub fn with_writer_version(mut self, writer_version: ParquetWriterVersion) -> Self {
         self.writer_version = writer_version;
         self
@@ -153,6 +167,12 @@ impl ParquetSink {
             .set_statistics_enabled(options.statistics.into())
             .set_dictionary_enabled(!options.no_dictionary);
 
+        writer_builder = Self::apply_column_dictionary(
+            writer_builder,
+            &options.column_dictionary,
+            &options.column_no_dictionary,
+        );
+
         writer_builder =
             Self::apply_encodings(writer_builder, &options.encoding, &options.column_encodings);
 
@@ -187,6 +207,24 @@ impl ParquetSink {
         Ok(sink)
     }
 
+    fn apply_column_dictionary(
+        mut builder: WriterPropertiesBuilder,
+        column_dictionary: &[String],
+        column_no_dictionary: &[String],
+    ) -> WriterPropertiesBuilder {
+        for col_name in column_dictionary {
+            let col_path = ColumnPath::from(col_name.as_str());
+            builder = builder.set_column_dictionary_enabled(col_path, true);
+        }
+
+        for col_name in column_no_dictionary {
+            let col_path = ColumnPath::from(col_name.as_str());
+            builder = builder.set_column_dictionary_enabled(col_path, false);
+        }
+
+        builder
+    }
+
     fn apply_encodings(
         mut builder: WriterPropertiesBuilder,
         default_encoding: &Option<ParquetEncoding>,