fix hardcoded dtype

Signed-off-by: Connor Tsui <[email protected]>

Author: connortsui20

vortex-layout/src/layouts/dict/writer.rs

@@ -27,7 +27,6 @@ use vortex_array::builders::dict::DictEncoder;
 use vortex_array::builders::dict::dict_encoder;
 use vortex_btrblocks::BtrBlocksCompressor;
 use vortex_dtype::DType;
-use vortex_dtype::Nullability::NonNullable;
 use vortex_dtype::PType;
 use vortex_error::VortexError;
 use vortex_error::VortexResult;
@@ -49,10 +48,12 @@ use crate::sequence::SequentialStream;
 use crate::sequence::SequentialStreamAdapter;
 use crate::sequence::SequentialStreamExt;
 
+/// Constraints for dictionary layout encoding.
 #[derive(Clone)]
 pub struct DictLayoutConstraints {
+    /// Maximum size of the dictionary in bytes.
     pub max_bytes: usize,
-    // Dict layout codes currently only support u16 codes
+    /// Maximum dictionary length.
     pub max_len: u16,
 }
 
@@ -396,13 +397,15 @@ impl Stream for DictionaryTransformer {
                         self.active_codes_tx = Some(codes_tx.clone());
                         self.active_values_tx = Some(values_tx);
 
-                        // Send first codes
+                        let codes_dtype = codes.1.dtype().clone();
+
+                        // Send first codes.
                         self.pending_send =
                             Some(Box::pin(async move { codes_tx.send(Ok(codes)).await }));
 
-                        // Create output streams
+                        // Create output streams.
                         let codes_stream = SequentialStreamAdapter::new(
-                            DType::Primitive(PType::U16, NonNullable),
+                            codes_dtype,
                             codes_rx.into_stream().boxed(),
                         )
                         .sendable();
@@ -514,3 +517,106 @@ fn encode_chunk(mut encoder: Box<dyn DictEncoder>, chunk: &dyn Array) -> Encodin
 fn remainder(array: &dyn Array, encoded_len: usize) -> Option<ArrayRef> {
     (encoded_len < array.len()).then(|| array.slice(encoded_len..array.len()))
 }
+
+#[cfg(test)]
+mod tests {
+    use futures::StreamExt;
+    use vortex_array::IntoArray;
+    use vortex_array::arrays::VarBinArray;
+    use vortex_array::builders::dict::DictConstraints;
+    use vortex_dtype::DType;
+    use vortex_dtype::Nullability::NonNullable;
+    use vortex_dtype::PType;
+
+    use super::DictionaryTransformer;
+    use super::dict_encode_stream;
+    use crate::sequence::SequenceId;
+    use crate::sequence::SequentialStream;
+    use crate::sequence::SequentialStreamAdapter;
+    use crate::sequence::SequentialStreamExt;
+
+    /// Regression test for a bug where the codes stream dtype was hardcoded to U16 instead of
+    /// using the actual codes dtype from the array. When `max_len <= 255`, the dict encoder
+    /// produces U8 codes, but the stream was incorrectly typed as U16, causing a dtype mismatch
+    /// assertion failure in [`SequentialStreamAdapter`].
+    #[tokio::test]
+    async fn test_dict_transformer_uses_u8_for_small_dictionaries() {
+        // Use max_len = 100 to force U8 codes (since 100 <= 255).
+        let constraints = DictConstraints {
+            max_bytes: 1024 * 1024,
+            max_len: 100,
+        };
+
+        // Create a simple string array with a few unique values.
+        let arr = VarBinArray::from(vec!["hello", "world", "hello", "world"]).into_array();
+
+        // Wrap into a sequential stream.
+        let mut pointer = SequenceId::root();
+        let input_stream = SequentialStreamAdapter::new(
+            arr.dtype().clone(),
+            futures::stream::once(async move { Ok((pointer.advance(), arr)) }),
+        )
+        .sendable();
+
+        // Encode into dict chunks.
+        let dict_stream = dict_encode_stream(input_stream, constraints);
+
+        // Transform into codes/values streams.
+        let mut transformer = DictionaryTransformer::new(dict_stream);
+
+        // Get the first (and only) run.
+        let (codes_stream, _values_fut) = transformer
+            .next()
+            .await
+            .expect("expected at least one dictionary run");
+
+        // The key assertion: codes stream dtype should be U8, not U16.
+        assert_eq!(
+            codes_stream.dtype(),
+            &DType::Primitive(PType::U8, NonNullable),
+            "codes stream should use U8 dtype for small dictionaries, not U16"
+        );
+    }
+
+    /// Test that the codes stream uses U16 dtype when the dictionary has more than 255 entries.
+    #[tokio::test]
+    async fn test_dict_transformer_uses_u16_for_large_dictionaries() {
+        // Use max_len = 1000 to allow U16 codes (since 1000 > 255).
+        let constraints = DictConstraints {
+            max_bytes: 1024 * 1024,
+            max_len: 1000,
+        };
+
+        // Create an array with more than 255 distinct values to force U16 codes.
+        let values: Vec<String> = (0..300).map(|i| format!("value_{i}")).collect();
+        let arr =
+            VarBinArray::from(values.iter().map(|s| s.as_str()).collect::<Vec<_>>()).into_array();
+
+        // Wrap into a sequential stream.
+        let mut pointer = SequenceId::root();
+        let input_stream = SequentialStreamAdapter::new(
+            arr.dtype().clone(),
+            futures::stream::once(async move { Ok((pointer.advance(), arr)) }),
+        )
+        .sendable();
+
+        // Encode into dict chunks.
+        let dict_stream = dict_encode_stream(input_stream, constraints);
+
+        // Transform into codes/values streams.
+        let mut transformer = DictionaryTransformer::new(dict_stream);
+
+        // Get the first (and only) run.
+        let (codes_stream, _values_fut) = transformer
+            .next()
+            .await
+            .expect("expected at least one dictionary run");
+
+        // Codes stream dtype should be U16 since we have more than 255 distinct values.
+        assert_eq!(
+            codes_stream.dtype(),
+            &DType::Primitive(PType::U16, NonNullable),
+            "codes stream should use U16 dtype for dictionaries with >255 entries"
+        );
+    }
+}