init mvp

Author: suxiaogang223

examples/row_selection_example.rs

@@ -0,0 +1,147 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+//   Unless required by applicable law or agreed to in writing,
+//   software distributed under the License is distributed on an
+//   "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+//   KIND, either express or implied.  See the License for the
+//   specific language governing permissions and limitations
+//   under the License.
+
+//! Example demonstrating how to use RowSelection to skip rows when reading ORC files
+
+use std::fs::File;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, Int32Array, RecordBatch, StringArray};
+use arrow::datatypes::{DataType, Field, Schema};
+use orc_rust::arrow_reader::ArrowReaderBuilder;
+use orc_rust::arrow_writer::ArrowWriterBuilder;
+use orc_rust::row_selection::{RowSelection, RowSelector};
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    // Step 1: Create a sample ORC file with 100 rows
+    println!("Creating sample ORC file...");
+    let file_path = "/tmp/row_selection_example.orc";
+    create_sample_orc_file(file_path)?;
+
+    // Step 2: Read the file without row selection (baseline)
+    println!("\n=== Reading all rows (no selection) ===");
+    let file = File::open(file_path)?;
+    let reader = ArrowReaderBuilder::try_new(file)?.build();
+    let mut total_rows = 0;
+    for batch in reader {
+        let batch = batch?;
+        total_rows += batch.num_rows();
+    }
+    println!("Total rows read: {}", total_rows);
+
+    // Step 3: Read with row selection - skip first 30 rows, select next 40, skip rest
+    println!("\n=== Reading with row selection ===");
+    let file = File::open(file_path)?;
+    
+    // Create a selection: skip 30, select 40, skip 30
+    let selection = vec![
+        RowSelector::skip(30),
+        RowSelector::select(40),
+        RowSelector::skip(30),
+    ]
+    .into();
+    
+    let reader = ArrowReaderBuilder::try_new(file)?
+        .with_row_selection(selection)
+        .build();
+    
+    let mut selected_rows = 0;
+    let mut batches = Vec::new();
+    for batch in reader {
+        let batch = batch?;
+        selected_rows += batch.num_rows();
+        batches.push(batch);
+    }
+    
+    println!("Total rows selected: {}", selected_rows);
+    println!("Expected: 40, Actual: {}", selected_rows);
+    
+    // Display some of the selected data
+    if let Some(first_batch) = batches.first() {
+        let id_col = first_batch
+            .column(0)
+            .as_any()
+            .downcast_ref::<Int32Array>()
+            .unwrap();
+        let name_col = first_batch
+            .column(1)
+            .as_any()
+            .downcast_ref::<StringArray>()
+            .unwrap();
+        
+        println!("\nFirst 5 selected rows:");
+        for i in 0..5.min(first_batch.num_rows()) {
+            println!(
+                "  id: {}, name: {}",
+                id_col.value(i),
+                name_col.value(i)
+            );
+        }
+    }
+
+    // Step 4: Read with multiple non-consecutive selections
+    println!("\n=== Reading with multiple selections ===");
+    let file = File::open(file_path)?;
+    
+    // Select rows 10-20 and 60-70
+    let selection = RowSelection::from_consecutive_ranges(
+        vec![10..20, 60..70].into_iter(),
+        100,
+    );
+    
+    let reader = ArrowReaderBuilder::try_new(file)?
+        .with_row_selection(selection)
+        .build();
+    
+    let mut selected_rows = 0;
+    for batch in reader {
+        let batch = batch?;
+        selected_rows += batch.num_rows();
+    }
+    
+    println!("Total rows selected: {}", selected_rows);
+    println!("Expected: 20 (10 from each range)");
+
+    println!("\n✓ Row selection example completed successfully!");
+    
+    Ok(())
+}
+
+fn create_sample_orc_file(path: &str) -> Result<(), Box<dyn std::error::Error>> {
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("id", DataType::Int32, false),
+        Field::new("name", DataType::Utf8, false),
+    ]));
+
+    // Create 100 rows
+    let ids: ArrayRef = Arc::new(Int32Array::from((0..100).collect::<Vec<i32>>()));
+    let names: ArrayRef = Arc::new(StringArray::from(
+        (0..100)
+            .map(|i| format!("name_{}", i))
+            .collect::<Vec<String>>(),
+    ));
+
+    let batch = RecordBatch::try_new(schema.clone(), vec![ids, names])?;
+
+    let file = File::create(path)?;
+    let mut writer = ArrowWriterBuilder::new(file, schema).try_build()?;
+    writer.write(&batch)?;
+    writer.close()?;
+
+    Ok(())
+}
+

