Merge branch 'main' into check-duplicate

Author: jonathanc-n

.github/workflows/audit.yml

@@ -31,6 +31,7 @@ on:
     paths:
       - "**/Cargo.toml"
       - "**/Cargo.lock"
+      - ".github/workflows/audit.yml"
 
   schedule:
     - cron: '0 0 * * *'
@@ -40,6 +41,10 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
+      - name: Setup Rust toolchain
+        uses: ./.github/actions/setup-builder
+        with:
+          rust-version: stable
       - uses: rustsec/audit-check@v2.0.0
         with:
           token: ${{ secrets.GITHUB_TOKEN }}

crates/catalog/glue/src/error.rs

@@ -25,7 +25,7 @@ pub(crate) fn from_aws_sdk_error<T>(error: aws_sdk_glue::error::SdkError<T>) ->
 where T: Debug {
     Error::new(
         ErrorKind::Unexpected,
-        "Operation failed for hitting aws skd error".to_string(),
+        "Operation failed for hitting aws sdk error".to_string(),
     )
     .with_source(anyhow!("aws sdk error: {:?}", error))
 }

crates/catalog/s3tables/src/catalog.rs

@@ -501,7 +501,7 @@ pub(crate) fn from_aws_sdk_error<T>(error: aws_sdk_s3tables::error::SdkError<T>)
 where T: std::fmt::Debug {
     Error::new(
         ErrorKind::Unexpected,
-        "Operation failed for hitting aws skd error".to_string(),
+        "Operation failed for hitting aws sdk error".to_string(),
     )
     .with_source(anyhow!("aws sdk error: {:?}", error))
 }

crates/iceberg/src/arrow/reader.rs

@@ -22,7 +22,7 @@ use std::ops::Range;
 use std::str::FromStr;
 use std::sync::Arc;
 
