fix quantiles

Eventual-Inc · Oct 18, 2023 · 4bcded3 · 4bcded3
1 parent 9566ae6
commit 4bcded3
Showing 1 changed file with 4 additions and 6 deletions.
diff --git a/src/daft-micropartition/src/ops/take.rs b/src/daft-micropartition/src/ops/take.rs
@@ -28,12 +28,11 @@ impl MicroPartition {
         let tables = self.concat_or_get()?;
         if let [single] = tables.as_slice() {
             let taken = single.sample(num)?;
+            let taken_len = taken.len();
             Ok(Self::new(
                 self.schema.clone(),
                 TableState::Loaded(Arc::new(vec![taken])),
-                TableMetadata {
-                    length: num.min(self.len()),
-                },
+                TableMetadata { length: taken_len },
                 self.statistics.clone(),
             ))
         } else {
@@ -45,12 +44,11 @@ impl MicroPartition {
         let tables = self.concat_or_get()?;
         if let [single] = tables.as_slice() {
             let taken = single.quantiles(num)?;
+            let taken_len = taken.len();
             Ok(Self::new(
                 self.schema.clone(),
                 TableState::Loaded(Arc::new(vec![taken])),
-                TableMetadata {
-                    length: (num - 1).max(0),
-                },
+                TableMetadata { length: taken_len },
                 self.statistics.clone(),
             ))
         } else {