feat(geo-example): add up to write to table

Author: nuts-rice

crates/examples/src/geo_rest_catalog.rs

@@ -18,23 +18,23 @@
 use std::collections::HashMap;
 use std::sync::Arc;
 
-use arrow_array::{ArrayRef, Float64Array, Int32Array, LargeBinaryArray, StringArray};
+use arrow_array::{ArrayRef, Float64Array, Int32Array, LargeBinaryArray, RecordBatch, StringArray};
 use geo_types::{Geometry, Point};
 use iceberg::spec::{NestedField, PrimitiveType, Schema, Type};
-use iceberg::writer::IcebergWriterBuilder;
 use iceberg::writer::base_writer::data_file_writer::DataFileWriterBuilder;
 use iceberg::writer::file_writer::ParquetWriterBuilder;
 use iceberg::writer::file_writer::location_generator::{
     DefaultFileNameGenerator, DefaultLocationGenerator,
 };
 use iceberg::writer::file_writer::rolling_writer::RollingFileWriterBuilder;
+use iceberg::writer::{IcebergWriter, IcebergWriterBuilder};
 use iceberg::{Catalog, CatalogBuilder, NamespaceIdent, TableCreation, TableIdent};
 use iceberg_catalog_rest::{REST_CATALOG_PROP_URI, RestCatalogBuilder};
 use parquet::file::properties::WriterProperties;
 
 static REST_URI: &str = "http://localhost:8181";
 static NAMESPACE: &str = "ns1";
-static TABLE_NAME: &str = "cities_table2";
+static TABLE_NAME: &str = "cities_table3";
 
 //This is an example of creating and loading an table using a schema with
 // geo types via the Iceberg REST Catalog.
@@ -143,41 +143,28 @@ async fn main() {
             HashMap::from([(REST_CATALOG_PROP_URI.to_string(), REST_URI.to_string())]),
         )
         .await
-        .map_err(|e| {
-            eprintln!("Failed to connect to REST catalog: {:?}", e);
-            eprintln!("Error: {}", e);
-            e
-        })
         .unwrap();
     println!("Connected to REST Catalog at {}", REST_URI);
 
     let namespace_ident = NamespaceIdent::from_vec(vec![NAMESPACE.to_string()]).unwrap();
-    let table_ident = TableIdent::new(namespace_ident.clone(), TABLE_NAME.to_string());
-
-    println!("Checking if table exists...");
-    let table_exists = catalog
-        .table_exists(&table_ident)
-        .await
-        .map_err(|e| {
-            eprintln!("Failed to check if table exists: {:?}", e);
-            eprintln!("Error: {}", e);
-            e
-        })
-        .unwrap();
 
