reifydb_core/value/column/transform/

take.rs

// SPDX-License-Identifier: Apache-2.0
// Copyright (c) 2025 ReifyDB

use reifydb_type::{Result, util::cowvec::CowVec};

use crate::value::column::columns::Columns;

impl Columns {
	pub fn take(&mut self, n: usize) -> Result<()> {
		// Take the first n encoded numbers
		if !self.row_numbers.is_empty() {
			let actual_n = n.min(self.row_numbers.len());
			let new_row_numbers: Vec<_> = self.row_numbers.iter().take(actual_n).copied().collect();
			self.row_numbers = CowVec::new(new_row_numbers);
		}

		// Take the first n created_at timestamps
		if !self.created_at.is_empty() {
			let actual_n = n.min(self.created_at.len());
			let new_created_at: Vec<_> = self.created_at.iter().take(actual_n).copied().collect();
			self.created_at = CowVec::new(new_created_at);
		}

		// Take the first n updated_at timestamps
		if !self.updated_at.is_empty() {
			let actual_n = n.min(self.updated_at.len());
			let new_updated_at: Vec<_> = self.updated_at.iter().take(actual_n).copied().collect();
			self.updated_at = CowVec::new(new_updated_at);
		}

		// Take the first n rows from columns
		let mut new_buffers = Vec::with_capacity(self.len());

		for data in self.columns.iter() {
			new_buffers.push(data.take(n));
		}

		self.columns = CowVec::new(new_buffers);

		Ok(())
	}
}