-use arrow_arith::boolean::{and, is_not_null, is_null, not, or};
+use arrow_arith::boolean::{and, and_kleene, is_not_null, is_null, not, or, or_kleene};
 use arrow_array::{Array, ArrayRef, BooleanArray, RecordBatch};
 use arrow_ord::cmp::{eq, gt, gt_eq, lt, lt_eq, neq};
 use arrow_schema::{
@@ -827,7 +827,7 @@ impl<'a> BoundPredicateVisitor for PredicateConverter<'a> {
         Ok(Box::new(move |batch| {
             let left = lhs(batch.clone())?;
             let right = rhs(batch)?;
-            and(&left, &right)
+            and_kleene(&left, &right)
         }))
     }
 
@@ -839,7 +839,7 @@ impl<'a> BoundPredicateVisitor for PredicateConverter<'a> {
         Ok(Box::new(move |batch| {
             let left = lhs(batch.clone())?;
             let right = rhs(batch)?;
-            or(&left, &right)
+            or_kleene(&left, &right)
         }))
     }
 
@@ -1165,18 +1165,29 @@ impl<R: FileRead> AsyncFileReader for ArrowFileReader<R> {
 #[cfg(test)]
 mod tests {
     use std::collections::{HashMap, HashSet};
+    use std::fs::File;
     use std::sync::Arc;
 
+    use arrow_array::cast::AsArray;
+    use arrow_array::{ArrayRef, RecordBatch, StringArray};
     use arrow_schema::{DataType, Field, Schema as ArrowSchema, TimeUnit};
-    use parquet::arrow::ProjectionMask;
+    use futures::TryStreamExt;
+    use parquet::arrow::{ArrowWriter, ProjectionMask};
+    use parquet::basic::Compression;
+    use parquet::file::properties::WriterProperties;
     use parquet::schema::parser::parse_message_type;
     use parquet::schema::types::SchemaDescriptor;
+    use tempfile::TempDir;
 
     use crate::arrow::reader::{CollectFieldIdVisitor, PARQUET_FIELD_ID_META_KEY};
-    use crate::arrow::ArrowReader;
+    use crate::arrow::{ArrowReader, ArrowReaderBuilder};
     use crate::expr::visitors::bound_predicate_visitor::visit;
-    use crate::expr::{Bind, Reference};
-    use crate::spec::{NestedField, PrimitiveType, Schema, SchemaRef, Type};
+    use crate::expr::{Bind, Predicate, Reference};
+    use crate::io::FileIO;
+    use crate::scan::{FileScanTask, FileScanTaskStream};
+    use crate::spec::{
+        DataContentType, DataFileFormat, Datum, NestedField, PrimitiveType, Schema, SchemaRef, Type,
+    };
     use crate::ErrorKind;
 
     fn table_schema_simple() -> SchemaRef {
@@ -1336,4 +1347,138 @@ message schema {
                 .expect("Some ProjectionMask");
         assert_eq!(mask, ProjectionMask::leaves(&parquet_schema, vec![0]));
     }
+
+    #[tokio::test]
+    async fn test_kleene_logic_or_behaviour() {
+        // a IS NULL OR a = 'foo'
+        let predicate = Reference::new("a")
+            .is_null()
+            .or(Reference::new("a").equal_to(Datum::string("foo")));
+
+        // Table data: [NULL, "foo", "bar"]
+        let data_for_col_a = vec![None, Some("foo".to_string()), Some("bar".to_string())];
+
+        // Expected: [NULL, "foo"].
+        let expected = vec![None, Some("foo".to_string())];
+
+        let (file_io, schema, table_location, _temp_dir) = setup_kleene_logic(data_for_col_a);
+        let reader = ArrowReaderBuilder::new(file_io).build();
+
+        let result_data = test_perform_read(predicate, schema, table_location, reader).await;
+
+        assert_eq!(result_data, expected);
+    }
+
+    #[tokio::test]
+    async fn test_kleene_logic_and_behaviour() {
+        // a IS NOT NULL AND a != 'foo'
+        let predicate = Reference::new("a")
+            .is_not_null()
+            .and(Reference::new("a").not_equal_to(Datum::string("foo")));
+
+        // Table data: [NULL, "foo", "bar"]
+        let data_for_col_a = vec![None, Some("foo".to_string()), Some("bar".to_string())];
+
+        // Expected: ["bar"].
+        let expected = vec![Some("bar".to_string())];
+
+        let (file_io, schema, table_location, _temp_dir) = setup_kleene_logic(data_for_col_a);
+        let reader = ArrowReaderBuilder::new(file_io).build();
+
+        let result_data = test_perform_read(predicate, schema, table_location, reader).await;
+
+        assert_eq!(result_data, expected);
+    }
+
+    async fn test_perform_read(
+        predicate: Predicate,
+        schema: SchemaRef,
+        table_location: String,
+        reader: ArrowReader,
+    ) -> Vec<Option<String>> {
+        let tasks = Box::pin(futures::stream::iter(
+            vec![Ok(FileScanTask {
+                start: 0,
+                length: 0,
+                record_count: None,
+                data_file_path: format!("{}/1.parquet", table_location),
+                data_file_content: DataContentType::Data,
+                data_file_format: DataFileFormat::Parquet,
+                schema: schema.clone(),
+                project_field_ids: vec![1],
+                predicate: Some(predicate.bind(schema, true).unwrap()),
+                deletes: vec![],
+            })]
+            .into_iter(),
+        )) as FileScanTaskStream;
+
+        let result = reader
+            .read(tasks)
+            .await
+            .unwrap()
+            .try_collect::<Vec<RecordBatch>>()
+            .await
+            .unwrap();
+
+        let result_data = result[0].columns()[0]
+            .as_string_opt::<i32>()
+            .unwrap()
+            .iter()
+            .map(|v| v.map(ToOwned::to_owned))
+            .collect::<Vec<_>>();
+
+        result_data
+    }
+
+    fn setup_kleene_logic(
+        data_for_col_a: Vec<Option<String>>,
+    ) -> (FileIO, SchemaRef, String, TempDir) {
+        let schema = Arc::new(
+            Schema::builder()
+                .with_schema_id(1)
+                .with_fields(vec![NestedField::optional(
+                    1,
+                    "a",
+                    Type::Primitive(PrimitiveType::String),
+                )
+                .into()])
+                .build()
+                .unwrap(),
+        );
+
+        let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
+            "a",
+            DataType::Utf8,
+            true,
+        )
+        .with_metadata(HashMap::from([(
+            PARQUET_FIELD_ID_META_KEY.to_string(),
+            "1".to_string(),
+        )]))]));
+
+        let tmp_dir = TempDir::new().unwrap();
+        let table_location = tmp_dir.path().to_str().unwrap().to_string();
+
+        let file_io = FileIO::from_path(&table_location).unwrap().build().unwrap();
+
+        let col = Arc::new(StringArray::from(data_for_col_a)) as ArrayRef;
+
+        let to_write = RecordBatch::try_new(arrow_schema.clone(), vec![col]).unwrap();
+
+        // Write the Parquet files
+        let props = WriterProperties::builder()
+            .set_compression(Compression::SNAPPY)
+            .build();
+
+        let file = File::create(format!("{}/1.parquet", &table_location)).unwrap();
+        let mut writer =
+            ArrowWriter::try_new(file, to_write.schema(), Some(props.clone())).unwrap();
+
+        writer.write(&to_write).expect("Writing batch");
+
+        // writer must be closed to write footer
+        writer.close().unwrap();
+
+        (file_io, schema, table_location, tmp_dir)
+    }
 }

