hudi-core 0.6.0

/*
 * 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.
 */

use std::collections::{HashMap, HashSet};

use arrow::array::StringArray;

use crate::Result;
use crate::error::CoreError;
use crate::index::{HoodieIndex, HoodieKey, RecordLocation};
use crate::keygen::record_key::record_keys_for_batch;
use crate::metadata::meta_field::MetaField;
use crate::table::{ReadOptions, Table};

/// A single-node index that scans latest file slices to locate record keys.
#[derive(Clone, Debug, Default)]
pub struct SimpleIndex;

impl HoodieIndex for SimpleIndex {
    async fn tag_location(
        &self,
        table: &Table,
        keys: &[HoodieKey],
    ) -> Result<HashMap<HoodieKey, Option<RecordLocation>>> {
        let requested = keys
            .iter()
            .map(|key| key.record_key.clone())
            .collect::<HashSet<_>>();
        let mut locations = keys
            .iter()
            .cloned()
            .map(|key| (key, None))
            .collect::<HashMap<_, _>>();
        if requested.is_empty() {
            return Ok(locations);
        }
        let slices = table.get_file_slices(&ReadOptions::new()).await?;
        if slices.is_empty() {
            return Ok(locations);
        }
        let reader = table
            .create_file_group_reader_with_options(None, std::iter::empty::<(&str, &str)>())
            .await?;
        let mut by_record_key: HashMap<String, Vec<RecordLocation>> = HashMap::new();
        for slice in slices {
            let batch = reader.read_file_slice(&slice, &ReadOptions::new()).await?;
            let location = RecordLocation {
                file_id: slice.file_id().to_string(),
                instant_time: table.timeline.get_latest_commit_timestamp()?,
                partition_path: slice.partition_path.clone(),
            };
            let mut tag = |key: &str| {
                if requested.contains(key) {
                    by_record_key
                        .entry(key.to_string())
                        .or_default()
                        .push(location.clone());
                }
            };
            match batch.column_by_name(MetaField::RecordKey.as_ref()) {
                Some(column) => column
                    .as_any()
                    .downcast_ref::<StringArray>()
                    .ok_or_else(|| {
                        CoreError::Unsupported(format!(
                            "record key field '{}' must be Utf8",
                            MetaField::RecordKey.as_ref()
                        ))
                    })?
                    .iter()
                    .flatten()
                    .for_each(tag),
                None => record_keys_for_batch(&table.hudi_configs, &batch)?
                    .iter()
                    .for_each(|key| tag(key)),
            }
        }
        // A key stored in several partitions (a non-global Java index allows
        // it) resolves to the copy in the incoming record's partition. With no
        // copy there the lookup stays global, as for a single copy: the record
        // moves from the copy with the greatest partition path and the others
        // stay. Ties break on file id, since slice listing order varies
        // between processes.
        for (hoodie_key, slot) in locations.iter_mut() {
            if let Some(candidates) = by_record_key.get(&hoodie_key.record_key) {
                *slot = candidates
                    .iter()
                    .filter(|location| location.partition_path == hoodie_key.partition_path)
                    .max_by_key(|location| &location.file_id)
                    .or_else(|| {
                        candidates
                            .iter()
                            .max_by_key(|location| (&location.partition_path, &location.file_id))
                    })
                    .cloned();
            }
        }
        Ok(locations)
    }
}