fix scalars

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

vortex-tensor/src/encodings/norm/tests.rs

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use vortex::array::IntoArray;
 use vortex::array::VortexSessionExecute;
 use vortex::array::arrays::Extension;
 use vortex::error::VortexResult;
@@ -108,3 +109,27 @@ fn execute_round_trip_zero_vector() -> VortexResult<()> {
 
     Ok(())
 }
+
+#[test]
+fn scalar_at_returns_original_vector() -> VortexResult<()> {
+    let arr = vector_array(
+        2,
+        &[
+            3.0, 4.0, // norm = 5.0
+            6.0, 8.0, // norm = 10.0
+        ],
+    )?;
+
+    let encoded = NormVectorArray::compress(arr)?;
+
+    // `scalar_at` on the NormVectorArray should match `scalar_at` on the decompressed result.
+    let mut ctx = vortex::array::LEGACY_SESSION.create_execution_ctx();
+    let decompressed = encoded.decompress(&mut ctx)?;
+
+    let norm_array = encoded.into_array();
+    for i in 0..2 {
+        assert_eq!(norm_array.scalar_at(i)?, decompressed.scalar_at(i)?);
+    }
+
+    Ok(())
+}

vortex-tensor/src/encodings/norm/vtable/operations.rs

@@ -1,15 +1,66 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use vortex::array::IntoArray;
+use vortex::array::arrays::ConstantArray;
+use vortex::array::arrays::FixedSizeList;
+use vortex::array::builtins::ArrayBuiltins;
 use vortex::array::vtable::OperationsVTable;
+use vortex::dtype::Nullability;
 use vortex::error::VortexResult;
+use vortex::error::vortex_err;
 use vortex::scalar::Scalar;
+use vortex::scalar_fn::fns::operators::Operator;
 
 use crate::encodings::norm::array::NormVectorArray;
 use crate::encodings::norm::vtable::NormVector;
+use crate::utils::extension_list_size;
+use crate::utils::extension_storage;
 
 impl OperationsVTable<NormVector> for NormVector {
     fn scalar_at(array: &NormVectorArray, index: usize) -> VortexResult<Scalar> {
-        array.vector_array().scalar_at(index)
+        let ext = array
+            .vector_array()
+            .dtype()
+            .as_extension_opt()
+            .ok_or_else(|| {
+                vortex_err!(
+                    "expected Vector extension dtype, got {}",
+                    array.vector_array().dtype()
+                )
+            })?;
+        let list_size = extension_list_size(ext)?;
+
+        // Get the storage (FixedSizeList) and slice out the elements for this row.
+        let storage = extension_storage(array.vector_array())?;
+        let fsl = storage
+            .as_opt::<FixedSizeList>()
+            .ok_or_else(|| vortex_err!("expected FixedSizeList storage"))?;
+        let row_elements = fsl.fixed_size_list_elements_at(index)?;
+
+        // Multiply all elements by the norm using a ConstantArray broadcast.
+        let norm_scalar = array.norms().scalar_at(index)?;
+        let norm_broadcast = ConstantArray::new(norm_scalar, list_size).into_array();
+        let scaled = row_elements.binary(norm_broadcast, Operator::Mul)?;
+
+        // Rebuild the FSL scalar, then wrap in the extension type.
+        let element_dtype = ext
+            .storage_dtype()
+            .as_fixed_size_list_element_opt()
+            .ok_or_else(|| {
+                vortex_err!(
+                    "expected FixedSizeList storage dtype, got {}",
+                    ext.storage_dtype()
+                )
+            })?;
+
+        let children: Vec<Scalar> = (0..list_size)
+            .map(|i| scaled.scalar_at(i))
+            .collect::<VortexResult<_>>()?;
+
+        let fsl_scalar =
+            Scalar::fixed_size_list(element_dtype.clone(), children, Nullability::NonNullable);
+
+        Ok(Scalar::extension_ref(ext.clone(), fsl_scalar))
     }
 }