alopex-core 0.8.0

//! KVS ブリッジ: カラムナーセグメントをトランザクション経由で読み書きする。

use std::sync::Arc;
use std::time::Duration;

use bincode::Options;

use crate::columnar::error::{ColumnarError, Result};
use crate::columnar::segment_v08::{
    ChunkedSegmentAccessV08, RangeAddressableSegmentProvider, SegmentReferenceV08,
    DEFAULT_V08_METADATA_LIMIT_BYTES,
};
use crate::columnar::segment_v2::{
    ColumnSegmentV2, InMemorySegmentSource, RecordBatch, SegmentMetaV2, SegmentReaderV2,
};
use crate::kv::any::{AnyKV, AnyKVTransaction};
use crate::kv::{KVStore, KVTransaction};
use crate::storage::format::bincode_config;
use crate::txn::TxnManager;
use crate::types::TxnMode;

/// キーレイアウトおよびプレフィックス。
pub mod key_layout {
    /// テーブルメタデータ（予約）。
    pub const PREFIX_TABLE_META: u8 = 0x10;
    /// カラムナーセグメント本体。
    pub const PREFIX_COLUMN_SEGMENT: u8 = 0x11;
    /// セグメントインデックス。
    pub const PREFIX_SEGMENT_INDEX: u8 = 0x12;
    /// 統計情報。
    pub const PREFIX_STATISTICS: u8 = 0x13;
    /// RowGroup 単位の付加情報。
    pub const PREFIX_ROW_GROUP: u8 = 0x14;
    /// Provisioned V08 range-addressable columnar segment objects.
    pub const PREFIX_COLUMN_SEGMENT_V08: u8 = 0x15;

    const V08_OBJECT_HEADER: u8 = 0;
    const V08_OBJECT_SCHEMA: u8 = 1;
    const V08_OBJECT_DIRECTORY: u8 = 2;
    const V08_OBJECT_CHUNK: u8 = 3;

    /// カラムナーセグメントのキーを生成する。
    pub fn column_segment_key(table_id: u32, segment_id: u64, column_idx: u16) -> Vec<u8> {
        let mut key = Vec::with_capacity(1 + 4 + 8 + 2);
        key.push(PREFIX_COLUMN_SEGMENT);
        key.extend_from_slice(&table_id.to_le_bytes());
        key.extend_from_slice(&segment_id.to_le_bytes());
        key.extend_from_slice(&column_idx.to_le_bytes());
        key
    }

    /// テーブル単位のセグメントインデックスキーを生成する。
    pub fn segment_index_key(table_id: u32) -> Vec<u8> {
        let mut key = Vec::with_capacity(1 + 4);
        key.push(PREFIX_SEGMENT_INDEX);
        key.extend_from_slice(&table_id.to_le_bytes());
        key
    }

    /// 統計情報キーを生成する。
    pub fn statistics_key(table_id: u32, segment_id: u64) -> Vec<u8> {
        let mut key = Vec::with_capacity(1 + 4 + 8);
        key.push(PREFIX_STATISTICS);
        key.extend_from_slice(&table_id.to_le_bytes());
        key.extend_from_slice(&segment_id.to_le_bytes());
        key
    }

    /// RowGroup 統計情報キーを生成する。
    pub fn row_group_stats_key(table_id: u32, segment_id: u64) -> Vec<u8> {
        let mut key = Vec::with_capacity(1 + 4 + 8);
        key.push(PREFIX_ROW_GROUP);
        key.extend_from_slice(&table_id.to_le_bytes());
        key.extend_from_slice(&segment_id.to_le_bytes());
        key
    }

    /// Generates the independent V08 header key.
    pub fn v08_header_key(table_id: u32, segment_id: u64) -> Vec<u8> {
        v08_object_key(table_id, segment_id, V08_OBJECT_HEADER)
    }

    /// Generates the independent V08 schema key.
    pub fn v08_schema_key(table_id: u32, segment_id: u64) -> Vec<u8> {
        v08_object_key(table_id, segment_id, V08_OBJECT_SCHEMA)
    }

    /// Generates the independent V08 directory key.
    pub fn v08_directory_key(table_id: u32, segment_id: u64) -> Vec<u8> {
        v08_object_key(table_id, segment_id, V08_OBJECT_DIRECTORY)
    }

