1use std::{
5 collections::HashMap,
6 fmt::{self, Display, Formatter},
7 slice,
8 sync::{Arc, LazyLock},
9};
10
11use arrow_array::{
12 Array, ArrayRef, FixedSizeBinaryArray, RecordBatch, RecordBatchOptions, TimestampNanosecondArray, UInt64Array,
13};
14use arrow_schema::{ArrowError, Field, FieldRef, Schema};
15
16use crate::{
17 Result,
18 error::Error,
19 value::{
20 column_view::ColumnView,
21 container::{temporal_array::datetimes, wide_int_array::wides},
22 datetime::DateTime,
23 partition::Partition,
24 row_number::RowNumber,
25 value_type::{
26 ValueType,
27 field::{FieldType, field_error, is_sumtype_tag, to_field},
28 },
29 },
30};
31
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33pub enum SystemColumn {
34 RowNumbers,
35 Partitions,
36 CreatedAt,
37 UpdatedAt,
38 Time,
39 CommitVersion,
40 Sequence,
41}
42
43impl SystemColumn {
44 pub const ALL: [SystemColumn; 7] = [
45 SystemColumn::RowNumbers,
46 SystemColumn::Partitions,
47 SystemColumn::CreatedAt,
48 SystemColumn::UpdatedAt,
49 SystemColumn::Time,
50 SystemColumn::CommitVersion,
51 SystemColumn::Sequence,
52 ];
53
54 pub const fn name(self) -> &'static str {
55 match self {
56 SystemColumn::RowNumbers => "#rownum",
57 SystemColumn::Partitions => "#partition",
58 SystemColumn::CreatedAt => "#created_at",
59 SystemColumn::UpdatedAt => "#updated_at",
60 SystemColumn::Time => "#time",
61 SystemColumn::CommitVersion => "#commit_version",
62 SystemColumn::Sequence => "#sequence",
63 }
64 }
65
66 pub const fn ty(self) -> ValueType {
67 match self {
68 SystemColumn::RowNumbers => ValueType::Uint8,
69 SystemColumn::Partitions => ValueType::Uint16,
70 SystemColumn::CreatedAt => ValueType::DateTime,
71 SystemColumn::UpdatedAt => ValueType::DateTime,
72 SystemColumn::Time => ValueType::DateTime,
73 SystemColumn::CommitVersion => ValueType::Uint8,
74 SystemColumn::Sequence => ValueType::Uint8,
75 }
76 }
77
78 pub fn from_name(name: &str) -> Option<SystemColumn> {
79 SystemColumn::ALL.into_iter().find(|column| column.name() == name)
80 }
81
82 fn rank(self) -> usize {
83 SystemColumn::ALL.iter().position(|column| *column == self).expect("every system column is in ALL")
84 }
85}
86
87static SYSTEM_FIELDS: LazyLock<[[FieldRef; 2]; 7]> = LazyLock::new(|| {
88 SystemColumn::ALL.map(|column| {
89 [false, true].map(|holds_nones| {
90 let value_type = match holds_nones {
91 true => ValueType::Option(Box::new(column.ty())),
92 false => column.ty(),
93 };
94 Arc::new(to_field(
95 column.name(),
96 &FieldType {
97 value_type: Some(value_type),
98 ..FieldType::default()
99 },
100 ))
101 })
102 })
103});
104
105pub fn system_field(column: SystemColumn, holds_nones: bool) -> FieldRef {
106 SYSTEM_FIELDS[column.rank()][usize::from(holds_nones)].clone()
107}
108
109impl Display for SystemColumn {
110 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
111 f.write_str(self.name())
112 }
113}
114
115pub fn system_column(batch: &RecordBatch, column: SystemColumn) -> Option<&ArrayRef> {
116 position(batch, column.name()).map(|index| batch.column(index))
117}
118
119pub fn row_numbers(batch: &RecordBatch) -> Result<&[RowNumber]> {
120 let Some(array) = present::<UInt64Array>(batch, SystemColumn::RowNumbers)? else {
121 return Ok(&[]);
122 };
123 let values: &[u64] = array.values();
124 Ok(unsafe { slice::from_raw_parts(values.as_ptr().cast::<RowNumber>(), values.len()) })
126}
127
128pub fn sequences(batch: &RecordBatch) -> Result<&[u64]> {
129 let Some(array) = present::<UInt64Array>(batch, SystemColumn::Sequence)? else {
130 return Ok(&[]);
131 };
132 Ok(array.values())
133}
134
135pub fn partitions(batch: &RecordBatch) -> Result<Vec<Partition>> {
136 let Some(array) = present::<FixedSizeBinaryArray>(batch, SystemColumn::Partitions)? else {
137 return Ok(Vec::new());
138 };
139 Ok(wides::<u128>(array).into_iter().map(Partition).collect())
140}
141
142pub fn created_at(batch: &RecordBatch) -> Result<&[DateTime]> {
143 datetime_column(batch, SystemColumn::CreatedAt)
144}
145
146pub fn updated_at(batch: &RecordBatch) -> Result<&[DateTime]> {
147 datetime_column(batch, SystemColumn::UpdatedAt)
148}
149
150pub fn time(batch: &RecordBatch) -> Result<&[DateTime]> {
151 datetime_column(batch, SystemColumn::Time)
152}
153
154#[derive(Debug, Clone, Copy, PartialEq, Eq)]
155pub struct RowTimes<'a> {
156 values: &'a [DateTime],
157 defined: Option<(&'a [u8], usize)>,
158}
159
160impl<'a> RowTimes<'a> {
161 pub const EMPTY: Self = Self {
162 values: &[],
163 defined: None,
164 };
165
166 pub fn new(values: &'a [DateTime], defined: Option<(&'a [u8], usize)>) -> Self {
167 Self {
168 values,
169 defined,
170 }
171 }
172
173 pub fn len(&self) -> usize {
174 self.values.len()
175 }
176
177 pub fn is_empty(&self) -> bool {
178 self.values.is_empty()
179 }
180
181 pub fn get(&self, index: usize) -> Option<DateTime> {
182 let value = *self.values.get(index)?;
183 match self.defined {
184 Some((bits, offset)) => {
185 let bit = offset + index;
186 bits.get(bit / 8).is_some_and(|byte| byte & (1 << (bit % 8)) != 0).then_some(value)
187 }
188 None => Some(value),
189 }
190 }
191
192 pub fn iter(&self) -> impl Iterator<Item = Option<DateTime>> + '_ {
193 (0..self.len()).map(|index| self.get(index))
194 }
195
196 pub fn values(&self) -> &'a [DateTime] {
197 self.values
198 }
199
200 pub fn defined(&self) -> Option<(&'a [u8], usize)> {
201 self.defined
202 }
203}
204
205#[allow(clippy::disallowed_methods)]
206pub fn row_times(batch: &RecordBatch) -> Result<RowTimes<'_>> {
207 let Some(array) = system_column(batch, SystemColumn::Time) else {
208 return Ok(RowTimes::EMPTY);
209 };
210 let array = array.as_any().downcast_ref::<TimestampNanosecondArray>().ok_or_else(|| {
211 field_error(format!(
212 "system column {} holds arrow type {}",
213 SystemColumn::Time.name(),
214 array.data_type()
215 ))
216 })?;
217 Ok(RowTimes::new(datetimes(array), array.nulls().map(|nulls| (nulls.validity(), nulls.offset()))))
218}
219
220pub fn require_row_numbers(batch: &RecordBatch) -> Result<&[RowNumber]> {
221 require(batch, SystemColumn::RowNumbers)?;
222 row_numbers(batch)
223}
224
225pub fn require_created_at(batch: &RecordBatch) -> Result<&[DateTime]> {
226 require(batch, SystemColumn::CreatedAt)?;
227 created_at(batch)
228}
229
230pub fn require_updated_at(batch: &RecordBatch) -> Result<&[DateTime]> {
231 require(batch, SystemColumn::UpdatedAt)?;
232 updated_at(batch)
233}
234
235pub fn require_time(batch: &RecordBatch) -> Result<&[DateTime]> {
236 require(batch, SystemColumn::Time)?;
237 time(batch)
238}
239
240pub fn commit_versions(batch: &RecordBatch) -> Result<&[u64]> {
241 match present::<UInt64Array>(batch, SystemColumn::CommitVersion)? {
242 Some(array) => Ok(array.values()),
243 None => Ok(&[]),
244 }
245}
246
247pub fn with_system_column(batch: RecordBatch, column: SystemColumn, array: ArrayRef) -> Result<RecordBatch> {
248 stamp_system_columns(batch, vec![(column, array)])
249}
250
251pub fn stamp_system_columns(batch: RecordBatch, stamps: Vec<(SystemColumn, ArrayRef)>) -> Result<RecordBatch> {
252 let schema = batch.schema();
253 let fields: Vec<FieldRef> = schema.fields().iter().cloned().collect();
254 stamped(schema.metadata().clone(), fields, batch.columns().to_vec(), batch.num_rows(), stamps)
255}
256
257pub fn restamp_row_numbers(batch: &RecordBatch, keep: &[SystemColumn], row_numbers: ArrayRef) -> Result<RecordBatch> {
258 let schema = batch.schema_ref();
259 let (fields, columns): (Vec<FieldRef>, Vec<ArrayRef>) = schema
260 .fields()
261 .iter()
262 .zip(batch.columns())
263 .filter(|(field, _)| SystemColumn::from_name(field.name()).is_none_or(|column| keep.contains(&column)))
264 .map(|(field, array)| (field.clone(), array.clone()))
265 .unzip();
266 stamped(
267 schema.metadata().clone(),
268 fields,
269 columns,
270 batch.num_rows(),
271 vec![(SystemColumn::RowNumbers, row_numbers)],
272 )
273}
274
275fn stamped(
276 metadata: HashMap<String, String>,
277 mut fields: Vec<FieldRef>,
278 mut columns: Vec<ArrayRef>,
279 num_rows: usize,
280 stamps: Vec<(SystemColumn, ArrayRef)>,
281) -> Result<RecordBatch> {
282 let mut row_count = num_rows;
283 for (column, array) in stamps {
284 let field = system_field(column, array.logical_null_count() > 0);
285 if fields.is_empty() && row_count == 0 {
286 row_count = array.len();
287 }
288 match fields.iter().position(|other| other.name() == column.name()) {
289 Some(index) => {
290 fields[index] = field;
291 columns[index] = array;
292 }
293 None => {
294 let index = fields
295 .iter()
296 .position(|other| {
297 SystemColumn::from_name(other.name())
298 .is_some_and(|other| other.rank() > column.rank())
299 })
300 .unwrap_or(fields.len());
301 fields.insert(index, field);
302 columns.insert(index, array);
303 }
304 }
305 }
306 RecordBatch::try_new_with_options(
307 Arc::new(Schema::new_with_metadata(fields, metadata)),
308 columns,
309 &RecordBatchOptions::new().with_row_count(Some(row_count)),
310 )
311 .map_err(arrow_error)
312}
313
314pub fn keep_system_columns(batch: &RecordBatch, keep: &[SystemColumn]) -> Result<RecordBatch> {
315 let indices: Vec<usize> = batch
316 .schema_ref()
317 .fields()
318 .iter()
319 .enumerate()
320 .filter(|(_, field)| SystemColumn::from_name(field.name()).is_none_or(|column| keep.contains(&column)))
321 .map(|(index, _)| index)
322 .collect();
323 batch.project(&indices).map_err(arrow_error)
324}
325
326pub const SCALAR_SOURCE_NAME: &str = "reifydb.scalar_source_name";
327
328pub fn scalar_source_name(field: &Field) -> Option<&str> {
329 field.metadata().get(SCALAR_SOURCE_NAME).map(String::as_str)
330}
331
332pub fn is_system_field(field: &Field) -> bool {
333 field.name().starts_with('#')
334}
335
336pub fn user_columns(batch: &RecordBatch) -> impl Iterator<Item = (&FieldRef, &ArrayRef)> {
337 batch.schema_ref().fields().iter().zip(batch.columns()).filter(|(field, _)| !is_system_field(field))
338}
339
340pub fn column_view<'a>(batch: &'a RecordBatch, name: &str) -> Result<Option<ColumnView<'a>>> {
341 match position(batch, name) {
342 Some(index) => ColumnView::try_from((batch.column(index), batch.schema_ref().field(index))).map(Some),
343 None => Ok(None),
344 }
345}
346
347pub fn resolve_column(batch: &RecordBatch, name: &str) -> Option<usize> {
348 let bare = name.strip_prefix('#').unwrap_or(name);
349 SystemColumn::ALL
350 .into_iter()
351 .find(|column| &column.name()[1..] == bare)
352 .and_then(|column| position(batch, column.name()))
353 .or_else(|| position(batch, name))
354 .or_else(|| sumtype_tag_position(batch, name))
355}
356
357pub fn sumtype_tag_position(batch: &RecordBatch, name: &str) -> Option<usize> {
358 let tag = format!("{name}_tag");
359 batch.schema_ref().fields().iter().position(|field| field.name() == &tag && is_sumtype_tag(field))
360}
361
362pub fn check_user_columns(batch: &RecordBatch) -> Result<()> {
363 if batch.num_rows() > 0 && batch.num_columns() == 0 {
364 return Err(field_error(format!("a batch of {} rows has no columns", batch.num_rows())));
365 }
366 Ok(())
367}
368
369fn position(batch: &RecordBatch, name: &str) -> Option<usize> {
370 batch.schema_ref().fields().iter().position(|field| field.name() == name)
371}
372
373fn present<T: 'static>(batch: &RecordBatch, column: SystemColumn) -> Result<Option<&T>> {
374 let Some(array) = system_column(batch, column) else {
375 return Ok(None);
376 };
377 if array.logical_null_count() > 0 {
378 return Err(field_error(format!(
379 "system column {} holds {} none rows",
380 column.name(),
381 array.logical_null_count()
382 )));
383 }
384 array.as_any().downcast_ref::<T>().map(Some).ok_or_else(|| {
385 field_error(format!("system column {} holds arrow type {}", column.name(), array.data_type()))
386 })
387}
388
389fn require(batch: &RecordBatch, column: SystemColumn) -> Result<()> {
390 match system_column(batch, column) {
391 Some(_) => Ok(()),
392 None => Err(field_error(format!("system column {} is missing", column.name()))),
393 }
394}
395
396fn datetime_column(batch: &RecordBatch, column: SystemColumn) -> Result<&[DateTime]> {
397 match present::<TimestampNanosecondArray>(batch, column)? {
398 Some(array) => Ok(datetimes(array)),
399 None => Ok(&[]),
400 }
401}
402
403fn arrow_error(error: ArrowError) -> Error {
404 field_error(error.to_string())
405}
406
407#[cfg(test)]
408mod tests {
409 use arrow_array::{Int32Array, UInt64Array};
410 use arrow_schema::DataType;
411
412 use super::*;
413 use crate::value::{
414 Value,
415 container::{temporal_array::datetime_array, wide_int_array::wide_array},
416 };
417
418 fn user(name: &str, values: Vec<i32>) -> (FieldRef, ArrayRef) {
419 let field = to_field(
420 name,
421 &FieldType {
422 value_type: Some(ValueType::Int4),
423 ..FieldType::default()
424 },
425 );
426 (Arc::new(field), Arc::new(Int32Array::from(values)))
427 }
428
429 fn batch(columns: Vec<(FieldRef, ArrayRef)>) -> RecordBatch {
430 let (fields, arrays): (Vec<FieldRef>, Vec<ArrayRef>) = columns.into_iter().unzip();
431 RecordBatch::try_new(Arc::new(Schema::new(fields)), arrays).unwrap()
432 }
433
434 fn names(batch: &RecordBatch) -> Vec<String> {
435 batch.schema_ref().fields().iter().map(|field| field.name().clone()).collect()
436 }
437
438 fn stamps(values: Vec<i64>) -> ArrayRef {
439 Arc::new(datetime_array(values.into_iter().map(DateTime::from_nanos)))
440 }
441
442 #[test]
443 fn row_numbers_read_the_rownum_column_and_are_empty_when_it_is_absent() {
444 let plain = batch(vec![user("a", vec![1, 2])]);
446 assert!(row_numbers(&plain).unwrap().is_empty());
447 let numbered =
448 with_system_column(plain, SystemColumn::RowNumbers, Arc::new(UInt64Array::from(vec![7u64, 9])))
449 .unwrap();
450 assert_eq!(row_numbers(&numbered).unwrap(), &[RowNumber(7), RowNumber(9)]);
451 }
452
453 #[test]
454 fn partitions_decode_the_sixteen_byte_column() {
455 let wide: ArrayRef = Arc::new(wide_array([1u128 << 100, 3]));
457 let with =
458 with_system_column(batch(vec![user("a", vec![1, 2])]), SystemColumn::Partitions, wide).unwrap();
459 assert_eq!(partitions(&with).unwrap(), vec![Partition(1u128 << 100), Partition(3)]);
460 }
461
462 #[test]
463 fn datetime_getters_read_their_own_column() {
464 let base = batch(vec![user("a", vec![1])]);
466 let base = with_system_column(base, SystemColumn::CreatedAt, stamps(vec![10])).unwrap();
467 let base = with_system_column(base, SystemColumn::UpdatedAt, stamps(vec![20])).unwrap();
468 let base = with_system_column(base, SystemColumn::Time, stamps(vec![30])).unwrap();
469 let base =
470 with_system_column(base, SystemColumn::CommitVersion, Arc::new(UInt64Array::from(vec![5u64])))
471 .unwrap();
472 assert_eq!(created_at(&base).unwrap(), &[DateTime::from_nanos(10)]);
473 assert_eq!(updated_at(&base).unwrap(), &[DateTime::from_nanos(20)]);
474 assert_eq!(time(&base).unwrap(), &[DateTime::from_nanos(30)]);
475 assert_eq!(commit_versions(&base).unwrap(), &[5]);
476 }
477
478 #[test]
479 fn a_getter_rejects_a_none_row() {
480 let mixed: ArrayRef =
482 Arc::new(TimestampNanosecondArray::from(vec![Some(1), None]).with_timezone("+00:00"));
483 let with = with_system_column(batch(vec![user("a", vec![1, 2])]), SystemColumn::Time, mixed).unwrap();
484 assert!(with.schema_ref().field_with_name("#time").unwrap().is_nullable());
485 assert!(time(&with).is_err());
486 }
487
488 #[test]
489 fn a_getter_rejects_a_column_of_the_wrong_arrow_type() {
490 let (_, wrong) = user("#rownum", vec![1]);
492 let field = Arc::new(Field::new("#rownum", wrong.data_type().clone(), false));
493 let bad = batch(vec![(field, wrong)]);
494 assert!(row_numbers(&bad).is_err());
495 }
496
497 #[test]
498 fn system_columns_are_inserted_after_user_columns_in_all_order() {
499 let base = batch(vec![user("a", vec![1]), user("b", vec![2])]);
501 let base = with_system_column(base, SystemColumn::Time, stamps(vec![3])).unwrap();
502 let base = with_system_column(base, SystemColumn::RowNumbers, Arc::new(UInt64Array::from(vec![1u64])))
503 .unwrap();
504 let base = with_system_column(base, SystemColumn::CreatedAt, stamps(vec![4])).unwrap();
505 assert_eq!(names(&base), ["a", "b", "#rownum", "#created_at", "#time"]);
506 }
507
508 #[test]
509 fn a_present_system_column_is_replaced_in_place() {
510 let base = with_system_column(
512 batch(vec![user("a", vec![1])]),
513 SystemColumn::RowNumbers,
514 Arc::new(UInt64Array::from(vec![1u64])),
515 )
516 .unwrap();
517 let replaced =
518 with_system_column(base, SystemColumn::RowNumbers, Arc::new(UInt64Array::from(vec![8u64])))
519 .unwrap();
520 assert_eq!(names(&replaced), ["a", "#rownum"]);
521 assert_eq!(row_numbers(&replaced).unwrap(), &[RowNumber(8)]);
522 }
523
524 #[test]
525 fn a_system_column_on_an_empty_batch_sets_the_row_count() {
526 let empty = RecordBatch::new_empty(Arc::new(Schema::empty()));
528 let with = with_system_column(
529 empty,
530 SystemColumn::RowNumbers,
531 Arc::new(UInt64Array::from(vec![1u64, 2, 3])),
532 )
533 .unwrap();
534 assert_eq!(with.num_rows(), 3);
535 }
536
537 #[test]
538 fn a_system_column_of_the_wrong_length_is_rejected() {
539 let result = with_system_column(
541 batch(vec![user("a", vec![1, 2])]),
542 SystemColumn::RowNumbers,
543 Arc::new(UInt64Array::from(vec![1u64])),
544 );
545 assert!(result.is_err());
546 }
547
548 #[test]
549 fn keep_drops_only_unnamed_system_columns() {
550 let op = Arc::new(Field::new("#op", DataType::Int32, false));
552 let base = batch(vec![user("a", vec![1]), (op, Arc::new(Int32Array::from(vec![1])))]);
553 let base = with_system_column(base, SystemColumn::RowNumbers, Arc::new(UInt64Array::from(vec![1u64])))
554 .unwrap();
555 let base = with_system_column(base, SystemColumn::CreatedAt, stamps(vec![1])).unwrap();
556 let base = with_system_column(base, SystemColumn::Time, stamps(vec![1])).unwrap();
557 let kept = keep_system_columns(&base, &[SystemColumn::RowNumbers, SystemColumn::Time]).unwrap();
558 assert_eq!(names(&kept), ["a", "#op", "#rownum", "#time"]);
559 assert_eq!(kept.num_rows(), 1);
560 }
561
562 #[test]
563 fn a_bare_name_resolves_to_the_system_column_first() {
564 let base = batch(vec![user("rownum", vec![5]), user("x", vec![6])]);
566 let base = with_system_column(base, SystemColumn::RowNumbers, Arc::new(UInt64Array::from(vec![1u64])))
567 .unwrap();
568 assert_eq!(resolve_column(&base, "rownum"), Some(2));
569 assert_eq!(resolve_column(&base, "#rownum"), Some(2));
570 assert_eq!(resolve_column(&base, "x"), Some(1));
571 assert_eq!(resolve_column(&base, "created_at"), None);
572 assert_eq!(resolve_column(&base, "missing"), None);
573 }
574
575 #[test]
576 fn rows_without_any_column_are_rejected() {
577 let counted = RecordBatch::try_new_with_options(
579 Arc::new(Schema::empty()),
580 vec![],
581 &RecordBatchOptions::new().with_row_count(Some(2)),
582 )
583 .unwrap();
584 assert!(check_user_columns(&counted).is_err());
585 let named = with_system_column(
586 RecordBatch::new_empty(Arc::new(Schema::empty())),
587 SystemColumn::RowNumbers,
588 Arc::new(UInt64Array::from(vec![1u64, 2])),
589 )
590 .unwrap();
591 assert!(check_user_columns(&named).is_ok());
592 assert!(check_user_columns(&RecordBatch::new_empty(Arc::new(Schema::empty()))).is_ok());
593 }
594
595 #[test]
596 fn a_hash_prefix_marks_a_system_field() {
597 assert!(is_system_field(&Field::new("#rownum", DataType::UInt64, false)));
599 assert!(is_system_field(&Field::new("#op", DataType::Int32, false)));
600 assert!(!is_system_field(&Field::new("rownum", DataType::UInt64, false)));
601 }
602
603 #[test]
604 fn column_view_finds_a_column_by_its_exact_name() {
605 let base = batch(vec![user("rownum", vec![5])]);
607 let base = with_system_column(base, SystemColumn::RowNumbers, Arc::new(UInt64Array::from(vec![9u64])))
608 .unwrap();
609 let user_view = column_view(&base, "rownum").unwrap().unwrap();
610 assert_eq!(user_view.get_value(0), Value::Int4(5));
611 let system_view = column_view(&base, "#rownum").unwrap().unwrap();
612 assert_eq!(system_view.get_value(0), Value::Uint8(9));
613 assert!(column_view(&base, "missing").unwrap().is_none());
614 }
615
616 #[test]
617 fn require_getters_fail_naming_the_absent_column() {
618 let plain = batch(vec![user("a", vec![1])]);
620 for (result, name) in [
621 (require_row_numbers(&plain).map(|_| ()), "#rownum"),
622 (require_created_at(&plain).map(|_| ()), "#created_at"),
623 (require_updated_at(&plain).map(|_| ()), "#updated_at"),
624 (require_time(&plain).map(|_| ()), "#time"),
625 ] {
626 let message = format!("{:?}", result.unwrap_err());
627 assert!(message.contains(name), "{name} missing from {message}");
628 }
629 }
630
631 #[test]
632 fn require_getters_read_a_present_column() {
633 let base = batch(vec![user("a", vec![1])]);
635 let base = with_system_column(base, SystemColumn::RowNumbers, Arc::new(UInt64Array::from(vec![4u64])))
636 .unwrap();
637 let base = with_system_column(base, SystemColumn::CreatedAt, stamps(vec![10])).unwrap();
638 let base = with_system_column(base, SystemColumn::UpdatedAt, stamps(vec![20])).unwrap();
639 let base = with_system_column(base, SystemColumn::Time, stamps(vec![30])).unwrap();
640 assert_eq!(require_row_numbers(&base).unwrap(), &[RowNumber(4)]);
641 assert_eq!(require_created_at(&base).unwrap(), &[DateTime::from_nanos(10)]);
642 assert_eq!(require_updated_at(&base).unwrap(), &[DateTime::from_nanos(20)]);
643 assert_eq!(require_time(&base).unwrap(), &[DateTime::from_nanos(30)]);
644 }
645
646 #[test]
647 fn require_getters_still_reject_a_none_row() {
648 let mixed: ArrayRef =
650 Arc::new(TimestampNanosecondArray::from(vec![Some(1), None]).with_timezone("+00:00"));
651 let with = with_system_column(batch(vec![user("a", vec![1, 2])]), SystemColumn::Time, mixed).unwrap();
652 assert!(require_time(&with).is_err());
653 }
654
655 #[test]
656 fn user_columns_skip_every_hash_column_in_order() {
657 let op = Arc::new(Field::new("#op", DataType::Int32, false));
659 let base =
660 batch(vec![user("a", vec![1]), (op, Arc::new(Int32Array::from(vec![2]))), user("b", vec![3])]);
661 let base = with_system_column(base, SystemColumn::RowNumbers, Arc::new(UInt64Array::from(vec![1u64])))
662 .unwrap();
663 let names: Vec<&str> = user_columns(&base).map(|(field, _)| field.name().as_str()).collect();
664 assert_eq!(names, ["a", "b"]);
665 let (_, b) = user_columns(&base).nth(1).unwrap();
666 assert_eq!(b.as_any().downcast_ref::<Int32Array>().unwrap().value(0), 3);
667 assert_eq!(user_columns(&RecordBatch::new_empty(Arc::new(Schema::empty()))).count(), 0);
668 }
669}