//! Logical and physical plan nodes backing the `TRIMSTRINGS <byte_len>` command.
//!
//! On DataFusion, `TRIMSTRINGS` is modeled as a real plan operator end to end:
//! - `TrimStringsLogicalNode` is a `UserDefinedLogicalNode` carrying the byte-length
//!   constant and a passthrough schema.
//! - The physical operator (`TrimStringsExec`) lives separately and is constructed
//!   by an `ExtensionPlanner` registered on the `SessionState`.

use std::any::Any;
use std::cmp::Ordering;
use std::fmt;
use std::hash::{Hash, Hasher};
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};

use arrow_data::ByteView;
use async_trait::async_trait;
use datafusion::arrow::array::{
    make_array, Array, ArrayRef, DictionaryArray, GenericListArray, GenericStringArray, MapArray,
    OffsetSizeTrait, StringViewArray, StructArray,
};
use datafusion::arrow::buffer::NullBuffer;
use datafusion::arrow::datatypes::{
    ArrowDictionaryKeyType, ArrowNativeType, DataType, Field, FieldRef, Fields, Int16Type,
    Int32Type, Int64Type, Int8Type, SchemaRef, UInt16Type, UInt32Type, UInt64Type, UInt8Type,
};
use datafusion::arrow::record_batch::RecordBatch;
use datafusion::common::{DFSchemaRef, Result as DfResult};
use datafusion::error::DataFusionError;
use datafusion::execution::{SessionState, TaskContext};
use datafusion::logical_expr::{
    Expr, LogicalPlan, UserDefinedLogicalNode, UserDefinedLogicalNodeCore,
};
use datafusion::physical_plan::metrics::{
    BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet,
};
use datafusion::physical_plan::ExecutionPlanProperties;
use datafusion::physical_plan::{
    DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, RecordBatchStream,
    SendableRecordBatchStream,
};
use datafusion::physical_planner::{ExtensionPlanner, PhysicalPlanner};
use datafusion::scalar::ScalarValue;
use futures::stream::{Stream, StreamExt};

/// Logical-plan extension node for `TRIMSTRINGS <byte_len>`.
///
/// Schema is the input schema (TRIMSTRINGS preserves column names and types — only the
/// byte-length of string values changes).
#[derive(Debug, Clone)]
pub struct TrimStringsLogicalNode {
    pub input: LogicalPlan,
    pub byte_len: u64,
    pub schema: DFSchemaRef,
}

impl TrimStringsLogicalNode {
    pub fn new(input: LogicalPlan, byte_len: u64) -> Self {
        let schema = input.schema().clone();
        Self {
            input,
            byte_len,
            schema,
        }
    }
}

// PartialEq, Eq, PartialOrd, and Hash all key off byte_len only — not the input plan.
//
// DataFusion's optimizer uses these to dedupe/cache plan nodes. The input plan is
// reachable via inputs() and is compared / hashed separately by the framework when
// it walks the tree, so we only need this node's own discriminator (byte_len) here.
// Including input in PartialEq while excluding it from Hash would violate the
// Eq ⇒ Hash invariant, so we keep all four impls in lockstep on byte_len only.
impl PartialEq for TrimStringsLogicalNode {
    fn eq(&self, other: &Self) -> bool {
        self.byte_len == other.byte_len
    }
}

impl Eq for TrimStringsLogicalNode {}

impl PartialOrd for TrimStringsLogicalNode {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.byte_len.cmp(&other.byte_len))
    }
}

impl Hash for TrimStringsLogicalNode {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.byte_len.hash(state);
    }
}

impl UserDefinedLogicalNodeCore for TrimStringsLogicalNode {
    fn name(&self) -> &str {
        "TrimStrings"
    }

    fn inputs(&self) -> Vec<&LogicalPlan> {
        vec![&self.input]
    }

    fn schema(&self) -> &DFSchemaRef {
        &self.schema
    }

    fn expressions(&self) -> Vec<Expr> {
        vec![Expr::Literal(
            ScalarValue::UInt64(Some(self.byte_len)),
            None,
        )]
    }