    /// Generates one independently addressable V08 column chunk key.
    pub fn v08_chunk_key(
        table_id: u32,
        segment_id: u64,
        row_group_index: u32,
        column_index: u16,
    ) -> Vec<u8> {
        let mut key = v08_object_key(table_id, segment_id, V08_OBJECT_CHUNK);
        key.extend_from_slice(&row_group_index.to_le_bytes());
        key.extend_from_slice(&column_index.to_le_bytes());
        key
    }

    fn v08_object_key(table_id: u32, segment_id: u64, object: u8) -> Vec<u8> {
        let mut key = Vec::with_capacity(1 + 4 + 8 + 1 + 4 + 2);
        key.push(PREFIX_COLUMN_SEGMENT_V08);
        key.extend_from_slice(&table_id.to_le_bytes());
        key.extend_from_slice(&segment_id.to_le_bytes());
        key.push(object);
        key
    }
}

/// Preflight result for DataFrame streaming over an embedded segment.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum StreamingSegmentLayoutPreflightV08 {
    /// A separately-addressable V08 header is provisioned.
    Available,
    /// The legacy V2 blob exists but cannot be used for bounded streaming.
    RequiresV08ChunkedLayout,
    /// Neither V08 nor legacy V2 data exists for this segment reference.
    NotFound,
}

/// カラムナーセグメントを KVS に永続化するブリッジ。
#[derive(Clone)]
pub struct ColumnarKvsBridge {
    store: Arc<AnyKV>,
    max_retries: usize,
}

impl ColumnarKvsBridge {
    /// 新規ブリッジを生成する。
    pub fn new(store: Arc<AnyKV>) -> Self {
        Self {
            store,
            max_retries: 10,
        }
    }

    fn read_v08_object(&self, key: Vec<u8>) -> Result<Vec<u8>> {
        let mut txn = self.store.begin(TxnMode::ReadOnly)?;
        txn.get(&key)?.ok_or(ColumnarError::NotFound)
    }

