be/src/runtime/variant-value.h

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

#pragma once

#include <ostream>
#include <string>
#include <string_view>

#include "common/status.h"
#include "runtime/string-value.h"

namespace impala_udf {
class FunctionContext;
struct StringVal;
}

namespace impala {

// Physical type tags for the variant binary encoding.
// Corresponds to basic_type=0 (primitive) with type_info encoding the physical type.
enum class VariantPhysicalType : uint8_t {
  VNULL = 0,
  BOOLEAN_TRUE = 1,
  BOOLEAN_FALSE = 2,
  INT8 = 3,
  INT16 = 4,
  INT32 = 5,
  INT64 = 6,
  DOUBLE = 7,
  DECIMAL4 = 8,
  DECIMAL8 = 9,
  DECIMAL16 = 10,
  DATE = 11,
  TIMESTAMPTZ = 12,
  TIMESTAMPNTZ = 13,
  FLOAT = 14,
  BINARY = 15,
  STRING = 16,
  TIME = 17,
  TIMESTAMPTZ_NANOS = 18,
  TIMESTAMPNTZ_NANOS = 19,
  UUID = 20
};

inline std::ostream& operator<<(std::ostream& os, VariantPhysicalType pt) {
  switch (pt) {
    case VariantPhysicalType::VNULL: return os << "VNULL";
    case VariantPhysicalType::BOOLEAN_TRUE: return os << "BOOLEAN_TRUE";
    case VariantPhysicalType::BOOLEAN_FALSE: return os << "BOOLEAN_FALSE";
    case VariantPhysicalType::INT8: return os << "INT8";
    case VariantPhysicalType::INT16: return os << "INT16";
    case VariantPhysicalType::INT32: return os << "INT32";
    case VariantPhysicalType::INT64: return os << "INT64";
    case VariantPhysicalType::DOUBLE: return os << "DOUBLE";
    case VariantPhysicalType::DECIMAL4: return os << "DECIMAL4";
    case VariantPhysicalType::DECIMAL8: return os << "DECIMAL8";
    case VariantPhysicalType::DECIMAL16: return os << "DECIMAL16";
    case VariantPhysicalType::DATE: return os << "DATE";
    case VariantPhysicalType::TIMESTAMPTZ: return os << "TIMESTAMPTZ";
    case VariantPhysicalType::TIMESTAMPNTZ: return os << "TIMESTAMPNTZ";
    case VariantPhysicalType::FLOAT: return os << "FLOAT";
    case VariantPhysicalType::BINARY: return os << "BINARY";
    case VariantPhysicalType::STRING: return os << "STRING";
    case VariantPhysicalType::TIME: return os << "TIME";
    case VariantPhysicalType::TIMESTAMPTZ_NANOS: return os << "TIMESTAMPTZ_NANOS";
    case VariantPhysicalType::TIMESTAMPNTZ_NANOS: return os << "TIMESTAMPNTZ_NANOS";
    case VariantPhysicalType::UUID: return os << "UUID";
  }
  return os << "UNKNOWN(" << static_cast<int>(pt) << ")";
}

// Basic type tags encoded in bits 0-1 of the value header byte.
enum class VariantBasicType : uint8_t {
  PRIMITIVE = 0,
  SHORT_STRING = 1,
  OBJECT = 2,
  ARRAY = 3
};

inline std::ostream& operator<<(std::ostream& os, VariantBasicType bt) {
  switch (bt) {
    case VariantBasicType::PRIMITIVE: return os << "PRIMITIVE";
    case VariantBasicType::SHORT_STRING: return os << "SHORT_STRING";
    case VariantBasicType::OBJECT: return os << "OBJECT";
    case VariantBasicType::ARRAY: return os << "ARRAY";
  }
  return os << "UNKNOWN(" << static_cast<int>(bt) << ")";
}

// Result of looking up an object field, an array element or a path in a variant value.
enum class VariantLookupResult : uint8_t {
  // '*result' holds the value.
  FOUND,
  // No such field, the array index is out of range, a path step does not match the type
  // of the value (e.g. a field name on an array or on a scalar), or the path is
  // malformed. Callers map this to SQL NULL.
  NOT_FOUND,
  // The value, or the entry that was looked up, is malformed (e.g. a truncated header, a
  // field id outside the dictionary, or a field whose encoding does not fit).
  CORRUPT
};

inline std::ostream& operator<<(std::ostream& os, VariantLookupResult r) {
  switch (r) {
    case VariantLookupResult::FOUND: return os << "FOUND";
    case VariantLookupResult::NOT_FOUND: return os << "NOT_FOUND";
    case VariantLookupResult::CORRUPT: return os << "CORRUPT";
  }
  return os << "UNKNOWN(" << static_cast<int>(r) << ")";
}

// Parses and provides access to the variant metadata blob (field name dictionary).
// The metadata format:
//   - Header byte: version (bits 0-3), sorted flag (bit 4),
//     offset_size_minus_one (bits 6-7)
//   - Dictionary size: variable-width integer (1-4 bytes based on offset size)
//   - Offsets: (dict_size + 1) entries of offset_size bytes each
//   - String data: concatenated field name strings
class VariantMetadata {
 public:
  VariantMetadata() = default;