    fn fmt_for_explain(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "TrimStrings: byte_len={}", self.byte_len)
    }

    /// Forward parent column requirements to our single child verbatim.
    ///
    /// TRIMSTRINGS preserves the row-shape and the column schema — it only
    /// truncates string-typed leaves byte-wise inside the columns it's
    /// asked to emit. So whatever output indices our parent demands, we
    /// need exactly the same input indices from our child, no more.
    ///
    /// Without this, the trait's default returns `None`, which makes
    /// DataFusion's `OptimizeProjections` bail out with `Transformed::no`
    /// for our entire subtree — every column upstream stays alive,
    /// including dead scan columns that should have been pruned. That
    /// in turn blocks `JsonRewrite` from fusing EXPLODE chains onto the
    /// json scan when a SET-form connector view leaves a top-level
    /// `line` sibling alive at the seed projection (SIM-4563).
    fn necessary_children_exprs(&self, output_columns: &[usize]) -> Option<Vec<Vec<usize>>> {
        Some(vec![output_columns.to_vec()])
    }

    /// TRIMSTRINGS mutates string values in place; it never adds or
    /// drops rows. A LIMIT above us trims N rows when pushed below
    /// instead of trimming all of them then taking the top N.
    fn supports_limit_pushdown(&self) -> bool {
        true
    }

    fn with_exprs_and_inputs(&self, exprs: Vec<Expr>, inputs: Vec<LogicalPlan>) -> DfResult<Self> {
        let byte_len = match exprs.as_slice() {
            [Expr::Literal(ScalarValue::UInt64(Some(n)), _)] => *n,
            _ => {
                return Err(DataFusionError::Plan(
                    "TrimStrings expects a single UInt64 literal expression".to_string(),
                ));
            }
        };
        let input = inputs.into_iter().next().ok_or_else(|| {
            DataFusionError::Plan("TrimStrings expects exactly one input".to_string())
        })?;
        Ok(Self::new(input, byte_len))
    }
}

// ============================================================================
// Byte-truncate kernel
// ============================================================================

/// True if the data type contains any string-typed leaf (recursively, through structs,
/// lists, and dictionaries).
fn type_has_string(dt: &DataType) -> bool {
    match dt {
        DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => true,
        DataType::Struct(fields) => fields.iter().any(|f| type_has_string(f.data_type())),
        DataType::List(f) | DataType::LargeList(f) => type_has_string(f.data_type()),
        DataType::Map(entry_field, _) => match entry_field.data_type() {
            // Map entries are Struct<key, value>; only check the value field.
            // Keys are identity/lookup strings and must not be truncated.
            DataType::Struct(fields) => fields
                .get(1)
                .is_some_and(|f| type_has_string(f.data_type())),
            _ => false,
        },
        DataType::Dictionary(_, value_type) => type_has_string(value_type),
        _ => false,
    }
}

/// Truncate a single string to at most `byte_len` bytes on a UTF-8 character boundary.
/// Returns the truncated slice and whether truncation happened (i.e. clamp count contributor).
#[inline]
fn truncate_str(s: &str, byte_len: usize) -> (&str, bool) {
    if s.len() <= byte_len {
        return (s, false);
    }
    (&s[..s.floor_char_boundary(byte_len)], true)
}

/// Truncate every string-typed leaf in `array` to `byte_len` bytes, returning the new array
/// and the number of clamped values. Recurses into structs and lists. Non-string columns
/// are returned unchanged.
///
/// `ancestor_validity` is the row-aligned validity bitmap from any enclosing struct
/// ancestors (AND-merged at each struct level). Leaves use it to *skip clamp counting*
/// for rows that are logically null due to an ancestor — even though the physical child
/// slot may hold a non-null value. The output array shape is unaffected; only
/// `trim_clamped` is corrected.
///
/// Note: list ancestry is *not* propagated through this mask (element offsets aren't
/// row-aligned with parent list rows, so AND-merging would require offset arithmetic).
/// In practice list-null rows from DataFusion's own producers carry zero-length offset
/// ranges, so no element positions would be counted anyway.
fn truncate_array(
    array: &ArrayRef,
    byte_len: usize,
    ancestor_validity: Option<&NullBuffer>,
) -> DfResult<(ArrayRef, u64)> {
    if !type_has_string(array.data_type()) {
        return Ok((array.clone(), 0));
    }
    match array.data_type() {
        DataType::Utf8 => truncate_generic_string::<i32>(array, byte_len, ancestor_validity),
        DataType::LargeUtf8 => truncate_generic_string::<i64>(array, byte_len, ancestor_validity),
        DataType::Utf8View => truncate_string_view(array, byte_len, ancestor_validity),
        DataType::Struct(_) => truncate_struct(array, byte_len, ancestor_validity),
        DataType::List(_) => truncate_list::<i32>(array, byte_len),
        DataType::LargeList(_) => truncate_list::<i64>(array, byte_len),
        DataType::Dictionary(_, _) => truncate_dictionary(array, byte_len, ancestor_validity),
        DataType::Map(_, _) => truncate_map(array, byte_len),
        other => Err(DataFusionError::Internal(format!(
            "TrimStringsExec: unsupported nested type {other:?}"
        ))),
    }
}