    fn load_index<'a>(txn: &mut AnyKVTransaction<'a>, table_id: u32) -> Result<Vec<u64>> {
        let key = key_layout::segment_index_key(table_id);
        if let Some(bytes) = txn.get(&key)? {
            bincode_config()
                .deserialize(&bytes)
                .map_err(|e| ColumnarError::InvalidFormat(e.to_string()))
        } else {
            Ok(Vec::new())
        }
    }

    fn persist_index<'a>(
        txn: &mut AnyKVTransaction<'a>,
        table_id: u32,
        index: &[u64],
    ) -> Result<()> {
        let key = key_layout::segment_index_key(table_id);
        let bytes = bincode_config()
            .serialize(index)
            .map_err(|e| ColumnarError::InvalidFormat(e.to_string()))?;
        txn.put(key, bytes)?;
        Ok(())
    }

    /// セグメントを書き込み、割り当てたセグメントIDを返す。
    ///
    /// 競合が起きた場合は簡易リトライを行う。
    pub fn write_segment(&self, table_id: u32, segment: &ColumnSegmentV2) -> Result<u64> {
        let mut attempts = 0usize;
        loop {
            attempts += 1;
            let store = self.store.clone();
            let manager = store.txn_manager();
            let mut txn = manager.begin(TxnMode::ReadWrite)?;
            let mut index = Self::load_index(&mut txn, table_id)?;
            let segment_id = index.last().copied().unwrap_or(0);
            let next_id = if index.is_empty() {
                0
            } else {
                segment_id.saturating_add(1)
            };

            let segment_key = key_layout::column_segment_key(table_id, next_id, 0);
            let bytes = bincode_config()
                .serialize(segment)
                .map_err(|e| ColumnarError::InvalidFormat(e.to_string()))?;
            txn.put(segment_key, bytes)?;

            // 統計は SegmentMetaV2 をそのまま保持。
            let stats_key = key_layout::statistics_key(table_id, next_id);
            let stats_bytes = bincode_config()
                .serialize(&segment.meta)
                .map_err(|e| ColumnarError::InvalidFormat(e.to_string()))?;
            txn.put(stats_key, stats_bytes)?;

            index.push(next_id);
            Self::persist_index(&mut txn, table_id, &index)?;

            let commit_result = manager.commit(txn).map_err(ColumnarError::from);

            match commit_result {
                Ok(()) => return Ok(next_id),
                Err(ColumnarError::TxnConflict) if attempts < self.max_retries => {
                    let backoff_ms = 10u64.saturating_mul(attempts as u64);
                    std::thread::sleep(Duration::from_millis(backoff_ms));
                    continue;
                }
                Err(e) => return Err(e),
            }
        }
    }

    /// セグメントを読み取り、指定カラムのみを返す。
    pub fn read_segment(
        &self,
        table_id: u32,
        segment_id: u64,
        columns: &[usize],
    ) -> Result<Vec<RecordBatch>> {
        let reader = self.open_segment_reader(table_id, segment_id)?;
        reader.read_columns(columns)
    }

    /// セグメント本体を取得する。
    pub fn read_segment_raw(&self, table_id: u32, segment_id: u64) -> Result<ColumnSegmentV2> {
        let key = key_layout::column_segment_key(table_id, segment_id, 0);
        let mut txn = self.store.begin(TxnMode::ReadOnly)?;
        let bytes = txn.get(&key)?.ok_or(ColumnarError::NotFound)?;
        bincode_config()
            .deserialize(&bytes)
            .map_err(|e| ColumnarError::InvalidFormat(e.to_string()))
    }

    /// Classifies whether this segment can be opened without a V2 blob read.
    pub fn streaming_layout_preflight_v08(
        &self,
        table_id: u32,
        segment_id: u64,
    ) -> Result<StreamingSegmentLayoutPreflightV08> {
        let mut txn = self.store.begin(TxnMode::ReadOnly)?;
        if txn
            .get(&key_layout::v08_header_key(table_id, segment_id))?
            .is_some()
        {
            return Ok(StreamingSegmentLayoutPreflightV08::Available);
        }
        // V2 persists a separate small statistics object.  Looking up this
        // marker avoids fetching its monolithic segment blob merely to decide
        // that bounded streaming must reject it.
        if txn
            .get(&key_layout::statistics_key(table_id, segment_id))?
            .is_some()
        {
            return Ok(StreamingSegmentLayoutPreflightV08::RequiresV08ChunkedLayout);
        }
        Ok(StreamingSegmentLayoutPreflightV08::NotFound)
    }

    /// Opens a provisioned V08 read-only chunk provider.
    ///
    /// Legacy V2 data is classified before it is read and is never converted
    /// into artificial streaming batches.
    pub fn open_v08_segment(
        &self,
        table_id: u32,
        segment_id: u64,
    ) -> Result<ChunkedSegmentAccessV08> {
        self.open_v08_segment_with_metadata_limit(
            table_id,
            segment_id,
            DEFAULT_V08_METADATA_LIMIT_BYTES,
        )
    }

    /// Opens a provisioned V08 reader with the caller's metadata bound.
    pub fn open_v08_segment_with_metadata_limit(
        &self,
        table_id: u32,
        segment_id: u64,
        metadata_limit_bytes: u64,
    ) -> Result<ChunkedSegmentAccessV08> {
        match self.streaming_layout_preflight_v08(table_id, segment_id)? {
            StreamingSegmentLayoutPreflightV08::Available => {
                ChunkedSegmentAccessV08::open_with_metadata_limit(
                    Arc::new(self.clone()),
                    SegmentReferenceV08::new(table_id, segment_id),
                    metadata_limit_bytes,
                )
            }
            StreamingSegmentLayoutPreflightV08::RequiresV08ChunkedLayout => {
                Err(ColumnarError::RequiresV08ChunkedLayout)
            }
            StreamingSegmentLayoutPreflightV08::NotFound => Err(ColumnarError::NotFound),
        }
    }

    /// セグメントリーダーを開く。
    pub fn open_segment_reader(&self, table_id: u32, segment_id: u64) -> Result<SegmentReaderV2> {
        let segment = self.read_segment_raw(table_id, segment_id)?;
        SegmentReaderV2::open(Box::new(InMemorySegmentSource::new(segment.data)))
    }

    /// セグメントメタデータのみ取得する（統計用）。
    pub fn read_statistics(&self, table_id: u32, segment_id: u64) -> Result<Vec<u8>> {
        let key = key_layout::statistics_key(table_id, segment_id);
        let mut txn = self.store.begin(TxnMode::ReadOnly)?;
        let bytes = txn.get(&key)?.ok_or(ColumnarError::NotFound)?;
        Ok(bytes)
    }

    /// インデックスを取得する（テスト用ユーティリティ）。
    pub fn segment_index(&self, table_id: u32) -> Result<Vec<u64>> {
        let mut txn = self.store.begin(TxnMode::ReadOnly)?;
        let index = Self::load_index(&mut txn, table_id)?;
        Ok(index)
    }

    /// カラム数を取得する（統計メタから取得）。
    pub fn column_count(&self, table_id: u32, segment_id: u64) -> Result<usize> {
        let stats = self.read_statistics(table_id, segment_id)?;
        let meta: SegmentMetaV2 = bincode_config()
            .deserialize(&stats)
            .map_err(|e| ColumnarError::InvalidFormat(e.to_string()))?;
        Ok(meta.schema.column_count())
    }

    /// すべてのセグメント (table_id, segment_id) を返す。
    ///
    /// セグメントインデックスキーをスキャンして全テーブルのセグメントを収集する。
    pub fn list_segments(&self) -> Result<Vec<(u32, u64)>> {
        let mut txn = self.store.begin(TxnMode::ReadOnly)?;
        let mut result = Vec::new();

        // PREFIX_SEGMENT_INDEX (0x12) で始まるキーをスキャンし、
        // 各テーブルのインデックスを読み取る
        let prefix = vec![key_layout::PREFIX_SEGMENT_INDEX];
        for (key, value) in txn.scan_prefix(&prefix)? {
            if key.len() >= 5 {
                // キー形式: [prefix(1) + table_id(4)]
                let table_id = u32::from_le_bytes([key[1], key[2], key[3], key[4]]);
                let index: Vec<u64> = bincode_config()
                    .deserialize(&value)
                    .map_err(|e| ColumnarError::InvalidFormat(e.to_string()))?;
                for seg_id in index {
                    result.push((table_id, seg_id));
                }
            }
        }
        Ok(result)
    }
}

