Support copying from glob patterns

aykut-bozkurt · aykut-bozkurt · commit fa24da5e2c01 · 2025-03-19T04:53:31.000+03:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -27,6 +27,7 @@ aws-config = { version = "1", default-features = false, features = ["rustls","rt
 aws-credential-types = {version = "1", default-features = false}
 azure_storage = {version = "0.21", default-features = false}
 futures = "0.3"
+glob = "0.3"
 home = "0.5"
 object_store = {version = "0.11", default-features = false, features = ["aws", "azure", "http"]}
 once_cell = "1"
diff --git a/src/arrow_parquet/parquet_reader.rs b/src/arrow_parquet/parquet_reader.rs
@@ -1,7 +1,11 @@
-use std::sync::Arc;
+use std::{
+    ops::{Deref, DerefMut},
+    sync::Arc,
+};
 
 use arrow::array::RecordBatch;
 use arrow_cast::{cast_with_options, CastOptions};
+use arrow_schema::SchemaRef;
 use futures::StreamExt;
 use parquet::arrow::async_reader::{ParquetObjectReader, ParquetRecordBatchStream};
 use pgrx::{
@@ -22,6 +26,7 @@ use crate::{
             parquet_schema_string_from_attributes,
         },
     },
+    parquet_udfs::list::list_uri,
     pgrx_utils::{collect_attributes_for, CollectAttributesFor},
     type_compat::{geometry::reset_postgis_context, map::reset_map_context},
     PG_BACKEND_TOKIO_RUNTIME,
@@ -36,18 +41,113 @@ use super::{
     uri_utils::{parquet_reader_from_uri, ParsedUriInfo},
 };
 
+pub(crate) struct SingleParquetReader {
+    reader: ParquetRecordBatchStream<ParquetObjectReader>,
+    attribute_contexts: Vec<ArrowToPgAttributeContext>,
+    match_by: MatchBy,
+}
+
+impl SingleParquetReader {
+    fn new(
+        uri_info: ParsedUriInfo,
+        match_by: MatchBy,
+        tupledesc_schema: SchemaRef,
+        attributes: &[FormData_pg_attribute],
+    ) -> Self {
+        let reader = parquet_reader_from_uri(uri_info);
+
+        // Ensure that the file schema matches the tupledesc schema.
+        // Gets cast_to_types for each attribute if a cast is needed for the attribute's columnar array
+        // to match the expected columnar array for its tupledesc type.
+        let cast_to_types = ensure_file_schema_match_tupledesc_schema(
+            reader.schema().clone(),
+            tupledesc_schema.clone(),
+            attributes,
+            match_by,
+        );
+
+        let attribute_contexts = collect_arrow_to_pg_attribute_contexts(
+            attributes,
+            &tupledesc_schema.fields,
+            Some(cast_to_types),
+        );
+
+        SingleParquetReader {
+            reader,
+            attribute_contexts,
+            match_by,
+        }
+    }
+
+    fn attribute_count(&self) -> usize {
+        self.attribute_contexts.len()
+    }
+
+    fn record_batch_to_tuple_datums(&self, record_batch: RecordBatch) -> Vec<Option<Datum>> {
+        let mut datums = vec![];
+
+        for (attribute_idx, attribute_context) in self.attribute_contexts.iter().enumerate() {
+            let name = attribute_context.name();
+
+            let column_array = match self.match_by {
+                MatchBy::Position => record_batch
+                    .columns()
+                    .get(attribute_idx)
+                    .unwrap_or_else(|| panic!("column {} not found", name)),
+
+                MatchBy::Name => record_batch
+                    .column_by_name(name)
+                    .unwrap_or_else(|| panic!("column {} not found", name)),
+            };
+
+            let datum = if attribute_context.needs_cast() {
+                // should fail instead of returning None if the cast fails at runtime
+                let cast_options = CastOptions {
+                    safe: false,
+                    ..Default::default()
+                };
+
+                let casted_column_array =
+                    cast_with_options(&column_array, attribute_context.data_type(), &cast_options)
+                        .unwrap_or_else(|e| panic!("failed to cast column {}: {}", name, e));
+
+                to_pg_datum(casted_column_array.to_data(), attribute_context)
+            } else {
+                to_pg_datum(column_array.to_data(), attribute_context)
+            };
+
+            datums.push(datum);
+        }
+
+        datums
+    }
+}
+
 pub(crate) struct ParquetReaderContext {
     buffer: Vec<u8>,
     offset: usize,
     started: bool,
     finished: bool,
-    parquet_reader: ParquetRecordBatchStream<ParquetObjectReader>,
-    attribute_contexts: Vec<ArrowToPgAttributeContext>,
+    parquet_readers: Vec<SingleParquetReader>,
+    current_parquet_reader_idx: usize,
     binary_out_funcs: Vec<PgBox<FmgrInfo>>,
-    match_by: MatchBy,
     per_row_memory_ctx: PgMemoryContexts,
 }
 
+impl Deref for ParquetReaderContext {
+    type Target = SingleParquetReader;
+
+    fn deref(&self) -> &Self::Target {
+        &self.parquet_readers[self.current_parquet_reader_idx]
+    }
+}
+
+impl DerefMut for ParquetReaderContext {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.parquet_readers[self.current_parquet_reader_idx]
+    }
+}
+
 impl ParquetReaderContext {
     pub(crate) fn new(uri_info: ParsedUriInfo, match_by: MatchBy, tupledesc: &PgTupleDesc) -> Self {
         // Postgis and Map contexts are used throughout reading the parquet file.
@@ -57,10 +157,6 @@ impl ParquetReaderContext {
 
         error_if_copy_from_match_by_position_with_generated_columns(tupledesc, match_by);
 
-        let parquet_reader = parquet_reader_from_uri(uri_info);
-
-        let parquet_file_schema = parquet_reader.schema();
-
         let attributes = collect_attributes_for(CollectAttributesFor::CopyFrom, tupledesc);
 
         pgrx::debug2!(
@@ -72,21 +168,29 @@ impl ParquetReaderContext {
 
         let tupledesc_schema = Arc::new(tupledesc_schema);
 
-        // Ensure that the file schema matches the tupledesc schema.
-        // Gets cast_to_types for each attribute if a cast is needed for the attribute's columnar array
-        // to match the expected columnar array for its tupledesc type.
-        let cast_to_types = ensure_file_schema_match_tupledesc_schema(
-            parquet_file_schema.clone(),
-            tupledesc_schema.clone(),
-            &attributes,
-            match_by,
-        );
-
-        let attribute_contexts = collect_arrow_to_pg_attribute_contexts(
-            &attributes,
-            &tupledesc_schema.fields,
-            Some(cast_to_types),
-        );
+        let parquet_readers = if uri_info.is_pattern() {
+            list_uri(uri_info)
+                .into_iter()
+                .map(|(file_uri, _)| {
+                    let file_uri_info = ParsedUriInfo::try_from(file_uri.as_str())
+                        .unwrap_or_else(|e| panic!("failed to parse file uri {}: {}", file_uri, e));
+
+                    SingleParquetReader::new(
+                        file_uri_info,
+                        match_by,
+                        tupledesc_schema.clone(),
+                        &attributes,
+                    )
+                })
+                .collect()
+        } else {
+            vec![SingleParquetReader::new(
+                uri_info,
+                match_by,
+                tupledesc_schema.clone(),
+                &attributes,
+            )]
+        };
 
         let binary_out_funcs = Self::collect_binary_out_funcs(&attributes);
 
@@ -95,10 +199,9 @@ impl ParquetReaderContext {
         ParquetReaderContext {
             buffer: Vec::new(),
             offset: 0,
-            attribute_contexts,
-            parquet_reader,
+            parquet_readers,
+            current_parquet_reader_idx: 0,
             binary_out_funcs,
-            match_by,
             started: false,
             finished: false,
             per_row_memory_ctx,
@@ -128,47 +231,8 @@ impl ParquetReaderContext {
         }
     }
 
-    fn record_batch_to_tuple_datums(
-        record_batch: RecordBatch,
-        attribute_contexts: &[ArrowToPgAttributeContext],
-        match_by: MatchBy,
-    ) -> Vec<Option<Datum>> {
-        let mut datums = vec![];
-
-        for (attribute_idx, attribute_context) in attribute_contexts.iter().enumerate() {
-            let name = attribute_context.name();
-
-            let column_array = match match_by {
-                MatchBy::Position => record_batch
-                    .columns()
-                    .get(attribute_idx)
-                    .unwrap_or_else(|| panic!("column {} not found", name)),
-
-                MatchBy::Name => record_batch
-                    .column_by_name(name)
-                    .unwrap_or_else(|| panic!("column {} not found", name)),
-            };
-
-            let datum = if attribute_context.needs_cast() {
-                // should fail instead of returning None if the cast fails at runtime
-                let cast_options = CastOptions {
-                    safe: false,
-                    ..Default::default()
-                };
-
-                let casted_column_array =
-                    cast_with_options(&column_array, attribute_context.data_type(), &cast_options)
-                        .unwrap_or_else(|e| panic!("failed to cast column {}: {}", name, e));
-
-                to_pg_datum(casted_column_array.to_data(), attribute_context)
-            } else {
-                to_pg_datum(column_array.to_data(), attribute_context)
-            };
-
-            datums.push(datum);
-        }
-
-        datums
+    fn has_more_parquet_readers(&self) -> bool {
+        self.current_parquet_reader_idx < self.parquet_readers.len() - 1
     }
 
     pub(crate) fn read_parquet(&mut self) -> bool {
@@ -183,7 +247,7 @@ impl ParquetReaderContext {
 
         // read a record batch from the parquet file. Record batch will contain
         // DEFAULT_BATCH_SIZE rows as we configured in the parquet reader.
-        let record_batch = PG_BACKEND_TOKIO_RUNTIME.block_on(self.parquet_reader.next());
+        let record_batch = PG_BACKEND_TOKIO_RUNTIME.block_on(self.reader.next());
 
         if let Some(batch_result) = record_batch {
             let record_batch =
@@ -198,6 +262,10 @@ impl ParquetReaderContext {
                 let record_batch = record_batch.slice(i, 1);
                 self.copy_row(record_batch);
             }
+        } else if self.has_more_parquet_readers() {
+            // move to the next parquet reader
+            self.current_parquet_reader_idx += 1;
+            self.read_parquet();
         } else {
             // finish PG copy
             self.copy_finish();
@@ -208,43 +276,39 @@ impl ParquetReaderContext {
 
     fn copy_row(&mut self, record_batch: RecordBatch) {
         unsafe {
-            self.per_row_memory_ctx.switch_to(|_context| {
-                /* 2 bytes: per-tuple header */
-                let natts = self.attribute_contexts.len() as i16;
-                let attnum_len_bytes = natts.to_be_bytes();
-                self.buffer.extend_from_slice(&attnum_len_bytes);
-
-                // convert the columnar arrays in record batch to tuple datums
-                let tuple_datums = Self::record_batch_to_tuple_datums(
-                    record_batch,
-                    &self.attribute_contexts,
-                    self.match_by,
-                );
-
-                // write the tuple datums to the ParquetReader's internal buffer in PG copy format
-                for (datum, out_func) in tuple_datums.into_iter().zip(self.binary_out_funcs.iter())
-                {
-                    if let Some(datum) = datum {
-                        let datum_bytes: *mut varlena = SendFunctionCall(out_func.as_ptr(), datum);
-
-                        /* 4 bytes: attribute's data size */
-                        let data_size = varsize_any_exhdr(datum_bytes);
-                        let data_size_bytes = (data_size as i32).to_be_bytes();
-                        self.buffer.extend_from_slice(&data_size_bytes);
-
-                        /* variable bytes: attribute's data */
-                        let data = vardata_any(datum_bytes) as _;
-                        let data_bytes = std::slice::from_raw_parts(data, data_size);
-                        self.buffer.extend_from_slice(data_bytes);
-                    } else {
-                        /* 4 bytes: null */
-                        let null_value = -1_i32;
-                        let null_value_bytes = null_value.to_be_bytes();
-                        self.buffer.extend_from_slice(&null_value_bytes);
-                    }
+            let mut old_ctx = self.per_row_memory_ctx.set_as_current();
+
+            /* 2 bytes: per-tuple header */
+            let natts = self.attribute_count() as i16;
+            let attnum_len_bytes = natts.to_be_bytes();
+            self.buffer.extend_from_slice(&attnum_len_bytes);
+
+            // convert the columnar arrays in record batch to tuple datums
+            let tuple_datums = self.record_batch_to_tuple_datums(record_batch);
+
+            // write the tuple datums to the ParquetReader's internal buffer in PG copy format
+            for (datum, out_func) in tuple_datums.into_iter().zip(self.binary_out_funcs.iter()) {
+                if let Some(datum) = datum {
+                    let datum_bytes: *mut varlena = SendFunctionCall(out_func.as_ptr(), datum);
+
+                    /* 4 bytes: attribute's data size */
+                    let data_size = varsize_any_exhdr(datum_bytes);
+                    let data_size_bytes = (data_size as i32).to_be_bytes();
+                    self.buffer.extend_from_slice(&data_size_bytes);
+
+                    /* variable bytes: attribute's data */
+                    let data = vardata_any(datum_bytes) as _;
+                    let data_bytes = std::slice::from_raw_parts(data, data_size);
+                    self.buffer.extend_from_slice(data_bytes);
+                } else {
+                    /* 4 bytes: null */
+                    let null_value = -1_i32;
+                    let null_value_bytes = null_value.to_be_bytes();
+                    self.buffer.extend_from_slice(&null_value_bytes);
                 }
-            });
+            }
 
+            old_ctx.set_as_current();
             self.per_row_memory_ctx.reset();
         };
     }
diff --git a/src/arrow_parquet/uri_utils.rs b/src/arrow_parquet/uri_utils.rs
@@ -73,6 +73,23 @@ impl ParsedUriInfo {
                             uri.scheme(), uri))
         }
     }
+
+    pub(crate) fn base_uri(&self) -> String {
+        if self.uri.scheme() == "file" {
+            // root path for local file
+            return "/".to_string();
+        }
+
+        format!(
+            "{}://{}",
+            self.uri.scheme(),
+            self.uri.host().expect("missing host")
+        )
+    }
+
+    pub(crate) fn is_pattern(&self) -> bool {
+        self.path.to_string().contains('*') || self.path.to_string().contains("**")
+    }
 }
 
 impl TryFrom<&str> for ParsedUriInfo {
diff --git a/src/parquet_udfs.rs b/src/parquet_udfs.rs
@@ -1,2 +1,3 @@
+pub(crate) mod list;
 pub(crate) mod metadata;
 pub(crate) mod schema;
diff --git a/src/parquet_udfs/list.rs b/src/parquet_udfs/list.rs

Original file line number	Diff line number	Diff line change
`@@ -1,2 +1,3 @@`
	`1`	`+pub(crate) mod list;`
`1`	`2`	`pub(crate) mod metadata;`
`2`	`3`	`pub(crate) mod schema;`