/// Rebuild a dictionary array with the values array truncated. Keys are preserved
/// as-is.
///
/// Counts clamps **per logical row**, not per distinct values-side string: for each
/// values-side index that clamps, we add the number of non-null keys pointing at it.
/// This honors the `trim_clamped` contract documented in `ExecutorStats`
/// ("column-values clamped" = sum across logical rows). Cost is O(distinct values +
/// rows) — one walk of values, one walk of keys.
///
/// `ancestor_validity` mirrors the same parameter in `truncate_array`: when an
/// enclosing struct is null at row i, that row's clamp must not be counted even
/// though the dictionary's keys/values may still hold non-null bytes there.
fn truncate_dictionary(
    array: &ArrayRef,
    byte_len: usize,
    ancestor_validity: Option<&NullBuffer>,
) -> DfResult<(ArrayRef, u64)> {
    /// Truncate dictionary values and return per-value clamp markers.
    /// Includes an offset-scan fast path: if every value is already within
    /// the limit the original values array is returned unchanged.
    fn truncate_dict_values<O: OffsetSizeTrait>(
        values: &ArrayRef,
        byte_len: usize,
    ) -> DfResult<(ArrayRef, Vec<bool>)> {
        let arr = values
            .as_any()
            .downcast_ref::<GenericStringArray<O>>()
            .ok_or_else(|| DataFusionError::Internal("expected GenericStringArray".to_string()))?;

        let all_within = arr
            .offsets()
            .windows(2)
            .all(|w| (w[1] - w[0]).as_usize() <= byte_len);
        if all_within {
            return Ok((values.clone(), vec![false; arr.len()]));
        }

        let mut markers = vec![false; arr.len()];
        let mut truncated = Vec::with_capacity(arr.len());
        for i in 0..arr.len() {
            let (trimmed, clamped) = truncate_str(arr.value(i), byte_len);
            markers[i] = clamped;
            truncated.push(trimmed);
        }
        let built: GenericStringArray<O> = truncated.into_iter().map(Some).collect();
        Ok((overlay_nulls(built, arr.nulls())?, markers))
    }

    /// Dictionary values stored as `Utf8View` (same layout as top-level string views).
    fn truncate_dict_values_view(
        values: &ArrayRef,
        byte_len: usize,
    ) -> DfResult<(ArrayRef, Vec<bool>)> {
        let arr = values
            .as_any()
            .downcast_ref::<StringViewArray>()
            .ok_or_else(|| DataFusionError::Internal("expected StringViewArray".to_string()))?;

        let limit = u32::try_from(byte_len).unwrap_or(u32::MAX);
        let all_within = arr
            .views()
            .iter()
            .all(|v| ByteView::from(*v).length <= limit);
        if all_within {
            return Ok((values.clone(), vec![false; arr.len()]));
        }

        let mut markers = vec![false; arr.len()];
        let mut truncated = Vec::with_capacity(arr.len());
        for i in 0..arr.len() {
            let (trimmed, clamped) = truncate_str(arr.value(i), byte_len);
            markers[i] = clamped;
            truncated.push(trimmed);
        }
        let built: StringViewArray = truncated.into_iter().map(Some).collect();
        Ok((overlay_nulls(built, arr.nulls())?, markers))
    }

    fn truncate_values_indexed(
        values: &ArrayRef,
        byte_len: usize,
    ) -> DfResult<(ArrayRef, Vec<bool>)> {
        match values.data_type() {
            DataType::Utf8 => truncate_dict_values::<i32>(values, byte_len),
            DataType::LargeUtf8 => truncate_dict_values::<i64>(values, byte_len),
            DataType::Utf8View => truncate_dict_values_view(values, byte_len),
            other => Err(DataFusionError::Internal(format!(
                "TrimStringsExec: unsupported dictionary value type {other:?}"
            ))),
        }
    }

    fn rebuild<K: ArrowDictionaryKeyType>(
        dict: &DictionaryArray<K>,
        byte_len: usize,
        ancestor_validity: Option<&NullBuffer>,
    ) -> DfResult<(ArrayRef, u64)> {
        let values = dict.values().clone();
        let (new_values, markers) = truncate_values_indexed(&values, byte_len)?;
        let keys = dict.keys();
        // Skip the per-row keys walk entirely when no values clamp — every key
        // would land on a `markers[idx] == false` and contribute 0 anyway. Common
        // case for ingest pipelines that aren't actually capping anything in this
        // batch.
        if !markers.iter().any(|&m| m) {
            let new_dict = DictionaryArray::<K>::try_new(keys.clone(), new_values)
                .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?;
            return Ok((Arc::new(new_dict) as ArrayRef, 0));
        }
        // Walk the keys: for each non-null key i pointing at a clamped values-side
        // index, contribute one clamp — but skip rows that are null in any ancestor
        // struct's validity bitmap. Branch on null_count and ancestor presence so
        // the common no-nulls / no-ancestor case stays a tight loop.
        let mut clamps: u64 = 0;
        match (keys.null_count() == 0, ancestor_validity) {
            (true, None) => {
                for i in 0..keys.len() {
                    let idx = keys.value(i).as_usize();
                    if idx < markers.len() && markers[idx] {
                        clamps += 1;
                    }
                }
            }
            (true, Some(ancestor)) => {
                for i in 0..keys.len() {
                    if !ancestor.is_valid(i) {
                        continue;
                    }
                    let idx = keys.value(i).as_usize();
                    if idx < markers.len() && markers[idx] {
                        clamps += 1;
                    }
                }
            }
            (false, None) => {
                for i in 0..keys.len() {
                    if keys.is_null(i) {
                        continue;
                    }
                    let idx = keys.value(i).as_usize();
                    if idx < markers.len() && markers[idx] {
                        clamps += 1;
                    }
                }
            }
            (false, Some(ancestor)) => {
                for i in 0..keys.len() {
                    if keys.is_null(i) || !ancestor.is_valid(i) {
                        continue;
                    }
                    let idx = keys.value(i).as_usize();
                    if idx < markers.len() && markers[idx] {
                        clamps += 1;
                    }
                }
            }
        }
        let new_dict = DictionaryArray::<K>::try_new(keys.clone(), new_values)
            .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?;
        Ok((Arc::new(new_dict) as ArrayRef, clamps))
    }

    fn downcast_dict<K: ArrowDictionaryKeyType>(array: &ArrayRef) -> DfResult<&DictionaryArray<K>> {
        array
            .as_any()
            .downcast_ref::<DictionaryArray<K>>()
            .ok_or_else(|| {
                DataFusionError::Internal(format!(
                    "TrimStringsExec: failed to downcast {:?} to DictionaryArray",
                    array.data_type()
                ))
            })
    }

    match array.data_type() {
        DataType::Dictionary(k, _) => match k.as_ref() {
            DataType::Int8 => rebuild::<Int8Type>(
                downcast_dict::<Int8Type>(array)?,
                byte_len,
                ancestor_validity,
            ),
            DataType::Int16 => rebuild::<Int16Type>(
                downcast_dict::<Int16Type>(array)?,
                byte_len,
                ancestor_validity,
            ),
            DataType::Int32 => rebuild::<Int32Type>(
                downcast_dict::<Int32Type>(array)?,
                byte_len,
                ancestor_validity,
            ),
            DataType::Int64 => rebuild::<Int64Type>(
                downcast_dict::<Int64Type>(array)?,
                byte_len,
                ancestor_validity,
            ),
            DataType::UInt8 => rebuild::<UInt8Type>(
                downcast_dict::<UInt8Type>(array)?,
                byte_len,
                ancestor_validity,
            ),
            DataType::UInt16 => rebuild::<UInt16Type>(
                downcast_dict::<UInt16Type>(array)?,
                byte_len,
                ancestor_validity,
            ),
            DataType::UInt32 => rebuild::<UInt32Type>(
                downcast_dict::<UInt32Type>(array)?,
                byte_len,
                ancestor_validity,
            ),
            DataType::UInt64 => rebuild::<UInt64Type>(
                downcast_dict::<UInt64Type>(array)?,
                byte_len,
                ancestor_validity,
            ),
            other => Err(DataFusionError::Internal(format!(
                "TrimStringsExec: unsupported dictionary key type {other:?}"
            ))),
        },
        other => Err(DataFusionError::Internal(format!(
            "TrimStringsExec: expected Dictionary, got {other:?}"
        ))),
    }
}