-    if table_exists {
-        println!("Table {TABLE_NAME} already exists, dropping now.");
+    // Create namespace if it doesn't exist
+    if !catalog.namespace_exists(&namespace_ident).await.unwrap() {
+        println!("Creating namespace {NAMESPACE}...");
         catalog
-            .drop_table(&table_ident)
+            .create_namespace(&namespace_ident, HashMap::new())
             .await
-            .map_err(|e| {
-                eprintln!("Failed to drop table: {:?}", e);
-                eprintln!("Error: {}", e);
-                e
-            })
             .unwrap();
     }
 
+    let table_ident = TableIdent::new(namespace_ident.clone(), TABLE_NAME.to_string());
+
+    println!("Checking if table exists...");
+    if catalog.table_exists(&table_ident).await.unwrap() {
+        println!("Table {TABLE_NAME} already exists, dropping now.");
+        catalog.drop_table(&table_ident).await.unwrap();
+    }
+
     let iceberg_schema = Schema::builder()
         .with_fields(vec![
             NestedField::required(1, "id".to_string(), Type::Primitive(PrimitiveType::Int)).into(),
@@ -251,15 +238,6 @@ async fn main() {
     let _created_table = catalog
         .create_table(&table_ident.namespace, table_creation)
         .await
-        .map_err(|e| {
-            eprintln!("\n=== FAILED TO CREATE TABLE ===");
-            eprintln!("Error type: {:?}", e);
-            eprintln!("Error message: {}", e);
-            eprintln!("Namespace: {:?}", table_ident.namespace);
-            eprintln!("Table name: {}", table_ident.name);
-            eprintln!("==============================\n");
-            e
-        })
         .unwrap();
     println!("Table {TABLE_NAME} created.");
     assert!(
@@ -269,6 +247,14 @@ async fn main() {
             .unwrap()
             .contains(&table_ident)
     );
+    let schema: Arc<arrow_schema::Schema> = Arc::new(
+        _created_table
+            .metadata()
+            .current_schema()
+            .as_ref()
+            .try_into()
+            .unwrap(),
+    );
     let location_generator =
         DefaultLocationGenerator::new(_created_table.metadata().clone()).unwrap();
     let file_name_generator = DefaultFileNameGenerator::new(
@@ -287,66 +273,65 @@ async fn main() {
         file_name_generator.clone(),
     );
     let data_file_writer_builder = DataFileWriterBuilder::new(rolling_file_writer_builder);
-    //let data_file_writer = data_file_writer_builder.build(None).await.unwrap();
+    let mut data_file_writer = data_file_writer_builder.build(None).await.unwrap();
 
     let features = mock_sample_features();
-    /*
-        let ids: ArrayRef = Arc::new(Int32Array::from_iter_values(features.iter().map(|f| f.id)));
-        let names: ArrayRef = Arc::new(StringArray::from_iter_values(
-            features.iter().map(|f| f.name.as_str()),
-        ));
-        let geometries_wkb: ArrayRef = Arc::new(LargeBinaryArray::from_iter_values(
-            features.iter().map(|f| f.to_wkb()),
-        ));
-        let geometry_types: ArrayRef = Arc::new(StringArray::from_iter_values(
-            features.iter().map(|f| f.geometry_type()),
-        ));
-        let srids: ArrayRef = Arc::new(Int32Array::from_iter_values(
-            features.iter().map(|f| f.srid),
-        ));
-        let bbox_min_xs: ArrayRef = Arc::new(Float64Array::from_iter_values(
-            features.iter().map(|f| f.bbox().0),
-        ));
-        let bbox_min_ys: ArrayRef = Arc::new(Float64Array::from_iter_values(
-            features.iter().map(|f| f.bbox().1),
-        ));
-        let bbox_max_xs: ArrayRef = Arc::new(Float64Array::from_iter_values(
-            features.iter().map(|f| f.bbox().2),
-        ));
-        let bbox_max_ys: ArrayRef = Arc::new(Float64Array::from_iter_values(
-            features.iter().map(|f| f.bbox().3),
-        ));
 
-        let countries: ArrayRef = Arc::new(StringArray::from_iter_values(
-            features
-                .iter()
-                .map(|f| f.properties.get("country").unwrap().as_str()),
-        ));
-        let populations: ArrayRef = Arc::new(StringArray::from_iter_values(
-            features
-                .iter()
-                .map(|f| f.properties.get("population").unwrap().as_str()),
-        ));
-    */
+    let ids: ArrayRef = Arc::new(Int32Array::from_iter_values(features.iter().map(|f| f.id)));
+    let names: ArrayRef = Arc::new(StringArray::from_iter_values(
+        features.iter().map(|f| f.name.as_str()),
+    ));
+    let geometries_wkb: ArrayRef = Arc::new(LargeBinaryArray::from_iter_values(
+        features.iter().map(|f| f.to_wkb()),
+    ));
+    let geometry_types: ArrayRef = Arc::new(StringArray::from_iter_values(
+        features.iter().map(|f| f.geometry_type()),
+    ));
+    let srids: ArrayRef = Arc::new(Int32Array::from_iter_values(
+        features.iter().map(|f| f.srid),
+    ));
+    let bbox_min_xs: ArrayRef = Arc::new(Float64Array::from_iter_values(
+        features.iter().map(|f| f.bbox().0),
+    ));
+    let bbox_min_ys: ArrayRef = Arc::new(Float64Array::from_iter_values(
+        features.iter().map(|f| f.bbox().1),
+    ));
+    let bbox_max_xs: ArrayRef = Arc::new(Float64Array::from_iter_values(
+        features.iter().map(|f| f.bbox().2),
+    ));
+    let bbox_max_ys: ArrayRef = Arc::new(Float64Array::from_iter_values(
+        features.iter().map(|f| f.bbox().3),
+    ));
+
+    let countries: ArrayRef = Arc::new(StringArray::from_iter_values(
+        features
+            .iter()
+            .map(|f| f.properties.get("country").unwrap().as_str()),
+    ));
+    let populations: ArrayRef = Arc::new(StringArray::from_iter_values(
+        features
+            .iter()
+            .map(|f| f.properties.get("population").unwrap().as_str()),
+    ));
+
     //TODO: make write with credentials
-    /*let record_batch = RecordBatch::try_new(schema.clone(), vec![
-            ids,
-            names,
-            geometries_wkb,
-            geometry_types,
-            srids,
-            bbox_min_xs,
-            bbox_min_ys,
-            bbox_max_xs,
-            bbox_max_ys,
-            countries,
-            populations,
-        ])
-        .unwrap();
+    let record_batch = RecordBatch::try_new(schema.clone(), vec![
+        ids,
+        names,
+        geometries_wkb,
+        geometry_types,
+        srids,
+        bbox_min_xs,
+        bbox_min_ys,
+        bbox_max_xs,
+        bbox_max_ys,
+        countries,
+        populations,
+    ])
+    .unwrap();
 
-        data_file_writer.write(record_batch.clone()).await.unwrap();
-        let data_file = data_file_writer.close().await.unwrap();
-    */
+    data_file_writer.write(record_batch.clone()).await.unwrap();
+    let data_file = data_file_writer.close().await.unwrap();
 
     let loaded_table = catalog.load_table(&table_ident).await.unwrap();
     println!("Table {TABLE_NAME} loaded!\n\nTable: {loaded_table:?}");