impl RangeAddressableSegmentProvider for ColumnarKvsBridge {
    fn read_header(&self, segment: SegmentReferenceV08) -> Result<Vec<u8>> {
        self.read_v08_object(key_layout::v08_header_key(
            segment.table_id,
            segment.segment_id,
        ))
    }

    fn read_schema(&self, segment: SegmentReferenceV08) -> Result<Vec<u8>> {
        self.read_v08_object(key_layout::v08_schema_key(
            segment.table_id,
            segment.segment_id,
        ))
    }

    fn read_directory(&self, segment: SegmentReferenceV08) -> Result<Vec<u8>> {
        self.read_v08_object(key_layout::v08_directory_key(
            segment.table_id,
            segment.segment_id,
        ))
    }

    fn read_chunk(
        &self,
        segment: SegmentReferenceV08,
        row_group_index: u32,
        column_index: u16,
    ) -> Result<Vec<u8>> {
        self.read_v08_object(key_layout::v08_chunk_key(
            segment.table_id,
            segment.segment_id,
            row_group_index,
            column_index,
        ))
    }
}

#[cfg(all(test, not(target_arch = "wasm32")))]
mod tests {
    use super::*;
    use crate::columnar::encoding::{Column, LogicalType};
    use crate::columnar::segment_v2::{ColumnSchema, RecordBatch, Schema};
    use crate::kv::memory::MemoryKV;
    use tempfile::tempdir;

    fn simple_schema() -> Schema {
        Schema {
            columns: vec![
                ColumnSchema {
                    name: "id".into(),
                    logical_type: LogicalType::Int64,
                    nullable: false,
                    fixed_len: None,
                },
                ColumnSchema {
                    name: "val".into(),
                    logical_type: LogicalType::Int64,
                    nullable: false,
                    fixed_len: None,
                },
            ],
        }
    }

    fn make_segment() -> ColumnSegmentV2 {
        let batch = RecordBatch::new(
            simple_schema(),
            vec![
                Column::Int64(vec![1, 2, 3]),
                Column::Int64(vec![10, 20, 30]),
            ],
            vec![None, None],
        );
        let mut writer = crate::columnar::segment_v2::SegmentWriterV2::new(Default::default());
        writer.write_batch(batch).unwrap();
        writer.finish().unwrap()
    }