/// Truncate every string value in `array`, returning the new array and clamped-row count.
///
/// The loop calls `value(i)` unconditionally — on null slots this returns `""` for
/// well-formed Arrow string arrays, so `truncate_str` is a no-op. The original null
/// bitmap is overlaid onto the result afterward, eliminating a per-row null branch
/// from the hot path.
fn truncate_generic_string<O: OffsetSizeTrait>(
    array: &ArrayRef,
    byte_len: usize,
    ancestor_validity: Option<&NullBuffer>,
) -> DfResult<(ArrayRef, u64)> {
    let arr = array
        .as_any()
        .downcast_ref::<GenericStringArray<O>>()
        .ok_or_else(|| DataFusionError::Internal("expected GenericStringArray".to_string()))?;

    // Fast path: scan the compact i32/i64 offsets buffer — no string data touched.
    // The tight integer-subtraction loop is auto-vectorisable by LLVM.
    let all_within = arr
        .offsets()
        .windows(2)
        .all(|w| (w[1] - w[0]).as_usize() <= byte_len);
    if all_within {
        return Ok((array.clone(), 0));
    }

    let (values, clamps) =
        truncate_string_values(arr.len(), byte_len, ancestor_validity, |i| arr.value(i));
    let built: GenericStringArray<O> = values.into_iter().map(Some).collect();
    Ok((overlay_nulls(built, arr.nulls())?, clamps))
}

fn truncate_string_view(
    array: &ArrayRef,
    byte_len: usize,
    ancestor_validity: Option<&NullBuffer>,
) -> DfResult<(ArrayRef, u64)> {
    let arr = array
        .as_any()
        .downcast_ref::<StringViewArray>()
        .ok_or_else(|| DataFusionError::Internal("expected StringViewArray".to_string()))?;

    // Fast path: in Arrow's StringView layout every u128 view stores the string
    // length, prefix, buffer index, and offset in fixed positions (see
    // `arrow_data::ByteView`). Use the typesafe accessor to extract just the
    // length, regardless of whether the value is inline (≤ 12 bytes) or
    // out-of-line. Scanning this contiguous fixed-stride buffer avoids touching
    // string data entirely.
    let limit = u32::try_from(byte_len).unwrap_or(u32::MAX);
    let all_within = arr
        .views()
        .iter()
        .all(|v| ByteView::from(*v).length <= limit);
    if all_within {
        return Ok((array.clone(), 0));
    }

    let (values, clamps) =
        truncate_string_values(arr.len(), byte_len, ancestor_validity, |i| arr.value(i));
    let built: StringViewArray = values.into_iter().map(Some).collect();
    Ok((overlay_nulls(built, arr.nulls())?, clamps))
}