#[cfg(test)]
pub mod tests {
	use reifydb_type::value::{Value, r#type::Type};

	use super::*;
	use crate::value::column::{ColumnBuffer, ColumnWithName};

	#[test]
	fn test_bool_column() {
		let mut test_instance = Columns::new(vec![ColumnWithName::bool_with_bitvec(
			"flag",
			[true, true, false],
			[false, true, true],
		)]);

		test_instance.take(1).unwrap();

		assert_eq!(test_instance[0], ColumnBuffer::bool_with_bitvec([true], [false]));
	}

	#[test]
	fn test_float4_column() {
		let mut test_instance = Columns::new(vec![ColumnWithName::float4_with_bitvec(
			"a",
			[1.0, 2.0, 3.0],
			[true, false, true],
		)]);

		test_instance.take(2).unwrap();

		assert_eq!(test_instance[0], ColumnBuffer::float4_with_bitvec([1.0, 2.0], [true, false]));
	}

	#[test]
	fn test_float8_column() {
		let mut test_instance = Columns::new(vec![ColumnWithName::float8_with_bitvec(
			"a",
			[1f64, 2.0, 3.0, 4.0],
			[true, true, false, true],
		)]);

		test_instance.take(2).unwrap();

		assert_eq!(test_instance[0], ColumnBuffer::float8_with_bitvec([1.0, 2.0], [true, true]));
	}

	#[test]
	fn test_int1_column() {
		let mut test_instance =
			Columns::new(vec![ColumnWithName::int1_with_bitvec("a", [1, 2, 3], [true, false, true])]);

		test_instance.take(2).unwrap();

		assert_eq!(test_instance[0], ColumnBuffer::int1_with_bitvec([1, 2], [true, false]));
	}

	#[test]
	fn test_int2_column() {
		let mut test_instance = Columns::new(vec![ColumnWithName::int2_with_bitvec(
			"a",
			[1, 2, 3, 4],
			[true, true, false, true],
		)]);

		test_instance.take(2).unwrap();

		assert_eq!(test_instance[0], ColumnBuffer::int2_with_bitvec([1, 2], [true, true]));
	}

	#[test]
	fn test_int4_column() {
		let mut test_instance =
			Columns::new(vec![ColumnWithName::int4_with_bitvec("a", [1, 2], [true, false])]);

		test_instance.take(1).unwrap();

		assert_eq!(test_instance[0], ColumnBuffer::int4_with_bitvec([1], [true]));
	}

	#[test]
	fn test_int8_column() {
		let mut test_instance =
			Columns::new(vec![ColumnWithName::int8_with_bitvec("a", [1, 2, 3], [false, true, true])]);

		test_instance.take(2).unwrap();

		assert_eq!(test_instance[0], ColumnBuffer::int8_with_bitvec([1, 2], [false, true]));
	}

	#[test]
	fn test_int16_column() {
		let mut test_instance =
			Columns::new(vec![ColumnWithName::int16_with_bitvec("a", [1, 2], [true, true])]);

		test_instance.take(1).unwrap();

		assert_eq!(test_instance[0], ColumnBuffer::int16_with_bitvec([1], [true]));
	}

	#[test]
	fn test_uint1_column() {
		let mut test_instance =
			Columns::new(vec![ColumnWithName::uint1_with_bitvec("a", [1, 2, 3], [false, false, true])]);

		test_instance.take(2).unwrap();

		assert_eq!(test_instance[0], ColumnBuffer::uint1_with_bitvec([1, 2], [false, false]));
	}

	#[test]
	fn test_uint2_column() {
		let mut test_instance =
			Columns::new(vec![ColumnWithName::uint2_with_bitvec("a", [1, 2], [true, false])]);

		test_instance.take(1).unwrap();

		assert_eq!(test_instance[0], ColumnBuffer::uint2_with_bitvec([1], [true]));
	}

	#[test]
	fn test_uint4_column() {
		let mut test_instance =
			Columns::new(vec![ColumnWithName::uint4_with_bitvec("a", [10, 20], [false, true])]);

		test_instance.take(1).unwrap();

		assert_eq!(test_instance[0], ColumnBuffer::uint4_with_bitvec([10], [false]));
	}

	#[test]
	fn test_uint8_column() {
		let mut test_instance =
			Columns::new(vec![ColumnWithName::uint8_with_bitvec("a", [10, 20, 30], [true, true, false])]);

		test_instance.take(2).unwrap();

		assert_eq!(test_instance[0], ColumnBuffer::uint8_with_bitvec([10, 20], [true, true]));
	}

	#[test]
	fn test_uint16_column() {
		let mut test_instance = Columns::new(vec![ColumnWithName::uint16_with_bitvec(
			"a",
			[100, 200, 300],
			[true, false, true],
		)]);

		test_instance.take(1).unwrap();

		assert_eq!(test_instance[0], ColumnBuffer::uint16_with_bitvec([100], [true]));
	}

	#[test]
	fn test_text_column() {
		let mut test_instance = Columns::new(vec![ColumnWithName::utf8_with_bitvec(
			"t",
			vec!["a".to_string(), "b".to_string(), "c".to_string()],
			vec![true, false, true],
		)]);

		test_instance.take(2).unwrap();

		assert_eq!(
			test_instance[0],
			ColumnBuffer::utf8_with_bitvec(["a".to_string(), "b".to_string()], [true, false])
		);
	}

	#[test]
	fn test_none_column() {
		let mut test_instance = Columns::new(vec![ColumnWithName::undefined_typed("u", Type::Boolean, 3)]);

		test_instance.take(2).unwrap();

		assert_eq!(test_instance[0].len(), 2);
		assert_eq!(test_instance[0].get_value(0), Value::none());
		assert_eq!(test_instance[0].get_value(1), Value::none());
	}

	#[test]
	fn test_handles_none() {
		let mut test_instance = Columns::new(vec![ColumnWithName::undefined_typed("u", Type::Boolean, 5)]);

		test_instance.take(3).unwrap();

		assert_eq!(test_instance[0].len(), 3);
		assert_eq!(test_instance[0].get_value(0), Value::none());
	}

	#[test]
	fn test_n_larger_than_len_is_safe() {
		let mut test_instance =
			Columns::new(vec![ColumnWithName::int2_with_bitvec("a", [10, 20], [true, false])]);

		test_instance.take(10).unwrap();

		assert_eq!(test_instance[0], ColumnBuffer::int2_with_bitvec([10, 20], [true, false]));
	}
}