src/array_decoder/mod.rs

@@ -238,18 +238,70 @@ pub struct NaiveStripeDecoder {
     index: usize,
     batch_size: usize,
     number_of_rows: usize,
+    row_selection: Option<crate::row_selection::RowSelection>,
+    selection_index: usize,
 }
 
 impl Iterator for NaiveStripeDecoder {
     type Item = Result<RecordBatch>;
 
     fn next(&mut self) -> Option<Self::Item> {
         if self.index < self.number_of_rows {
-            let record = self
-                .decode_next_batch(self.number_of_rows - self.index)
-                .transpose()?;
-            self.index += self.batch_size;
-            Some(record)
+            // Handle row selection if present
+            if self.row_selection.is_some() {
+                // Process selectors until we find rows to select or exhaust the selection
+                loop {
+                    let selector_info = {
+                        let selection = self.row_selection.as_ref().unwrap();
+                        let selectors = selection.selectors();
+                        if self.selection_index >= selectors.len() {
+                            return None;
+                        }
+                        let selector = selectors[self.selection_index];
+                        (selector.skip, selector.row_count)
+                    };
+                    
+                    let (is_skip, row_count) = selector_info;
+                    
+                    if is_skip {
+                        // Skip these rows by advancing the index
+                        self.index += row_count;
+                        self.selection_index += 1;
+                        
+                        // Decode and discard the skipped rows to advance the internal decoders
+                        if let Err(e) = self.skip_rows(row_count) {
+                            return Some(Err(e));
+                        }
+                    } else {
+                        // Select these rows
+                        let rows_to_read = row_count.min(self.batch_size);
+                        let remaining = self.number_of_rows - self.index;
+                        let actual_rows = rows_to_read.min(remaining);
+                        
+                        if actual_rows == 0 {
+                            self.selection_index += 1;
+                            continue;
+                        }
+                        
+                        let record = self.decode_next_batch(actual_rows).transpose()?;
+                        self.index += actual_rows;
+                        
+                        // Update selector to track progress
+                        if actual_rows >= row_count {
+                            self.selection_index += 1;
+                        }
+                        
+                        return Some(record);
+                    }
+                }
+            } else {
+                // No row selection - decode normally
+                let record = self
+                    .decode_next_batch(self.number_of_rows - self.index)
+                    .transpose()?;
+                self.index += self.batch_size;
+                Some(record)
+            }
         } else {
             None
         }
@@ -433,6 +485,15 @@ impl NaiveStripeDecoder {
     }
 
     pub fn new(stripe: Stripe, schema_ref: SchemaRef, batch_size: usize) -> Result<Self> {
+        Self::new_with_selection(stripe, schema_ref, batch_size, None)
+    }
+
+    pub fn new_with_selection(
+        stripe: Stripe,
+        schema_ref: SchemaRef,
+        batch_size: usize,
+        row_selection: Option<crate::row_selection::RowSelection>,
+    ) -> Result<Self> {
         let number_of_rows = stripe.number_of_rows();
         let decoders = stripe
             .columns()
@@ -448,6 +509,20 @@ impl NaiveStripeDecoder {
             index: 0,
             batch_size,
             number_of_rows,
+            row_selection,
+            selection_index: 0,
         })
     }
+
+    /// Skip the specified number of rows by decoding and discarding them
+    fn skip_rows(&mut self, count: usize) -> Result<()> {
+        // Decode in batches to avoid large memory allocations
+        let mut remaining = count;
+        while remaining > 0 {
+            let chunk = self.batch_size.min(remaining);
+            let _ = self.inner_decode_next_batch(chunk)?;
+            remaining -= chunk;
+        }
+        Ok(())
+    }
 }

src/arrow_reader.rs

@@ -28,6 +28,7 @@ use crate::error::Result;
 use crate::projection::ProjectionMask;
 use crate::reader::metadata::{read_metadata, FileMetadata};
 use crate::reader::ChunkReader;