/// Shared truncation loop for all string array types. Calls `value_fn(i)` for every
/// row unconditionally (callers guarantee that null slots return `""`).
/// Returns truncated slices and the number of clamped, logically-valid rows.
fn truncate_string_values<'a>(
    len: usize,
    byte_len: usize,
    ancestor_validity: Option<&NullBuffer>,
    value_fn: impl Fn(usize) -> &'a str,
) -> (Vec<&'a str>, u64) {
    let mut clamps: u64 = 0;
    let mut values = Vec::with_capacity(len);
    for i in 0..len {
        let (trimmed, clamped) = truncate_str(value_fn(i), byte_len);
        if clamped && ancestor_validity.map_or(true, |buf| buf.is_valid(i)) {
            clamps += 1;
        }
        values.push(trimmed);
    }
    (values, clamps)
}

/// Replace the null bitmap of `array` with `source_nulls`, preserving all other buffers.
fn overlay_nulls(
    array: impl Array + 'static,
    source_nulls: Option<&NullBuffer>,
) -> DfResult<ArrayRef> {
    let Some(nulls) = source_nulls else {
        return Ok(Arc::new(array));
    };
    let data = array
        .to_data()
        .into_builder()
        .nulls(Some(nulls.clone()))
        .build()
        .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?;
    Ok(make_array(data))
}

fn truncate_struct(
    array: &ArrayRef,
    byte_len: usize,
    ancestor_validity: Option<&NullBuffer>,
) -> DfResult<(ArrayRef, u64)> {
    let s = array
        .as_any()
        .downcast_ref::<StructArray>()
        .ok_or_else(|| DataFusionError::Internal("expected StructArray".to_string()))?;
    // AND-merge the ancestor validity with this struct's own nulls before recursing
    // into children. Children downstream will skip clamp counts for rows where any
    // ancestor (or this struct) is null.
    let merged_validity: Option<NullBuffer> = match (ancestor_validity, s.nulls()) {
        (None, None) => None,
        (Some(a), None) => Some(a.clone()),
        (None, Some(b)) => Some(b.clone()),
        (Some(a), Some(b)) => NullBuffer::union(Some(a), Some(b)),
    };
    let mut total_clamps: u64 = 0;
    let mut new_columns: Vec<(FieldRef, ArrayRef)> = Vec::with_capacity(s.num_columns());
    for (field, child) in s.fields().iter().zip(s.columns().iter()) {
        let (new_child, clamps) = truncate_array(child, byte_len, merged_validity.as_ref())?;
        total_clamps += clamps;
        new_columns.push((field.clone(), new_child));
    }
    let new_struct = StructArray::try_new(
        s.fields().clone(),
        new_columns.into_iter().map(|(_, a)| a).collect(),
        s.nulls().cloned(),
    )
    .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?;
    Ok((Arc::new(new_struct) as ArrayRef, total_clamps))
}

fn truncate_list<O: OffsetSizeTrait>(
    array: &ArrayRef,
    byte_len: usize,
) -> DfResult<(ArrayRef, u64)> {
    let list = array
        .as_any()
        .downcast_ref::<GenericListArray<O>>()
        .ok_or_else(|| DataFusionError::Internal("expected GenericListArray".to_string()))?;
    // Extract the element FieldRef from the array's data_type (List/LargeList store it there).
    let field: FieldRef = match list.data_type() {
        DataType::List(f) | DataType::LargeList(f) => f.clone(),
        other => {
            return Err(DataFusionError::Internal(format!(
                "TrimStringsExec: expected List or LargeList, got {other:?}"
            )));
        }
    };
    // List rows aren't 1:1 with the flat values array, so we don't propagate parent-row
    // validity into the recursive call; element-level validity is already handled by the
    // leaf string array's own null bitmap.
    let (new_values, clamps) = truncate_array(list.values(), byte_len, None)?;
    let new_list = GenericListArray::<O>::try_new(
        field,
        list.offsets().clone(),
        new_values,
        list.nulls().cloned(),
    )
    .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?;
    Ok((Arc::new(new_list) as ArrayRef, clamps))
}