  // Initialize from a metadata blob. Returns error if the blob is malformed.
  Status Init(const uint8_t* data, uint32_t len);

  // Returns the number of field names in the dictionary.
  uint32_t DictionarySize() const { return dict_size_; }

  // Returns the field name at the given dictionary index.
  // index must be in [0, DictionarySize()).
  std::string_view GetFieldName(uint32_t index) const;

  // True if the dictionary claims to be sorted and unique (the sorted_strings bit).
  // Lookups do not rely on it.
  bool IsSorted() const { return is_sorted_; }

  bool IsValid() const { return offsets_ != nullptr; }

 private:
  uint32_t ReadOffset(uint32_t index) const;

  const uint8_t* offsets_ = nullptr;
  const uint8_t* string_data_ = nullptr;
  uint32_t string_data_len_ = 0;  // number of bytes in the string-data region
  uint32_t dict_size_ = 0;
  uint8_t version_ = 0;
  uint8_t offset_size_ = 0;  // 1, 2, 3, or 4 bytes per offset
  bool is_sorted_ = false;
};

// Provides access to a variant value blob. Requires a VariantMetadata for field name
// lookups when navigating objects.
//
// Value format header byte:
//   bits 0-1: basic_type (PRIMITIVE=0, SHORT_STRING=1, OBJECT=2, ARRAY=3)
//   bits 2-7: type_info
//     - For PRIMITIVE: physical type id
//     - For SHORT_STRING: string length (0-63 bytes)
//     - For OBJECT: field_offset_size_minus_one (bits 2-3),
//                   field_id_size_minus_one (bits 4-5), is_large (bit 6)
//     - For ARRAY: offset_size_minus_one (bits 2-3), is_large (bit 4)
class VariantValue {
 public:
  VariantValue() = default;
  VariantValue(const uint8_t* data, uint32_t len, const VariantMetadata* metadata)
      : data_(data), len_(len), metadata_(metadata) {}

  // Returns the basic type of this value.
  VariantBasicType GetBasicType() const;

  // Returns the physical type (only valid for PRIMITIVE basic type).
  VariantPhysicalType GetPhysicalType() const;

  // Returns true if this value is null.
  bool IsNull() const;

  // Scalar accessors. Each returns false and leaves '*out' unchanged if the value is not
  // of the expected physical type or its payload would read out of bounds.
  [[nodiscard]] bool GetBoolean(bool* out) const;
  [[nodiscard]] bool GetInt8(int8_t* out) const;
  [[nodiscard]] bool GetInt16(int16_t* out) const;
  [[nodiscard]] bool GetInt32(int32_t* out) const;
  [[nodiscard]] bool GetInt64(int64_t* out) const;
  [[nodiscard]] bool GetFloat(float* out) const;
  [[nodiscard]] bool GetDouble(double* out) const;
  [[nodiscard]] bool GetString(StringValue* out) const;
  [[nodiscard]] bool GetBinary(StringValue* out) const;

