Skip to content

Commit 392e9c5

Browse files
jackkleemanjackkleeman
authored
Avoid many option checks for unwritten rows (#4335)
1 parent 26096b5 commit 392e9c5

File tree

1 file changed

+37
-50
lines changed

1 file changed

+37
-50
lines changed

‎crates/storage-query-datafusion/src/table_macro.rs‎

Lines changed: 37 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -123,7 +123,10 @@ macro_rules! define_fmt {
123123
use std::fmt::Write;
124124
let _ = write!(builder, "{d}");
125125
builder.append_value("");
126-
self.flags.$element = true;
126+
if !self.flags.$element {
127+
self.flags.$element = true;
128+
self.flags.written_count += 1;
129+
}
127130
}
128131
}
129132
};
@@ -134,7 +137,10 @@ macro_rules! define_fmt {
134137
use std::fmt::Write;
135138
let _ = write!(builder, "{d}");
136139
builder.append_value("");
137-
self.flags.$element = true;
140+
if !self.flags.$element {
141+
self.flags.$element = true;
142+
self.flags.written_count += 1;
143+
}
138144
}
139145
}
140146
};
@@ -382,6 +388,7 @@ macro_rules! document_type {
382388
/// ```ignore
383389
/// pub struct UserBuilder {
384390
/// rows_inserted_so_far: usize,
391+
/// projected_count: u16,
385392
/// projected_schema: SchemaRef,
386393
/// arrays: UserArrayBuilder,
387394
/// }
@@ -397,6 +404,7 @@ macro_rules! document_type {
397404
/// }
398405
/// #[derive(Default)]
399406
/// struct UserRowBuilderFlags {
407+
/// written_count: u16,
400408
/// name: bool,
401409
/// age: bool,
402410
/// secret: bool,
@@ -407,7 +415,10 @@ macro_rules! document_type {
407415
/// pub fn name(&mut self, value: impl AsRef<str>) {
408416
/// if let Some(builder) = self.builder.arrays.name.as_mut() {
409417
/// builder.append_value(value);
410-
/// self.flags.name = true;
418+
/// if !self.flags.name {
419+
/// self.flags.name = true;
420+
/// self.flags.written_count += 1;
421+
/// }
411422
/// }
412423
/// }
413424
/// #[inline]
@@ -416,71 +427,34 @@ macro_rules! document_type {
416427
/// use std::fmt::Write;
417428
/// let _ = builder.write_fmt(core::format_args!("{d}"));
418429
/// builder.append_value("");
419-
/// self.flags.name = true;
430+
/// if !self.flags.name {
431+
/// self.flags.name = true;
432+
/// self.flags.written_count += 1;
433+
/// }
420434
/// }
421435
/// }
422436
/// #[inline]
423437
/// pub fn is_name_defined(&self) -> bool {
424438
/// self.builder.arrays.name.is_some()
425439
/// }
426-
/// #[inline]
427-
/// pub fn age(&mut self, value: u32) {
428-
/// if let Some(builder) = self.builder.arrays.age.as_mut() {
429-
/// builder.append_value(value);
430-
/// self.flags.age = true;
431-
/// }
432-
/// }
433440
///
434441
/// #[inline]
435442
/// pub fn is_age_defined(&self) -> bool {
436443
/// self.builder.arrays.age.is_some()
437444
/// }
438-
/// #[inline]
439-
/// pub fn secret(&mut self, value: impl AsRef<[u8]>) {
440-
/// if let Some(builder) = self.builder.arrays.secret.as_mut() {
441-
/// builder.append_value(value);
442-
/// self.flags.secret = true;
443-
/// }
444-
/// }
445-
/// #[inline]
446-
/// pub fn is_secret_defined(&self) -> bool {
447-
/// self.builder.arrays.secret.is_some()
448-
/// }
449-
/// #[inline]
450-
/// pub fn birth_date(&mut self, value: i64) {
451-
/// if let Some(builder) = self.builder.arrays.birth_date.as_mut() {
452-
/// builder.append_value(value);
453-
/// self.flags.birth_date = true;
454-
/// }
455-
/// }
456-
///
457-
/// #[inline]
458-
/// pub fn is_birth_date_defined(&self) -> bool {
459-
/// self.builder.arrays.birth_date.is_some()
460-
/// }
445+
/// // ...
461446
/// }
462447
/// impl<'a> Drop for UserRowBuilder<'a> {
463448
/// fn drop(&mut self) {
449+
/// if self.flags.written_count == self.builder.projected_count {
450+
/// return;
451+
/// }
464452
/// if let Some(e) = self.builder.arrays.name.as_mut() {
465453
/// if (!self.flags.name) {
466454
/// e.append_null();
467455
/// }
468456
/// }
469-
/// if let Some(e) = self.builder.arrays.age.as_mut() {
470-
/// if (!self.flags.age) {
471-
/// e.append_null();
472-
/// }
473-
/// }
474-
/// if let Some(e) = self.builder.arrays.secret.as_mut() {
475-
/// if (!self.flags.secret) {
476-
/// e.append_null();
477-
/// }
478-
/// }
479-
/// if let Some(e) = self.builder.arrays.birth_date.as_mut() {
480-
/// if (!self.flags.birth_date) {
481-
/// e.append_null();
482-
/// }
483-
/// }
457+
/// // ...
484458
/// }
485459
/// }
486460
/// impl UserArrayBuilder {
@@ -552,8 +526,10 @@ macro_rules! document_type {
552526
/// }
553527
/// impl UserBuilder {
554528
/// pub fn new(projected_schema: &Schema) -> Self {
529+
/// let projected_count = projected_schema.fields().len() as u16;
555530
/// Self {
556531
/// rows_inserted_so_far: 0,
532+
/// projected_count,
557533
/// arrays: UserArrayBuilder::new(&projected_schema),
558534
/// projected_schema,
559535
/// }
@@ -616,6 +592,7 @@ macro_rules! define_table {
616592

617593
pub struct [< $table_name:camel Builder >] {
618594
rows_inserted_so_far: usize,
595+
projected_count: u16,
619596
projected_schema: ::datafusion::arrow::datatypes::SchemaRef,
620597
arrays: [< $table_name:camel ArrayBuilder >],
621598
}
@@ -634,6 +611,7 @@ macro_rules! define_table {
634611

635612
#[derive(Default)]
636613
struct [< $table_name:camel RowBuilderFlags >] {
614+
written_count: u16,
637615
$($element : bool, )+
638616
}
639617
// --------------------------------------------------------------------------
@@ -649,7 +627,10 @@ macro_rules! define_table {
649627
pub fn $element(&mut self, value: define_primitive_trait!($ty)) {
650628
if let Some(builder) = self.builder.arrays.$element.as_mut() {
651629
builder.append_value(value);
652-
self.flags.$element = true;
630+
if !self.flags.$element {
631+
self.flags.$element = true;
632+
self.flags.written_count += 1;
633+
}
653634
}
654635
}
655636

@@ -668,6 +649,10 @@ macro_rules! define_table {
668649
impl<'a> Drop for [< $table_name:camel RowBuilder >]<'a> {
669650

670651
fn drop(&mut self) {
652+
// Fast path: all projected fields were written
653+
if self.flags.written_count == self.builder.projected_count {
654+
return;
655+
}
671656

672657
$(
673658
if let Some(e) = self.builder.arrays.$element.as_mut() {
@@ -764,8 +749,10 @@ macro_rules! define_table {
764749
#[allow(clippy::all)]
765750
impl $crate::table_util::Builder for [< $table_name:camel Builder >] {
766751
fn new(projected_schema: ::datafusion::arrow::datatypes::SchemaRef) -> Self {
752+
let projected_count = projected_schema.fields().len() as u16;
767753
Self {
768754
rows_inserted_so_far: 0,
755+
projected_count,
769756
arrays: [< $table_name:camel ArrayBuilder >]::new(&projected_schema),
770757
projected_schema,
771758
}

0 commit comments

Comments
 (0)