/// Truncate string values inside a map column, leaving keys untouched.
///
/// A `MapArray` is physically `List<Struct<key, value>>`. We recurse into the
/// value column (index 1) of the entries struct only — keys are identity strings
/// used for lookups and must not be modified.
fn truncate_map(array: &ArrayRef, byte_len: usize) -> DfResult<(ArrayRef, u64)> {
    let map = array
        .as_any()
        .downcast_ref::<MapArray>()
        .ok_or_else(|| DataFusionError::Internal("expected MapArray".to_string()))?;

    let entries = map.entries();
    let keys = entries.column(0).clone();
    let values = entries.column(1);

    let (new_values, clamps) = truncate_array(values, byte_len, None)?;

    let DataType::Map(entry_field, sorted) = map.data_type() else {
        return Err(DataFusionError::Internal(
            "MapArray data_type is not Map".to_string(),
        ));
    };
    let DataType::Struct(entry_fields) = entry_field.data_type() else {
        return Err(DataFusionError::Internal(
            "expected Struct inside Map".to_string(),
        ));
    };
    let value_field = Arc::new(Field::new(
        entry_fields[1].name(),
        new_values.data_type().clone(),
        entry_fields[1].is_nullable(),
    ));
    let new_struct_fields = Fields::from(vec![entry_fields[0].clone(), value_field]);
    let new_entries = StructArray::new(
        new_struct_fields.clone(),
        vec![keys, new_values],
        entries.nulls().cloned(),
    );
    let new_map_field = Arc::new(Field::new(
        entry_field.name(),
        DataType::Struct(new_struct_fields),
        entry_field.is_nullable(),
    ));
    let new_map = MapArray::new(
        new_map_field,
        map.offsets().clone(),
        new_entries,
        map.nulls().cloned(),
        *sorted,
    );
    Ok((Arc::new(new_map) as ArrayRef, clamps))
}

fn truncate_record_batch(
    batch: &RecordBatch,
    byte_len: usize,
    schema: &SchemaRef,
) -> DfResult<(RecordBatch, u64)> {
    let mut new_cols: Vec<ArrayRef> = Vec::with_capacity(batch.num_columns());
    let mut clamps: u64 = 0;
    for col in batch.columns() {
        let (new_col, c) = truncate_array(col, byte_len, None)?;
        new_cols.push(new_col);
        clamps += c;
    }
    let new_batch = RecordBatch::try_new(schema.clone(), new_cols)
        .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?;
    Ok((new_batch, clamps))
}

// ============================================================================
// Physical operator
// ============================================================================

/// Physical operator backing `TRIMSTRINGS`.
///
/// Schema is the input schema unchanged — the kernel preserves every column's `DataType`
/// (`Utf8` stays `Utf8`, etc.) so downstream operators don't have to deal with type churn.
/// Only the byte-length of string values changes.
#[derive(Debug)]
pub struct TrimStringsExec {
    input: Arc<dyn ExecutionPlan>,
    byte_len: u64,
    metrics: ExecutionPlanMetricsSet,
    cache: Arc<PlanProperties>,
}

impl TrimStringsExec {
    pub fn new(input: Arc<dyn ExecutionPlan>, byte_len: u64) -> Self {
        let cache = Arc::new(Self::compute_properties(&input));
        Self {
            input,
            byte_len,
            metrics: ExecutionPlanMetricsSet::new(),
            cache,
        }
    }

    pub fn byte_len(&self) -> u64 {
        self.byte_len
    }

    fn compute_properties(input: &Arc<dyn ExecutionPlan>) -> PlanProperties {
        // TRIMSTRINGS is a row-preserving map: it doesn't reorder, repartition, or
        // change row count. Inherit every property from the input so the planner sees
        // the operator as a transparent passthrough (preserves orderings, propagates
        // bounded/unbounded correctly, etc.).
        PlanProperties::new(
            input.equivalence_properties().clone(),
            input.output_partitioning().clone(),
            input.pipeline_behavior(),
            input.boundedness(),
        )
    }
}

impl DisplayAs for TrimStringsExec {
    fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match t {
            DisplayFormatType::Default | DisplayFormatType::Verbose => {
                write!(f, "TrimStringsExec: byte_len={}", self.byte_len)
            }
            DisplayFormatType::TreeRender => {
                writeln!(f, "byte_len={}", self.byte_len)
            }
        }
    }
}

impl ExecutionPlan for TrimStringsExec {
    fn name(&self) -> &'static str {
        "TrimStringsExec"
    }

    fn as_any(&self) -> &dyn Any {
        self
    }

    fn properties(&self) -> &Arc<PlanProperties> {
        &self.cache
    }

    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
        vec![&self.input]
    }

    fn maintains_input_order(&self) -> Vec<bool> {
        vec![true]
    }

    fn benefits_from_input_partitioning(&self) -> Vec<bool> {
        vec![false]
    }

    fn with_new_children(
        self: Arc<Self>,
        children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> DfResult<Arc<dyn ExecutionPlan>> {
        Ok(Arc::new(TrimStringsExec::new(
            Arc::clone(&children[0]),
            self.byte_len,
        )))
    }

    fn execute(
        &self,
        partition: usize,
        context: Arc<TaskContext>,
    ) -> DfResult<SendableRecordBatchStream> {
        let input_stream = self.input.execute(partition, context)?;
        let baseline = BaselineMetrics::new(&self.metrics, partition);
        let trim_clamped: Count =
            MetricBuilder::new(&self.metrics).counter("trim_clamped", partition);
        let schema = self.schema();
        let byte_len = self.byte_len as usize;

        let stream = TrimStringsStream {
            input: input_stream,
            schema: schema.clone(),
            byte_len,
            baseline,
            trim_clamped,
        };
        Ok(Box::pin(stream))
    }

    fn metrics(&self) -> Option<MetricsSet> {
        Some(self.metrics.clone_inner())
    }
}