  // Object access by position. Each returns false if this value is not a well-formed
  // object (a corrupt or truncated header/offset table), if 'index' is out of range, or
  // if the field's encoding is out of bounds.
  [[nodiscard]] bool GetObjectSize(uint32_t* out) const;
  // Gets the field value by position index.
  [[nodiscard]] bool GetFieldByIndex(uint32_t index, VariantValue* result) const;
  // Gets the field name at position index in this object.
  [[nodiscard]] bool GetFieldNameByIndex(uint32_t index, std::string_view* out) const;
  // Gets the value of the field named 'name'. Searches this object's field ids by their
  // names, so it works with any dictionary (sorted or not, with or without duplicate
  // strings) and with field ids in UTF-16 order. NOT_FOUND if this is not an object or
  // has no such field; CORRUPT if this is an empty value, there is no metadata, the
  // object is malformed, or the field (or a field id compared before it) cannot be
  // read.
  [[nodiscard]] VariantLookupResult GetFieldByName(
      std::string_view name, VariantValue* result) const;

  // Array access. GetArraySize() returns false on a corrupt/truncated array.
  [[nodiscard]] bool GetArraySize(uint32_t* out) const;
  // NOT_FOUND if this is not an array or 'index' is out of range; CORRUPT if this is an
  // empty value, a malformed array, or the element cannot be read.
  [[nodiscard]] VariantLookupResult GetArrayElement(
      uint32_t index, VariantValue* result) const;

  // Navigates a path like "$.field.nested[0].value". '*result' is unspecified unless
  // FOUND. CORRUPT if this value (sized by ValueSize()) or a value looked up on the path
  // cannot be read. A malformed path gives NOT_FOUND (the FE rejects malformed constant
  // paths).
  [[nodiscard]] VariantLookupResult NavigatePath(
      std::string_view path, VariantValue* result) const;

  // Computes the encoded size of the value starting at 'data' from its header, like
  // parquet-java's VariantUtil.valueSize(): a fixed size for most primitives,
  // 1 + length for a short string, 5 + length for a long string or binary, and
  // header + tables + last offset for an object or array (nested values are not
  // validated). Returns false if 'data' is null or 'avail' is 0, the primitive type id
  // is unknown (21-63), an object or array header or offset table is malformed, or the
  // value would extend past 'avail' bytes.
  [[nodiscard]] static bool ValueSize(const uint8_t* data, uint32_t avail,
      uint32_t* size);

  // Serialize this variant value to JSON string.
  Status ToJson(std::string* json_out) const;
  Status ToJson(impala_udf::FunctionContext* ctx, impala_udf::StringVal* result) const;
  // Serialize this variant value as JSON directly into 'out'. The JSON is fully
  // buffered internally and only flushed to 'out' on success.
  Status ToJson(std::ostream* out) const;

  bool IsValid() const { return data_ != nullptr; }
  const uint8_t* Data() const { return data_; }
  uint32_t Len() const { return len_; }

  // Reads a primitive of type T from the payload at data_ + offset (offset defaults to 1,
  // just past the header byte). Returns false without modifying '*out' if the read would
  // extend past the end of the value buffer.
  template <typename T>
  [[nodiscard]] bool ReadValue(T* out, uint32_t offset = 1) const {
    if (offset > len_ || sizeof(T) > len_ - offset) return false;
    memcpy(out, data_ + offset, sizeof(T));
    return true;
  }

 private:
  // Helper to read a variable-width unsigned integer.
  static uint32_t ReadUint(const uint8_t* data, uint32_t size);

