Skip to content

Commit 4b09619

Browse files
authored
Merge pull request #1 from alamb/alamb/bound
Use newtype to get hash/eq for f64
2 parents edceba3 + 2a6f0d6 commit 4b09619

File tree

1 file changed

+41
-62
lines changed

1 file changed

+41
-62
lines changed

datafusion-examples/examples/table_sample.rs

Lines changed: 41 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -108,78 +108,57 @@ use rand::rngs::StdRng;
108108
use rand::{Rng, SeedableRng};
109109
use rand_distr::{Distribution, Poisson};
110110

111-
/// This example demonstrates the table sample support.
112111

113-
#[derive(Debug, Clone)]
114-
struct TableSamplePlanNode {
115-
inner_plan: LogicalPlan,
116112

117-
lower_bound: f64,
118-
upper_bound: f64,
119-
with_replacement: bool,
120-
seed: u64,
113+
/// This example demonstrates the table sample support.
114+
115+
/// Hashable and comparible f64 for sampling bounds
116+
#[derive(Debug, Clone, Copy, PartialOrd)]
117+
struct Bound(f64);
118+
impl PartialEq for Bound {
119+
fn eq(&self, other: &Self) -> bool {
120+
(self.0 - other.0).abs() < f64::EPSILON
121+
}
121122
}
122123

123-
impl Hash for TableSamplePlanNode {
124+
impl Eq for Bound {}
125+
126+
impl Hash for Bound {
124127
fn hash<H: Hasher>(&self, state: &mut H) {
125-
self.inner_plan.hash(state);
126-
self.lower_bound.to_bits().hash(state);
127-
self.upper_bound.to_bits().hash(state);
128-
self.with_replacement.hash(state);
129-
self.seed.hash(state);
128+
// Hash the bits of the f64
129+
self.0.to_bits().hash(state);
130130
}
131131
}
132132

133-
impl PartialEq for TableSamplePlanNode {
134-
fn eq(&self, other: &Self) -> bool {
135-
self.inner_plan == other.inner_plan
136-
&& (self.lower_bound - other.lower_bound).abs() < f64::EPSILON
137-
&& (self.upper_bound - other.upper_bound).abs() < f64::EPSILON
138-
&& self.with_replacement == other.with_replacement
139-
&& self.seed == other.seed
133+
impl From<f64> for Bound {
134+
fn from(value: f64) -> Self {
135+
Self(value)
136+
}
137+
}
138+
impl From<Bound> for f64 {
139+
fn from(value: Bound) -> Self {
140+
value.0
140141
}
141142
}
142143

143-
impl Eq for TableSamplePlanNode {}
144-
145-
impl PartialOrd for TableSamplePlanNode {
146-
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
147-
self.inner_plan
148-
.partial_cmp(&other.inner_plan)
149-
.and_then(|ord| {
150-
if ord != Ordering::Equal {
151-
Some(ord)
152-
} else {
153-
self.lower_bound
154-
.partial_cmp(&other.lower_bound)
155-
.and_then(|ord| {
156-
if ord != Ordering::Equal {
157-
Some(ord)
158-
} else {
159-
self.upper_bound.partial_cmp(&other.upper_bound).and_then(
160-
|ord| {
161-
if ord != Ordering::Equal {
162-
Some(ord)
163-
} else {
164-
self.with_replacement
165-
.partial_cmp(&other.with_replacement)
166-
.and_then(|ord| {
167-
if ord != Ordering::Equal {
168-
Some(ord)
169-
} else {
170-
self.seed.partial_cmp(&other.seed)
171-
}
172-
})
173-
}
174-
},
175-
)
176-
}
177-
})
178-
}
179-
})
144+
impl AsRef<f64> for Bound {
145+
fn as_ref(&self) -> &f64 {
146+
&self.0
180147
}
181148
}
182149

150+
151+
#[derive(Debug, Clone, Hash, Eq, PartialEq, PartialOrd)]
152+
struct TableSamplePlanNode {
153+
inner_plan: LogicalPlan,
154+
155+
lower_bound: Bound,
156+
upper_bound: Bound,
157+
with_replacement: bool,
158+
seed: u64,
159+
}
160+
161+
183162
impl UserDefinedLogicalNodeCore for TableSamplePlanNode {
184163
fn name(&self) -> &str {
185164
"TableSample"
@@ -234,8 +213,8 @@ impl TableSampleExtensionPlanner {
234213
Ok(Arc::new(SampleExec {
235214
input: physical_input.clone(),
236215
lower_bound: 0.0,
237-
upper_bound: specific_node.upper_bound,
238-
with_replacement: specific_node.with_replacement,
216+
upper_bound: specific_node.upper_bound.into(),
217+
with_replacement: specific_node.with_replacement.into(),
239218
seed: specific_node.seed,
240219
metrics: Default::default(),
241220
cache: SampleExec::compute_properties(&physical_input),
@@ -663,8 +642,8 @@ impl<'a, S: ContextProvider> TableSamplePlanner<'a, S> {
663642
) -> Result<LogicalPlan> {
664643
let node = TableSamplePlanNode {
665644
inner_plan: input,
666-
lower_bound: 0.0,
667-
upper_bound: fraction,
645+
lower_bound: Bound::from(0.0),
646+
upper_bound: Bound::from(fraction),
668647
with_replacement: with_replacement.unwrap_or(false),
669648
seed: seed.unwrap_or_else(rand::random),
670649
};

0 commit comments

Comments
 (0)