struct TrimStringsStream {
    input: SendableRecordBatchStream,
    schema: SchemaRef,
    byte_len: usize,
    baseline: BaselineMetrics,
    trim_clamped: Count,
}

impl Stream for TrimStringsStream {
    type Item = DfResult<RecordBatch>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let timer = self.baseline.elapsed_compute().clone();
        let poll = match self.input.poll_next_unpin(cx) {
            Poll::Ready(Some(Ok(batch))) => {
                let _t = timer.timer();
                match truncate_record_batch(&batch, self.byte_len, &self.schema) {
                    Ok((new_batch, clamps)) => {
                        if clamps > 0 {
                            self.trim_clamped.add(clamps as usize);
                        }
                        Poll::Ready(Some(Ok(new_batch)))
                    }
                    Err(e) => Poll::Ready(Some(Err(e))),
                }
            }
            other => other,
        };
        self.baseline.record_poll(poll)
    }
}

impl RecordBatchStream for TrimStringsStream {
    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }
}

// ============================================================================
// ExtensionPlanner: lower TrimStringsLogicalNode → TrimStringsExec
// ============================================================================

/// Planner that recognizes `TrimStringsLogicalNode` and lowers it to `TrimStringsExec`.
///
/// Register on the `SessionState` (via `DefaultPhysicalPlanner::with_extension_planners`)
/// so DataFusion's physical planner produces the right node at planning time.
#[derive(Debug)]
pub struct TrimStringsExtensionPlanner;