  // If this value is a PRIMITIVE with a readable header byte, sets '*pt' to its physical
  // type and returns true; otherwise returns false.
  [[nodiscard]] bool AsPrimitive(VariantPhysicalType* pt) const {
    if (data_ == nullptr || len_ < 1) return false;
    if (GetBasicType() != VariantBasicType::PRIMITIVE) return false;
    *pt = GetPhysicalType();
    return true;
  }

  // Checks this value is a PRIMITIVE of 'expected' physical type, then reads a T-sized
  // payload just past the header byte (see ReadValue()). Returns false without modifying
  // '*out' on a type mismatch or a truncated payload. Shared by the typed scalar getters.
  template <typename T>
  [[nodiscard]] bool ReadValueOfType(VariantPhysicalType expected, T* out) const {
    VariantPhysicalType pt;
    if (!AsPrimitive(&pt) || pt != expected) return false;
    return ReadValue(out);
  }

  // Validated internal layout of an object/array value. The pointers and lengths are
  // guaranteed to lie within the value buffer once the matching Parse*() returns true.
  struct ObjectLayout {
    uint32_t num_fields;
    const uint8_t* field_ids;
    const uint8_t* offsets;
    const uint8_t* data;
    // Size of the value data, i.e. the last offset, which is validated to fit in the
    // buffer. Bytes after it (e.g. padding) are not part of any value.
    uint32_t data_len;
  };
  struct ArrayLayout {
    uint32_t num_elems;
    const uint8_t* offsets;
    const uint8_t* data;
    // Size of the value data, i.e. the last offset, which is validated to fit in the
    // buffer. Bytes after it (e.g. padding) are not part of any value.
    uint32_t data_len;
  };
  // Parse and bounds-check the object/array header + offset table. Returns false (without
  // reading out of bounds) if the value is not that type or the header/tables do not fit.
  [[nodiscard]] bool ParseObjectLayout(ObjectLayout* out) const;
  [[nodiscard]] bool ParseArrayLayout(ArrayLayout* out) const;

  // Extract a single field/element from an already-parsed, validated layout. A value is
  // located by its entry's start offset, which is bounds-checked against data_len, and
  // sized from its own header; the next offset is not used, because values may be stored
  // in any order. For a field name, checks that its dictionary id is in range. Shared by
  // the public single-access accessors and by JSON serialization, which parses the
  // layout once and reuses it across the whole loop.
  [[nodiscard]] bool FieldFromLayout(const ObjectLayout& layout, uint32_t index,
      VariantValue* result) const;
  [[nodiscard]] bool FieldNameFromLayout(const ObjectLayout& layout, uint32_t index,
      std::string_view* out) const;
  // Sets '*index' to the position of the field named 'name' in 'layout'. NOT_FOUND if
  // there is no such field; CORRUPT if a field id it compares is outside the dictionary.
  [[nodiscard]] VariantLookupResult FindField(const ObjectLayout& layout,
      std::string_view name, uint32_t* index) const;
  [[nodiscard]] bool ElementFromLayout(const ArrayLayout& layout, uint32_t index,
      VariantValue* result) const;
  // Sets '*result' to the value that starts at offset 'start' of the value data
  // ['data', 'data' + 'data_len'), sized by ValueSize(). Returns false if it does not
  // fit.
  [[nodiscard]] bool SliceValue(const uint8_t* data, uint32_t data_len, uint32_t start,
      VariantValue* result) const;

  // JSON serialization lives in the .cc (so rapidjson stays out of this header) but needs
  // the private layout helpers above to parse each object/array once and reuse it.
  friend struct VariantJsonSerializer;

  // Bit-width helpers derived from the header byte (require a readable header).
  uint32_t ObjectFieldIdSize() const;
  uint32_t ObjectOffsetSize() const;
  uint32_t ArrayOffsetSize() const;

  const uint8_t* data_ = nullptr;
  uint32_t len_ = 0;
  const VariantMetadata* metadata_ = nullptr;
};

}  // namespace impala