Norm-extracted vectors

[connortsui20]

Summary

Tracking Issue: #6865

Related PR: #7018

Adds a new encoding specific to the vector extension type.

Note that we cannot actually utilize this by adding it to the compressor until we make it pluggable (see #7018). When that does eventually land, we can simply create a NormVectorScheme that uses the NormVectorArray.

Implementation

We currently do not have a good way of broadcasting a multiplication onto a FixedSizeList, so this implementation hand rolls the norm multiplication. Additionally, we still do not have #6717 so the scalar_at implementation is also slow.

API Changes

Adds a new encoding type NormVector

Testing

Some simple tests including roundtrips.

[codspeed-hq]

Merging this PR will degrade performance by 10.2%

❌ 1 regressed benchmark
✅ 1105 untouched benchmarks
⏩ 1522 skipped benchmarks¹

⚠️ Please fix the performance issues or acknowledge them on CodSpeed.

Performance Changes

	Mode	Benchmark	BASE	HEAD	Efficiency
❌	Simulation	map_each[BufferMut<i32>, 128]	770.6 ns	858.1 ns	-10.2%

---

_{Comparing ct/norm (e288bb1) with develop (ec2c602)}

Footnotes

1. 1522 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports. ↩

[connortsui20]

Note that once we get ScalarValue::Array (#6717), this will be a lot less ugly.

[gatesn]

Short cut would be L2Norm.new_expr(...)

[gatesn]

Even then, we should have a nicer way of doing this!!

[gatesn]

What is FlatElements?

[connortsui20]

It's a helper function I use for all of the other tensor types to make sure we can deal with ConstantArray correctly:

/// Extracts the flat primitive elements from a tensor storage array (FixedSizeList).
///
/// When the input is a [`ConstantArray`] (e.g., a literal query vector), only a single row is
/// materialized to avoid expanding it to the full column length.

[gatesn]

You said vector array can be nullable, but you don't check validity when iterating below

[connortsui20]

whoops

[gatesn]

This all seems a bit convoluted? Can you just run-end encode the norms by the vector lengths, then use a divide function?

[connortsui20]

So this is what the decompress path would look like (the compress path would be similar):

        // We need to multiply each vector element with its respective norm. We do not have any kind
        // of "broadcast" expression to each of the `FixedSizeList` elements, so we can mimic this
        // by multiplying the normalized vector array with a `RunEnd(Sequence)` array.
        let base: PValue = list_size.into();
        let multiplier: PValue = base;
        let ends_ptype = base.ptype();
        let ends_nullability = Nullability::NonNullable;
        let ends =
            SequenceArray::try_new(base, multiplier, ends_ptype, ends_nullability, num_vectors)?;

        let runend = RunEndArray::try_new(ends.into_array(), self.norms.clone())?;

        let storage = extension_storage(&self.vector_array)?;
        let fsl: FixedSizeListArray = storage.execute(ctx)?;
        let elements = fsl.elements();
        debug_assert!(elements.dtype().is_primitive());

        let decompress = elements.binary(runend.into_array(), Operator::Mul)?;

        let denormalized_elements: PrimitiveArray = decompress.execute(ctx)?;

        // SAFETY: TODO
        let fsl = unsafe {
            FixedSizeListArray::new_unchecked(
                denormalized_elements.into_array(),
                list_size,
                self.validity()?,
                num_vectors,
            )
        };

        Ok(ExtensionArray::new(ext.clone(), fsl.into_array()).into_array())
    }

I don't think that this is actually less convoluted? But maybe we do this regardless because the optimizer might be able to do something

[connortsui20]

ok just realized that a bunch of stuff was wrong because I changed from non-nullable to nullable halfway through, will fix everything

[connortsui20]

Also @gatesn the .binary(???, Operator::Div) won't work because we need a safe divide by 0, and Operator::Div will return an ArrowError that I can't easily intercept.

Do you think it is fine to just have this broadcast logic for decompression and not compression? (see the latest commit)