#[async_trait]
impl ExtensionPlanner for TrimStringsExtensionPlanner {
    async fn plan_extension(
        &self,
        _planner: &dyn PhysicalPlanner,
        node: &dyn UserDefinedLogicalNode,
        _logical_inputs: &[&LogicalPlan],
        physical_inputs: &[Arc<dyn ExecutionPlan>],
        _session_state: &SessionState,
    ) -> DfResult<Option<Arc<dyn ExecutionPlan>>> {
        let Some(trim) = node.as_any().downcast_ref::<TrimStringsLogicalNode>() else {
            return Ok(None);
        };
        let input = physical_inputs.first().cloned().ok_or_else(|| {
            DataFusionError::Plan("TrimStringsExec expects exactly one physical input".to_string())
        })?;
        Ok(Some(Arc::new(TrimStringsExec::new(input, trim.byte_len))))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use datafusion::arrow::array::{Int32Array, StringArray};
    use datafusion::arrow::buffer::OffsetBuffer;

    /// Build a nullable struct where some rows are null at the struct level but the
    /// child string array still holds long values at those positions. The clamp
    /// counter must skip those rows.
    #[test]
    fn struct_nulls_do_not_contribute_to_trim_clamped() {
        // Child string column: 3 rows, all non-null, all longer than 3 bytes.
        let child = StringArray::from(vec!["hidden_long", "real_long", "also_long"]);
        let fields: Fields = vec![Field::new("s", DataType::Utf8, false)].into();
        // Parent struct: row 0 null, rows 1 and 2 valid.
        let validity = NullBuffer::from(vec![false, true, true]);
        let struct_arr =
            StructArray::try_new(fields, vec![Arc::new(child) as ArrayRef], Some(validity))
                .expect("struct construction");
        let array: ArrayRef = Arc::new(struct_arr);

        let (_, clamps) = truncate_array(&array, 3, None).expect("kernel");

        // Two valid rows clamp; the null row's child must be skipped.
        assert_eq!(clamps, 2, "only valid struct rows should contribute clamps");
    }

    /// Without an enclosing struct, leaf nullability still works correctly:
    /// null leaf rows are not counted, non-null long rows are.
    #[test]
    fn leaf_nulls_skipped() {
        let arr = StringArray::from(vec![Some("longer"), None, Some("ok"), Some("toolong")]);
        let arr_ref: ArrayRef = Arc::new(arr);
        let (_, clamps) = truncate_array(&arr_ref, 3, None).expect("kernel");
        assert_eq!(
            clamps, 2,
            "only non-null rows over byte_len should contribute"
        );
    }

    /// A Dictionary(Int32, Utf8) column should have its values array truncated.
    /// Without the dictionary branch, this column would silently pass through unchanged.
    #[test]
    fn dictionary_string_values_are_truncated() {
        let values = StringArray::from(vec!["short", "verylongstring", "tiny"]);
        let keys = Int32Array::from(vec![0, 1, 2, 1, 0]);
        let dict = DictionaryArray::<Int32Type>::try_new(keys, Arc::new(values))
            .expect("dict construction");
        let array: ArrayRef = Arc::new(dict);

        let (new_array, clamps) = truncate_array(&array, 4, None).expect("kernel");

        // Per-row counting: keys = [0, 1, 2, 1, 0]. Indices 0 ('short' = 5) and 1
        // ('verylongstring' = 14) clamp; index 2 ('tiny' = 4) does not. So rows 0, 1, 3, 4
        // all clamp (4 total) and row 2 (key=2) doesn't.
        assert_eq!(clamps, 4, "dictionary clamp counts per logical row");

        // Round-trip the new dictionary to confirm its values were truncated.
        let new_dict = new_array
            .as_any()
            .downcast_ref::<DictionaryArray<Int32Type>>()
            .expect("output is dictionary");
        let new_values = new_dict
            .values()
            .as_any()
            .downcast_ref::<StringArray>()
            .expect("string values");
        assert_eq!(new_values.value(0), "shor"); // 'short' clamped to 4
        assert_eq!(new_values.value(1), "very"); // 'verylongstring' clamped to 4
        assert_eq!(new_values.value(2), "tiny"); // 'tiny' passes through
    }

    /// A dictionary column nested inside a struct must truncate its values
    /// like any other string-bearing column.
    #[test]
    fn struct_containing_dictionary_truncates_values() {
        let values = StringArray::from(vec!["short", "verylongstring"]);
        let keys = Int32Array::from(vec![0, 1, 0]);
        let dict = DictionaryArray::<Int32Type>::try_new(keys, Arc::new(values))
            .expect("dict construction");
        let dict_field = Field::new(
            "tag",
            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
            true,
        );
        let fields: Fields = vec![dict_field].into();
        let struct_arr = StructArray::try_new(fields, vec![Arc::new(dict) as ArrayRef], None)
            .expect("struct construction");
        let array: ArrayRef = Arc::new(struct_arr);

        let (_, clamps) = truncate_array(&array, 4, None).expect("kernel");

        // keys = [0, 1, 0]. Both values clamp at byte_len=4 ('short'=5, 'verylongstring'=14).
        // Three logical rows; all three reference clamped values.
        assert_eq!(
            clamps, 3,
            "all dictionary-referenced rows should contribute"
        );
    }

    /// A dictionary column inside a list element must truncate its values.
    #[test]
    fn list_containing_dictionary_truncates_values() {
        let values = StringArray::from(vec!["short", "verylongstring"]);
        let keys = Int32Array::from(vec![0, 1, 0, 1]);
        let dict = DictionaryArray::<Int32Type>::try_new(keys, Arc::new(values))
            .expect("dict construction");
        let element_field = Arc::new(Field::new_list_field(
            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
            true,
        ));
        // 2 list rows: row 0 = [keys[0], keys[1]] (length 2), row 1 = [keys[2], keys[3]] (length 2).
        let offsets = OffsetBuffer::new(vec![0i32, 2, 4].into());
        let list = GenericListArray::<i32>::try_new(
            element_field,
            offsets,
            Arc::new(dict) as ArrayRef,
            None,
        )
        .expect("list construction");
        let array: ArrayRef = Arc::new(list);

        let (_, clamps) = truncate_array(&array, 4, None).expect("kernel");

        // 4 element positions across both list rows; both dictionary values clamp,
        // so all 4 contribute.
        assert_eq!(clamps, 4);
    }

    /// Regression for the ancestor_validity threading on dictionary columns:
    /// a null parent struct row must not contribute clamps, even though the
    /// underlying dictionary keys/values may still hold non-null bytes there.
    #[test]
    fn struct_nulls_mask_dictionary_clamps() {
        // Three rows of a dictionary column: keys = [0, 0, 0], all pointing at a
        // long string that would clamp.
        let values = StringArray::from(vec!["verylongstring"]);
        let keys = Int32Array::from(vec![0, 0, 0]);
        let dict = DictionaryArray::<Int32Type>::try_new(keys, Arc::new(values))
            .expect("dict construction");
        let dict_field = Field::new(
            "tag",
            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
            true,
        );
        let fields: Fields = vec![dict_field].into();
        // Parent struct: row 0 null, rows 1 and 2 valid.
        let validity = NullBuffer::from(vec![false, true, true]);
        let struct_arr =
            StructArray::try_new(fields, vec![Arc::new(dict) as ArrayRef], Some(validity))
                .expect("struct construction");
        let array: ArrayRef = Arc::new(struct_arr);

        let (_, clamps) = truncate_array(&array, 4, None).expect("kernel");

        // Only the two valid struct rows should contribute, even though the
        // dictionary keys array holds non-null entries at all three positions.
        assert_eq!(
            clamps, 2,
            "null parent struct rows must mask dictionary clamps"
        );
    }
}