+use crate::row_selection::RowSelection;
 use crate::schema::RootDataType;
 use crate::stripe::{Stripe, StripeMetadata};
 
@@ -40,6 +41,7 @@ pub struct ArrowReaderBuilder<R> {
     pub(crate) projection: ProjectionMask,
     pub(crate) schema_ref: Option<SchemaRef>,
     pub(crate) file_byte_range: Option<Range<usize>>,
+    pub(crate) row_selection: Option<RowSelection>,
 }
 
 impl<R> ArrowReaderBuilder<R> {
@@ -51,6 +53,7 @@ impl<R> ArrowReaderBuilder<R> {
             projection: ProjectionMask::all(),
             schema_ref: None,
             file_byte_range: None,
+            row_selection: None,
         }
     }
 
@@ -79,6 +82,33 @@ impl<R> ArrowReaderBuilder<R> {
         self
     }
 
+    /// Set a [`RowSelection`] to filter rows
+    ///
+    /// The [`RowSelection`] specifies which rows should be decoded from the ORC file.
+    /// This can be used to skip rows that don't match predicates, reducing I/O and
+    /// improving query performance.
+    ///
+    /// # Example
+    ///
+    /// ```no_run
+    /// # use std::fs::File;
+    /// # use orc_rust::arrow_reader::ArrowReaderBuilder;
+    /// # use orc_rust::row_selection::{RowSelection, RowSelector};
+    /// let file = File::open("data.orc").unwrap();
+    /// let selection = vec![
+    ///     RowSelector::skip(100),
+    ///     RowSelector::select(50),
+    /// ].into();
+    /// let reader = ArrowReaderBuilder::try_new(file)
+    ///     .unwrap()
+    ///     .with_row_selection(selection)
+    ///     .build();
+    /// ```
+    pub fn with_row_selection(mut self, row_selection: RowSelection) -> Self {
+        self.row_selection = Some(row_selection);
+        self
+    }
+
     /// Returns the currently computed schema
     ///
     /// Unless [`with_schema`](Self::with_schema) was called, this is computed dynamically
@@ -124,6 +154,7 @@ impl<R: ChunkReader> ArrowReaderBuilder<R> {
             schema_ref,
             current_stripe: None,
             batch_size: self.batch_size,
+            row_selection: self.row_selection,
         }
     }
 }
@@ -133,6 +164,7 @@ pub struct ArrowReader<R> {
     schema_ref: SchemaRef,
     current_stripe: Option<Box<dyn Iterator<Item = Result<RecordBatch>> + Send>>,
     batch_size: usize,
+    row_selection: Option<RowSelection>,
 }
 
 impl<R> ArrowReader<R> {
@@ -146,8 +178,22 @@ impl<R: ChunkReader> ArrowReader<R> {
         let stripe = self.cursor.next().transpose()?;
         match stripe {
             Some(stripe) => {
-                let decoder =
-                    NaiveStripeDecoder::new(stripe, self.schema_ref.clone(), self.batch_size)?;
+                // Split off the row selection for this stripe
+                let stripe_rows = stripe.number_of_rows();
+                let selection = self.row_selection.as_mut().and_then(|s| {
+                    if s.row_count() > 0 {
+                        Some(s.split_off(stripe_rows))
+                    } else {
+                        None
+                    }
+                });
+
+                let decoder = NaiveStripeDecoder::new_with_selection(
+                    stripe,
+                    self.schema_ref.clone(),
+                    self.batch_size,
+                    selection,
+                )?;
                 self.current_stripe = Some(Box::new(decoder));
                 self.next().transpose()
             }

src/lib.rs

@@ -61,6 +61,7 @@ pub mod projection;
 #[allow(dead_code)]
 mod proto;
 pub mod reader;
+pub mod row_selection;
 pub mod schema;
 pub mod statistics;
 pub mod stripe;
@@ -70,3 +71,4 @@ pub use arrow_reader::{ArrowReader, ArrowReaderBuilder};
 pub use arrow_writer::{ArrowWriter, ArrowWriterBuilder};
 #[cfg(feature = "async")]
 pub use async_arrow_reader::ArrowStreamReader;
+pub use row_selection::{RowSelection, RowSelector};