crates/iceberg/src/writer/base_writer/data_file_writer.rs

@@ -144,6 +144,16 @@ mod test {
 
         let mut data_file_writer = DataFileWriterBuilder::new(pw, None).build().await.unwrap();
 
+        let arrow_schema = arrow_schema::Schema::new(vec![
+            Field::new("foo", DataType::Int32, false),
+            Field::new("bar", DataType::Utf8, false),
+        ]);
+        let batch = RecordBatch::try_new(Arc::new(arrow_schema.clone()), vec![
+            Arc::new(Int32Array::from(vec![1, 2, 3])),
+            Arc::new(StringArray::from(vec!["Alice", "Bob", "Charlie"])),
+        ])?;
+        data_file_writer.write(batch).await?;
+
         let data_files = data_file_writer.close().await.unwrap();
         assert_eq!(data_files.len(), 1);
 

crates/iceberg/src/writer/file_writer/parquet_writer.rs

@@ -84,24 +84,16 @@ impl<T: LocationGenerator, F: FileNameGenerator> FileWriterBuilder for ParquetWr
     type R = ParquetWriter;
 
     async fn build(self) -> crate::Result<Self::R> {
-        let arrow_schema: ArrowSchemaRef = Arc::new(self.schema.as_ref().try_into()?);
         let written_size = Arc::new(AtomicI64::new(0));
         let out_file = self.file_io.new_output(
             self.location_generator
                 .generate_location(&self.file_name_generator.generate_file_name()),
         )?;
-        let inner_writer = TrackWriter::new(out_file.writer().await?, written_size.clone());
-        let async_writer = AsyncFileWriter::new(inner_writer);
-        let writer =
-            AsyncArrowWriter::try_new(async_writer, arrow_schema.clone(), Some(self.props))
-                .map_err(|err| {
-                    Error::new(ErrorKind::Unexpected, "Failed to build parquet writer.")
-                        .with_source(err)
-                })?;
 
         Ok(ParquetWriter {
             schema: self.schema.clone(),
-            writer,
+            inner_writer: None,
+            writer_properties: self.props,
             written_size,
             current_row_num: 0,
             out_file,
@@ -226,7 +218,8 @@ impl SchemaVisitor for IndexByParquetPathName {
 pub struct ParquetWriter {
     schema: SchemaRef,
     out_file: OutputFile,
-    writer: AsyncArrowWriter<AsyncFileWriter<TrackWriter>>,
+    inner_writer: Option<AsyncArrowWriter<AsyncFileWriter<TrackWriter>>>,
+    writer_properties: WriterProperties,
     written_size: Arc<AtomicI64>,
     current_row_num: usize,
 }
@@ -520,8 +513,35 @@ impl ParquetWriter {
 
 impl FileWriter for ParquetWriter {
     async fn write(&mut self, batch: &arrow_array::RecordBatch) -> crate::Result<()> {
+        // Skip empty batch
+        if batch.num_rows() == 0 {
+            return Ok(());
+        }
+
         self.current_row_num += batch.num_rows();
-        self.writer.write(batch).await.map_err(|err| {
+
+        // Lazy initialize the writer
+        let writer = if let Some(writer) = &mut self.inner_writer {
+            writer
+        } else {
+            let arrow_schema: ArrowSchemaRef = Arc::new(self.schema.as_ref().try_into()?);
+            let inner_writer =
+                TrackWriter::new(self.out_file.writer().await?, self.written_size.clone());
+            let async_writer = AsyncFileWriter::new(inner_writer);
+            let writer = AsyncArrowWriter::try_new(
+                async_writer,
+                arrow_schema.clone(),
+                Some(self.writer_properties.clone()),
+            )
+            .map_err(|err| {
+                Error::new(ErrorKind::Unexpected, "Failed to build parquet writer.")
+                    .with_source(err)
+            })?;
+            self.inner_writer = Some(writer);
+            self.inner_writer.as_mut().unwrap()
+        };
+
+        writer.write(batch).await.map_err(|err| {
             Error::new(
                 ErrorKind::Unexpected,
                 "Failed to write using parquet writer.",
@@ -532,7 +552,10 @@ impl FileWriter for ParquetWriter {
     }
 
     async fn close(self) -> crate::Result<Vec<crate::spec::DataFileBuilder>> {
-        let metadata = self.writer.close().await.map_err(|err| {
+        let Some(writer) = self.inner_writer else {
+            return Ok(vec![]);
+        };
+        let metadata = writer.close().await.map_err(|err| {
             Error::new(ErrorKind::Unexpected, "Failed to close parquet writer.").with_source(err)
         })?;
 
@@ -1538,4 +1561,57 @@ mod tests {
 
         Ok(())
     }
+
+    #[tokio::test]
+    async fn test_empty_write() -> Result<()> {
+        let temp_dir = TempDir::new().unwrap();
+        let file_io = FileIOBuilder::new_fs_io().build().unwrap();
+        let location_gen =
+            MockLocationGenerator::new(temp_dir.path().to_str().unwrap().to_string());
+        let file_name_gen =
+            DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
+
+        // Test that file will create if data to write
+        let schema = {
+            let fields = vec![
+                arrow_schema::Field::new("col", arrow_schema::DataType::Int64, true).with_metadata(
+                    HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "0".to_string())]),
+                ),
+            ];
+            Arc::new(arrow_schema::Schema::new(fields))
+        };
+        let col = Arc::new(Int64Array::from_iter_values(0..1024)) as ArrayRef;
+        let to_write = RecordBatch::try_new(schema.clone(), vec![col]).unwrap();
+        let mut pw = ParquetWriterBuilder::new(
+            WriterProperties::builder().build(),
+            Arc::new(to_write.schema().as_ref().try_into().unwrap()),
+            file_io.clone(),
+            location_gen.clone(),
+            file_name_gen,
+        )
+        .build()
+        .await?;
+        pw.write(&to_write).await?;
+        let file_path = pw.out_file.location().to_string();
+        pw.close().await.unwrap();
+        assert!(file_io.exists(file_path).await.unwrap());
+
+        // Test that file will not create if no data to write
+        let file_name_gen =
+            DefaultFileNameGenerator::new("test_empty".to_string(), None, DataFileFormat::Parquet);
+        let pw = ParquetWriterBuilder::new(
+            WriterProperties::builder().build(),
+            Arc::new(to_write.schema().as_ref().try_into().unwrap()),
+            file_io.clone(),
+            location_gen,
+            file_name_gen,
+        )
+        .build()
+        .await?;
+        let file_path = pw.out_file.location().to_string();
+        pw.close().await.unwrap();
+        assert!(!file_io.exists(file_path).await.unwrap());
+
+        Ok(())
+    }
 }