    #[test]
    fn test_segment_atomic_write_transaction() {
        let store = AnyKV::Memory(MemoryKV::new_with_limit(Some(16)));
        let bridge = ColumnarKvsBridge::new(Arc::new(store));
        let segment = make_segment();

        // セグメントバイトが大きすぎてメモリリミットに抵触し、コミットが拒否される。
        let err = bridge.write_segment(1, &segment).unwrap_err();
        assert!(matches!(err, ColumnarError::MemoryLimitExceeded { .. }));

        // インデックスとデータは存在しない。
        let store = bridge.store.clone();
        let manager = store.txn_manager();
        let mut txn = manager.begin(TxnMode::ReadOnly).unwrap();
        let index_key = key_layout::segment_index_key(1);
        assert!(txn.get(&index_key).unwrap().is_none());
        manager.commit(txn).unwrap();
    }

    #[test]
    fn test_segment_write_read_via_kvs() {
        let store = Arc::new(AnyKV::Memory(MemoryKV::new()));
        let bridge = ColumnarKvsBridge::new(store.clone());
        let segment = make_segment();
        let id = bridge.write_segment(7, &segment).unwrap();
        assert_eq!(id, 0);

        let batches = bridge.read_segment(7, id, &[0, 1]).unwrap();
        assert_eq!(batches.len(), 1);
        if let Column::Int64(ids) = &batches[0].columns[0] {
            assert_eq!(ids, &vec![1, 2, 3]);
        } else {
            panic!("expected Column::Int64");
        }
    }

    #[test]
    fn test_segment_consistency_after_crash() {
        let dir = tempdir().unwrap();
        let wal_path = dir.path().join("wal.log");
        let segment = make_segment();
        {
            let store = Arc::new(AnyKV::Memory(MemoryKV::open(&wal_path).unwrap()));
            let bridge = ColumnarKvsBridge::new(store);
            bridge.write_segment(9, &segment).unwrap();
        }
        // WAL から復旧
        let reopened = Arc::new(AnyKV::Memory(MemoryKV::open(&wal_path).unwrap()));
        let bridge = ColumnarKvsBridge::new(reopened);
        let batches = bridge.read_segment(9, 0, &[1]).unwrap();
        if let Column::Int64(vals) = &batches[0].columns[0] {
            assert_eq!(vals, &vec![10, 20, 30]);
        } else {
            panic!("expected Column::Int64");
        }
    }

    #[test]
    fn test_multiple_segments_concurrent_access() {
        let store = Arc::new(AnyKV::Memory(MemoryKV::new()));
        let bridge = ColumnarKvsBridge::new(store.clone());
        let segment = make_segment();
        let bridge_arc = Arc::new(bridge);

        let mut handles = Vec::new();
        for _ in 0..4 {
            let b = bridge_arc.clone();
            let seg = segment.clone();
            handles.push(std::thread::spawn(move || {
                b.write_segment(3, &seg).unwrap();
            }));
        }
        for h in handles {
            h.join().unwrap();
        }

        let index = bridge_arc.segment_index(3).unwrap();
        assert_eq!(index.len(), 4);
        // すべて読めるか確認
        for id in index {
            let batches = bridge_arc.read_segment(3, id, &[0]).unwrap();
            assert_eq!(batches[0].num_rows(), 3);
        }
    }

    #[test]
    fn v08_preflight_rejects_legacy_v2_without_opening_it() {
        let store = Arc::new(AnyKV::Memory(MemoryKV::new()));
        let bridge = ColumnarKvsBridge::new(store);
        let segment_id = bridge.write_segment(17, &make_segment()).unwrap();

        assert_eq!(
            bridge
                .streaming_layout_preflight_v08(17, segment_id)
                .unwrap(),
            StreamingSegmentLayoutPreflightV08::RequiresV08ChunkedLayout
        );
        assert!(matches!(
            bridge.open_v08_segment(17, segment_id),
            Err(ColumnarError::RequiresV08ChunkedLayout)
        ));
    }

    #[test]
    fn v08_preflight_prioritizes_provisioned_header_key() {
        let store = Arc::new(AnyKV::Memory(MemoryKV::new()));
        let bridge = ColumnarKvsBridge::new(store.clone());
        let manager = store.txn_manager();
        let mut txn = manager.begin(TxnMode::ReadWrite).unwrap();
        txn.put(key_layout::v08_header_key(18, 4), vec![1, 2, 3])
            .unwrap();
        manager.commit(txn).unwrap();

        assert_eq!(
            bridge.streaming_layout_preflight_v08(18, 4).unwrap(),
            StreamingSegmentLayoutPreflightV08::Available
        );
    }
}