From 6efb1e3e341b18a4992359876f5d9fe79c584210 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mikl=C3=B3s=20Koren?= Date: Mon, 1 Jun 2026 12:04:42 +0200 Subject: [PATCH] Factor out dta into separate extension (codedthinking/duckdb-dta) Replace src/dta/ with duckdb-dta git submodule. The dta extension registers read_dta() and COPY TO (FORMAT dta) independently. Dodo compiles dta_reader.cpp for variable label extraction on use. Tests split: pure SQL tests in duckdb-dta, dodo-specific tests (save, codebook, label variable) in test/sql/dodo_dta.test. Co-Authored-By: Claude Opus 4.6 (1M context) --- .gitmodules | 3 + CMakeLists.txt | 13 +- Makefile | 6 +- duckdb-dta | 1 + extension_config.cmake | 5 + src/dta/dta_reader.cpp | 624 ------------------------------- src/dta/dta_reader.hpp | 118 ------ src/dta/dta_writer.cpp | 304 --------------- src/dta/dta_writer.hpp | 74 ---- src/dta/read_dta_function.cpp | 374 ------------------ src/dta/read_dta_function.hpp | 10 - src/dta/write_dta_function.cpp | 446 ---------------------- src/dta/write_dta_function.hpp | 10 - src/extension/dodo_extension.cpp | 6 - test/sql/dodo_dta.test | 96 +++++ test/sql/read_dta.test | 122 ------ test/sql/write_dta.test | 471 ----------------------- 17 files changed, 114 insertions(+), 2569 deletions(-) create mode 160000 duckdb-dta delete mode 100644 src/dta/dta_reader.cpp delete mode 100644 src/dta/dta_reader.hpp delete mode 100644 src/dta/dta_writer.cpp delete mode 100644 src/dta/dta_writer.hpp delete mode 100644 src/dta/read_dta_function.cpp delete mode 100644 src/dta/read_dta_function.hpp delete mode 100644 src/dta/write_dta_function.cpp delete mode 100644 src/dta/write_dta_function.hpp create mode 100644 test/sql/dodo_dta.test delete mode 100644 test/sql/read_dta.test delete mode 100644 test/sql/write_dta.test diff --git a/.gitmodules b/.gitmodules index 69cda36..e509213 100644 --- a/.gitmodules +++ b/.gitmodules @@ -9,3 +9,6 @@ [submodule "duckdb-ui"] path = duckdb-ui url = https://github.com/duckdb/duckdb-ui.git +[submodule "duckdb-dta"] + path = duckdb-dta + url = https://github.com/codedthinking/duckdb-dta diff --git a/CMakeLists.txt b/CMakeLists.txt index 22e1184..e77d255 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -12,15 +12,13 @@ set(CMAKE_CXX_STANDARD "17" CACHE STRING "C++ standard to enforce") set(CMAKE_CXX_STANDARD_REQUIRED ON) include_directories(${CMAKE_CURRENT_SOURCE_DIR}/src/extension ${CMAKE_CURRENT_SOURCE_DIR}/src/core - ${CMAKE_CURRENT_SOURCE_DIR}/src/dta) + ${CMAKE_CURRENT_SOURCE_DIR}/duckdb-dta/src/include + ${CMAKE_CURRENT_SOURCE_DIR}/duckdb-dta/src) set(EXTENSION_SOURCES src/extension/dodo_extension.cpp src/core/dodo_core.cpp - src/dta/dta_reader.cpp - src/dta/read_dta_function.cpp - src/dta/dta_writer.cpp - src/dta/write_dta_function.cpp) + duckdb-dta/src/dta_reader.cpp) build_static_extension(${TARGET_NAME} ${EXTENSION_SOURCES}) build_loadable_extension(${TARGET_NAME} " " ${EXTENSION_SOURCES}) @@ -32,6 +30,7 @@ install( ARCHIVE DESTINATION "${INSTALL_LIB_DIR}") # CLI tool (standalone, no DuckDB dependency) -add_executable(dodoc src/cli/dodoc.cpp src/core/dodo_core.cpp src/dta/dta_reader.cpp) -target_include_directories(dodoc PRIVATE src/core src/dta) +add_executable(dodoc src/cli/dodoc.cpp src/core/dodo_core.cpp + duckdb-dta/src/dta_reader.cpp) +target_include_directories(dodoc PRIVATE src/core duckdb-dta/src/include) target_compile_features(dodoc PRIVATE cxx_std_17) diff --git a/Makefile b/Makefile index 7fa8154..19819bc 100644 --- a/Makefile +++ b/Makefile @@ -9,8 +9,8 @@ include extension-ci-tools/makefiles/duckdb_extension.Makefile # ---- dodoc standalone CLI (no DuckDB dependency) ---- DODOC_BUILD_DIR := build/dodoc -DODOC_SOURCES := src/cli/dodoc.cpp src/core/dodo_core.cpp -DODOC_HEADERS := src/core/dodo_core.hpp src/core/string_utils.hpp +DODOC_SOURCES := src/cli/dodoc.cpp src/core/dodo_core.cpp duckdb-dta/src/dta_reader.cpp +DODOC_HEADERS := src/core/dodo_core.hpp src/core/string_utils.hpp duckdb-dta/src/include/dta_reader.hpp CXX ?= c++ CXXFLAGS ?= -O2 -std=c++17 @@ -20,7 +20,7 @@ dodoc: $(DODOC_BUILD_DIR)/dodoc $(DODOC_BUILD_DIR)/dodoc: $(DODOC_SOURCES) $(DODOC_HEADERS) @mkdir -p $(DODOC_BUILD_DIR) - $(CXX) $(CXXFLAGS) -Isrc/core -o $@ $(DODOC_SOURCES) + $(CXX) $(CXXFLAGS) -Isrc/core -Iduckdb-dta/src/include -o $@ $(DODOC_SOURCES) dodoc-install: dodoc install -m 755 $(DODOC_BUILD_DIR)/dodoc /usr/local/bin/dodoc diff --git a/duckdb-dta b/duckdb-dta new file mode 160000 index 0000000..5cd3d4d --- /dev/null +++ b/duckdb-dta @@ -0,0 +1 @@ +Subproject commit 5cd3d4d55bc4b716ef2e6f081e61a53365be2810 diff --git a/extension_config.cmake b/extension_config.cmake index 9cf55a0..a04d565 100644 --- a/extension_config.cmake +++ b/extension_config.cmake @@ -1,5 +1,10 @@ # This file is included by DuckDB's build system. It specifies which extension to load +# Load dta extension first (dependency) +duckdb_extension_load(dta + SOURCE_DIR ${CMAKE_CURRENT_LIST_DIR}/duckdb-dta +) + # Extension from this repo duckdb_extension_load(dodo SOURCE_DIR ${CMAKE_CURRENT_LIST_DIR} diff --git a/src/dta/dta_reader.cpp b/src/dta/dta_reader.cpp deleted file mode 100644 index 9866c3c..0000000 --- a/src/dta/dta_reader.cpp +++ /dev/null @@ -1,624 +0,0 @@ -#include "dta_reader.hpp" - -#include -#include -#include -#include -#include - -namespace dta { - -const std::string DtaReader::empty_strl_; - -// ─── Version params ───────────────────────────────────────────────────────── - -DtaVersionParams DtaVersionParams::ForVersion(int version) { - DtaVersionParams p; - p.version = version; - switch (version) { - case 117: - p.varname_len = 33; - p.sortlist_entry_size = 2; - p.fmt_len = 49; - p.label_name_len = 33; - p.var_label_len = 81; - p.k_field_size = 2; - p.n_field_size = 4; - p.dataset_label_len_size = 1; - p.has_alias_vars = false; - break; - case 118: - p.varname_len = 129; - p.sortlist_entry_size = 2; - p.fmt_len = 57; - p.label_name_len = 129; - p.var_label_len = 321; - p.k_field_size = 2; - p.n_field_size = 8; - p.dataset_label_len_size = 2; - p.has_alias_vars = false; - break; - case 119: - p.varname_len = 129; - p.sortlist_entry_size = 4; - p.fmt_len = 57; - p.label_name_len = 129; - p.var_label_len = 321; - p.k_field_size = 4; - p.n_field_size = 8; - p.dataset_label_len_size = 2; - p.has_alias_vars = false; - break; - case 120: - p.varname_len = 129; - p.sortlist_entry_size = 2; - p.fmt_len = 57; - p.label_name_len = 129; - p.var_label_len = 321; - p.k_field_size = 2; - p.n_field_size = 8; - p.dataset_label_len_size = 2; - p.has_alias_vars = true; - break; - case 121: - p.varname_len = 129; - p.sortlist_entry_size = 4; - p.fmt_len = 57; - p.label_name_len = 129; - p.var_label_len = 321; - p.k_field_size = 4; - p.n_field_size = 8; - p.dataset_label_len_size = 2; - p.has_alias_vars = true; - break; - default: - throw std::runtime_error("Unsupported .dta format version: " + std::to_string(version)); - } - return p; -} - -// ─── Missing value detection ──────────────────────────────────────────────── - -bool DtaMissing::IsMissingByte(int8_t val) { - return val > 100; // >= 101 (0x65) -} - -bool DtaMissing::IsMissingInt(int16_t val) { - return val > 32740; // >= 32741 (0x7fe5) -} - -bool DtaMissing::IsMissingLong(int32_t val) { - return val > 2147483620; // >= 0x7fffffe5 -} - -bool DtaMissing::IsMissingFloat(float val) { - // Float missing: z > +1.fffffeX+7e i.e. > max nonmissing - // Use raw bit pattern: missing starts at 0x7f000000 - uint32_t bits; - memcpy(&bits, &val, 4); - return (bits & 0x7fffffff) >= 0x7f000000; -} - -bool DtaMissing::IsMissingDouble(double val) { - // Double missing: z > +1.fffffffffffffX+3fe - // Missing starts at 0x7fe0000000000000 - uint64_t bits; - memcpy(&bits, &val, 8); - return (bits & 0x7fffffffffffffff) >= 0x7fe0000000000000ULL; -} - -// ─── Byte width for type codes ────────────────────────────────────────────── - -uint16_t DtaTypeByteWidth(uint16_t type_code) { - if (type_code >= 1 && type_code <= 2045) { - return type_code; // str# has width = # - } - switch (type_code) { - case 32768: return 8; // strL: 8-byte (v,o) reference - case 65525: return 0; // alias: no data - case 65526: return 8; // double - case 65527: return 4; // float - case 65528: return 4; // long - case 65529: return 2; // int - case 65530: return 1; // byte - default: - throw std::runtime_error("Unknown .dta type code: " + std::to_string(type_code)); - } -} - -// ─── Byte-swap helpers ────────────────────────────────────────────────────── - -static uint16_t swap16(uint16_t val) { - return (val >> 8) | (val << 8); -} - -static uint32_t swap32(uint32_t val) { - return ((val >> 24) & 0xff) | ((val >> 8) & 0xff00) | ((val << 8) & 0xff0000) | ((val << 24) & 0xff000000); -} - -static uint64_t swap64(uint64_t val) { - val = ((val >> 8) & 0x00ff00ff00ff00ffULL) | ((val << 8) & 0xff00ff00ff00ff00ULL); - val = ((val >> 16) & 0x0000ffff0000ffffULL) | ((val << 16) & 0xffff0000ffff0000ULL); - return (val >> 32) | (val << 32); -} - -template <> -uint16_t DtaReader::SwapIfNeeded(uint16_t val) const { -#if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__ - return msf_ ? swap16(val) : val; -#else - return msf_ ? val : swap16(val); -#endif -} - -template <> -uint32_t DtaReader::SwapIfNeeded(uint32_t val) const { -#if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__ - return msf_ ? swap32(val) : val; -#else - return msf_ ? val : swap32(val); -#endif -} - -template <> -uint64_t DtaReader::SwapIfNeeded(uint64_t val) const { -#if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__ - return msf_ ? swap64(val) : val; -#else - return msf_ ? val : swap64(val); -#endif -} - -template <> -int16_t DtaReader::SwapIfNeeded(int16_t val) const { - uint16_t u; - memcpy(&u, &val, 2); - u = SwapIfNeeded(u); - int16_t result; - memcpy(&result, &u, 2); - return result; -} - -template <> -int32_t DtaReader::SwapIfNeeded(int32_t val) const { - uint32_t u; - memcpy(&u, &val, 4); - u = SwapIfNeeded(u); - int32_t result; - memcpy(&result, &u, 4); - return result; -} - -template <> -float DtaReader::SwapIfNeeded(float val) const { - uint32_t u; - memcpy(&u, &val, 4); - u = SwapIfNeeded(u); - float result; - memcpy(&result, &u, 4); - return result; -} - -template <> -double DtaReader::SwapIfNeeded(double val) const { - uint64_t u; - memcpy(&u, &val, 8); - u = SwapIfNeeded(u); - double result; - memcpy(&result, &u, 8); - return result; -} - -// ─── I/O helpers ──────────────────────────────────────────────────────────── - -void DtaReader::ReadBytes(void *buf, size_t n) { - if (fread(buf, 1, n, fp_) != n) { - throw std::runtime_error("Unexpected end of .dta file"); - } -} - -uint16_t DtaReader::ReadU16() { - uint16_t val; - ReadBytes(&val, 2); - return SwapIfNeeded(val); -} - -uint32_t DtaReader::ReadU32() { - uint32_t val; - ReadBytes(&val, 4); - return SwapIfNeeded(val); -} - -uint64_t DtaReader::ReadU64() { - uint64_t val; - ReadBytes(&val, 8); - return SwapIfNeeded(val); -} - -std::string DtaReader::ReadFixedString(uint32_t len) { - std::vector buf(len); - ReadBytes(buf.data(), len); - // Find null terminator - auto end = std::find(buf.begin(), buf.end(), '\0'); - return std::string(buf.begin(), end); -} - -void DtaReader::ReadTag(const char *expected) { - size_t len = strlen(expected); - std::vector buf(len); - ReadBytes(buf.data(), len); - if (memcmp(buf.data(), expected, len) != 0) { - throw std::runtime_error("Expected tag '" + std::string(expected) + "' not found in .dta file"); - } -} - -// ─── Constructor ──────────────────────────────────────────────────────────── - -DtaReader::DtaReader(const std::string &path) - : fp_(nullptr), msf_(false), n_obs_(0), row_width_(0), data_offset_(0), strls_offset_(0), - value_labels_offset_(0) { - fp_ = fopen(path.c_str(), "rb"); - if (!fp_) { - throw std::runtime_error("Cannot open .dta file: " + path); - } - - ParseHeader(); - ParseMap(); - ParseVariableTypes(); - ParseVarnames(); - // Skip sortlist - uint32_t n_vars_for_sort = static_cast(columns_.size()); - ReadTag(""); - fseek(fp_, (n_vars_for_sort + 1) * params_.sortlist_entry_size, SEEK_CUR); - ReadTag(""); - ParseFormats(); - ParseValueLabelNames(); - ParseVariableLabels(); - SkipCharacteristics(); - - // Compute row width - row_width_ = 0; - for (auto &col : columns_) { - row_width_ += col.byte_width; - } -} - -DtaReader::~DtaReader() { - if (fp_) { - fclose(fp_); - } -} - -// ─── Header parsing ───────────────────────────────────────────────────────── - -void DtaReader::ParseHeader() { - ReadTag(""); - ReadTag("
"); - - // Release - ReadTag(""); - char release_buf[4] = {}; - ReadBytes(release_buf, 3); - ReadTag(""); - int version = atoi(release_buf); - params_ = DtaVersionParams::ForVersion(version); - - // Byteorder - ReadTag(""); - char bo_buf[4] = {}; - ReadBytes(bo_buf, 3); - ReadTag(""); - msf_ = (memcmp(bo_buf, "MSF", 3) == 0); - - // K (number of variables) - ReadTag(""); - uint32_t n_vars; - if (params_.k_field_size == 2) { - n_vars = ReadU16(); - } else { - n_vars = ReadU32(); - } - ReadTag(""); - - // N (number of observations) - ReadTag(""); - if (params_.n_field_size == 4) { - n_obs_ = ReadU32(); - } else { - n_obs_ = ReadU64(); - } - ReadTag(""); - - // Dataset label - ReadTag(""); - - // Timestamp (skip) - ReadTag(""); - uint8_t ts_len; - ReadBytes(&ts_len, 1); - if (ts_len > 0) { - fseek(fp_, ts_len, SEEK_CUR); - } - ReadTag(""); - - ReadTag("
"); - - // Pre-allocate columns - columns_.resize(n_vars); -} - -// ─── Map ──────────────────────────────────────────────────────────────────── - -void DtaReader::ParseMap() { - ReadTag(""); - uint64_t offsets[14]; - for (int i = 0; i < 14; i++) { - offsets[i] = ReadU64(); - } - ReadTag(""); - - // offsets[9] = , offsets[10] = , offsets[11] = - data_offset_ = offsets[9]; - strls_offset_ = offsets[10]; - value_labels_offset_ = offsets[11]; -} - -// ─── Variable types ───────────────────────────────────────────────────────── - -void DtaReader::ParseVariableTypes() { - ReadTag(""); - std::vector raw_types(columns_.size()); - for (size_t i = 0; i < columns_.size(); i++) { - raw_types[i] = ReadU16(); - } - ReadTag(""); - - // Filter out alias variables (type 65525) for formats 120/121 - if (params_.has_alias_vars) { - std::vector filtered; - for (size_t i = 0; i < raw_types.size(); i++) { - if (raw_types[i] != 65525) { - columns_[i].type_code = raw_types[i]; - columns_[i].byte_width = DtaTypeByteWidth(raw_types[i]); - filtered.push_back(columns_[i]); - } - } - // We need to track original indices for varnames/formats/labels parsing - // For now, store all columns then filter after all metadata is parsed - for (size_t i = 0; i < columns_.size(); i++) { - columns_[i].type_code = raw_types[i]; - columns_[i].byte_width = (raw_types[i] == 65525) ? 0 : DtaTypeByteWidth(raw_types[i]); - } - } else { - for (size_t i = 0; i < columns_.size(); i++) { - columns_[i].type_code = raw_types[i]; - columns_[i].byte_width = DtaTypeByteWidth(raw_types[i]); - } - } -} - -// ─── Variable names ───────────────────────────────────────────────────────── - -void DtaReader::ParseVarnames() { - ReadTag(""); - for (size_t i = 0; i < columns_.size(); i++) { - columns_[i].name = ReadFixedString(params_.varname_len); - } - ReadTag(""); -} - -// ─── Formats ──────────────────────────────────────────────────────────────── - -void DtaReader::ParseFormats() { - ReadTag(""); - for (size_t i = 0; i < columns_.size(); i++) { - columns_[i].format = ReadFixedString(params_.fmt_len); - } - ReadTag(""); -} - -// ─── Value-label names ────────────────────────────────────────────────────── - -void DtaReader::ParseValueLabelNames() { - ReadTag(""); - for (size_t i = 0; i < columns_.size(); i++) { - columns_[i].value_label_name = ReadFixedString(params_.label_name_len); - } - ReadTag(""); -} - -// ─── Variable labels ──────────────────────────────────────────────────────── - -void DtaReader::ParseVariableLabels() { - ReadTag(""); - for (size_t i = 0; i < columns_.size(); i++) { - columns_[i].label = ReadFixedString(params_.var_label_len); - } - ReadTag(""); - - // Now filter out alias variables if needed - if (params_.has_alias_vars) { - std::vector filtered; - for (auto &col : columns_) { - if (col.type_code != 65525) { - filtered.push_back(std::move(col)); - } - } - columns_ = std::move(filtered); - } -} - -// ─── Characteristics (skip) ───────────────────────────────────────────────── - -void DtaReader::SkipCharacteristics() { - // Use the data_offset_ from to skip directly past - // data_offset_ points to "", so we just seek there - fseek(fp_, static_cast(data_offset_), SEEK_SET); -} - -// ─── Data reading ─────────────────────────────────────────────────────────── - -size_t DtaReader::ReadRows(uint64_t start_row, uint32_t count, std::vector &buffer) { - if (start_row >= n_obs_) { - return 0; - } - uint32_t actual = static_cast(std::min(static_cast(count), n_obs_ - start_row)); - - // Seek past tag to the actual data content - // data_offset_ points to the start of "" - // "" is 6 bytes - uint64_t data_content_offset = data_offset_ + 6; - uint64_t byte_offset = data_content_offset + start_row * row_width_; - - fseek(fp_, static_cast(byte_offset), SEEK_SET); - - size_t total_bytes = static_cast(actual) * row_width_; - buffer.resize(total_bytes); - ReadBytes(buffer.data(), total_bytes); - return actual; -} - -// ─── strL support ─────────────────────────────────────────────────────────── - -void DtaReader::LoadStrLs() { - fseek(fp_, static_cast(strls_offset_), SEEK_SET); - ReadTag(""); - - while (true) { - // Try to read "GSO" or "" - char marker[4]; - size_t n = fread(marker, 1, 3, fp_); - if (n < 3) { - break; - } - - if (memcmp(marker, "GSO", 3) != 0) { - // Must be "" — seek back - fseek(fp_, -3, SEEK_CUR); - ReadTag(""); - return; - } - - // GSO format: GSO + v(4 bytes) + o(8 bytes) + t(1 byte) + len(4 bytes) + content - // v and o are encoded per byteorder - uint32_t v; - uint64_t o; - - if (params_.version == 117) { - // v117: v is 4 bytes, o is 4 bytes - v = ReadU32(); - o = ReadU32(); - } else { - // v118+: v is 4 bytes, o is 8 bytes - v = ReadU32(); - o = ReadU64(); - } - - uint8_t t; // type: 129=binary, 130=ASCII/UTF-8 - ReadBytes(&t, 1); - - uint32_t len = ReadU32(); - - std::string content(len, '\0'); - if (len > 0) { - ReadBytes(&content[0], len); - } - // Remove trailing null if present - if (!content.empty() && content.back() == '\0') { - content.pop_back(); - } - - // Pack (v, o) into a single 64-bit key - uint64_t key = (static_cast(v) << 32) | (o & 0xffffffffULL); - strl_table_[key] = std::move(content); - } -} - -const std::string &DtaReader::ResolveStrL(uint32_t v, uint64_t o) const { - uint64_t key = (static_cast(v) << 32) | (o & 0xffffffffULL); - auto it = strl_table_.find(key); - if (it != strl_table_.end()) { - return it->second; - } - return empty_strl_; -} - -// ─── Value labels ─────────────────────────────────────────────────────────── - -void DtaReader::LoadValueLabels() { - fseek(fp_, static_cast(value_labels_offset_), SEEK_SET); - ReadTag(""); - - while (true) { - // Try to read "" or "" - char marker[6]; - size_t n = fread(marker, 1, 5, fp_); - if (n < 5) { - break; - } - - if (memcmp(marker, "", 5) != 0) { - fseek(fp_, -static_cast(n), SEEK_CUR); - ReadTag(""); - return; - } - - // len (4 bytes) — total length of what follows until - uint32_t total_len = ReadU32(); - - // labname (label_name_len bytes, null-terminated) - std::string labname = ReadFixedString(params_.label_name_len); - - // 3 bytes padding - fseek(fp_, 3, SEEK_CUR); - - // value_label_table: n(4), txtlen(4), off[n](4*n), val[n](4*n), txt[txtlen] - uint32_t n_entries = ReadU32(); - uint32_t txtlen = ReadU32(); - - std::vector off(n_entries); - for (uint32_t i = 0; i < n_entries; i++) { - off[i] = ReadU32(); - } - - std::vector val(n_entries); - for (uint32_t i = 0; i < n_entries; i++) { - int32_t v; - ReadBytes(&v, 4); - val[i] = SwapIfNeeded(v); - } - - std::vector txt(txtlen); - if (txtlen > 0) { - ReadBytes(txt.data(), txtlen); - } - - DtaValueLabel vl; - vl.name = labname; - for (uint32_t i = 0; i < n_entries; i++) { - if (off[i] < txtlen) { - std::string label(&txt[off[i]]); - vl.mappings[val[i]] = std::move(label); - } - } - - value_labels_.push_back(std::move(vl)); - - ReadTag(""); - } -} - -} // namespace dta diff --git a/src/dta/dta_reader.hpp b/src/dta/dta_reader.hpp deleted file mode 100644 index 23e9414..0000000 --- a/src/dta/dta_reader.hpp +++ /dev/null @@ -1,118 +0,0 @@ -#pragma once - -#include -#include -#include -#include -#include - -namespace dta { - -// Version-dependent parameters for .dta formats 117-121 -struct DtaVersionParams { - int version; // 117-121 - uint32_t varname_len; // 33 (117) or 129 (118+) - uint32_t sortlist_entry_size; // 2 (117/118/120) or 4 (119/121) - uint32_t fmt_len; // 49 (117) or 57 (118+) - uint32_t label_name_len; // 33 (117) or 129 (118+) - uint32_t var_label_len; // 81 (117) or 321 (118+) - uint32_t k_field_size; // 2 (117/118/120) or 4 (119/121) - uint32_t n_field_size; // 4 (117) or 8 (118+) - uint32_t dataset_label_len_size; // 1 (117) or 2 (118+) - bool has_alias_vars; // true for 120, 121 - - static DtaVersionParams ForVersion(int version); -}; - -struct DtaColumn { - std::string name; - uint16_t type_code; // 1-2045=str, 32768=strL, 65525=alias, 65526-65530=numeric - uint16_t byte_width; // bytes per observation for this variable - std::string format; // display format (e.g. "%td", "%tc", "%9.0g") - std::string value_label_name; - std::string label; // variable label -}; - -struct DtaValueLabel { - std::string name; - std::unordered_map mappings; // val -> text -}; - -// Missing value sentinel boundaries (values >= these are missing) -struct DtaMissing { - static bool IsMissingByte(int8_t val); - static bool IsMissingInt(int16_t val); - static bool IsMissingLong(int32_t val); - static bool IsMissingFloat(float val); - static bool IsMissingDouble(double val); -}; - -class DtaReader { -public: - explicit DtaReader(const std::string &path); - ~DtaReader(); - - // Metadata (available after construction) - int Version() const { return params_.version; } - bool IsMSF() const { return msf_; } - uint64_t NumObs() const { return n_obs_; } - uint32_t NumVars() const { return static_cast(columns_.size()); } - const std::vector &Columns() const { return columns_; } - const std::string &DatasetLabel() const { return dataset_label_; } - - // Byte-swap a value according to the file's byte order - template - T Swap(T val) const { return SwapIfNeeded(val); } - - // Data reading - uint32_t RowWidth() const { return row_width_; } - uint64_t DataOffset() const { return data_offset_; } - size_t ReadRows(uint64_t start_row, uint32_t count, std::vector &buffer); - - // strL support - void LoadStrLs(); - const std::string &ResolveStrL(uint32_t v, uint64_t o) const; - - // Value labels - void LoadValueLabels(); - const std::vector &ValueLabels() const { return value_labels_; } - -private: - FILE *fp_; - DtaVersionParams params_; - bool msf_; - uint64_t n_obs_; - std::string dataset_label_; - std::vector columns_; - uint32_t row_width_; - uint64_t data_offset_; - uint64_t strls_offset_; - uint64_t value_labels_offset_; - std::unordered_map strl_table_; // packed(v,o) -> string - std::vector value_labels_; - static const std::string empty_strl_; - - void ParseHeader(); - void ParseMap(); - void ParseVariableTypes(); - void ParseVarnames(); - void ParseFormats(); - void ParseValueLabelNames(); - void ParseVariableLabels(); - void SkipCharacteristics(); - - // I/O helpers - template - T SwapIfNeeded(T val) const; - void ReadTag(const char *expected); - void ReadBytes(void *buf, size_t n); - uint16_t ReadU16(); - uint32_t ReadU32(); - uint64_t ReadU64(); - std::string ReadFixedString(uint32_t len); -}; - -// Byte width for a given type code -uint16_t DtaTypeByteWidth(uint16_t type_code); - -} // namespace dta diff --git a/src/dta/dta_writer.cpp b/src/dta/dta_writer.cpp deleted file mode 100644 index d9f0bf0..0000000 --- a/src/dta/dta_writer.cpp +++ /dev/null @@ -1,304 +0,0 @@ -#include "dta_writer.hpp" - -#include -#include -#include - -namespace dta { - -// ─── Helpers ──────────────────────────────────────────────────────────────── - -void DtaWriter::WriteTag(const char *tag) { - size_t len = strlen(tag); - if (fwrite(tag, 1, len, fp_) != len) { - throw std::runtime_error("Failed to write to .dta file"); - } -} - -void DtaWriter::WriteFixed(const std::string &s, uint32_t len) { - std::vector buf(len, 0); - size_t copy_len = std::min(static_cast(len - 1), s.size()); - memcpy(buf.data(), s.data(), copy_len); - if (fwrite(buf.data(), 1, len, fp_) != len) { - throw std::runtime_error("Failed to write to .dta file"); - } -} - -void DtaWriter::WriteU16(uint16_t v) { - // Always LSF (native little-endian) - if (fwrite(&v, 2, 1, fp_) != 1) { - throw std::runtime_error("Failed to write to .dta file"); - } -} - -void DtaWriter::WriteU32(uint32_t v) { - if (fwrite(&v, 4, 1, fp_) != 1) { - throw std::runtime_error("Failed to write to .dta file"); - } -} - -void DtaWriter::WriteU64(uint64_t v) { - if (fwrite(&v, 8, 1, fp_) != 1) { - throw std::runtime_error("Failed to write to .dta file"); - } -} - -// ─── Constructor / Destructor ─────────────────────────────────────────────── - -DtaWriter::DtaWriter(const std::string &path, const std::vector &columns, - const std::string &dataset_label) - : fp_(nullptr), columns_(columns), dataset_label_(dataset_label), row_width_(0), map_content_pos_(0) { - memset(offsets_, 0, sizeof(offsets_)); - - row_width_ = 0; - for (auto &col : columns_) { - row_width_ += col.byte_width; - } - - fp_ = fopen(path.c_str(), "wb"); - if (!fp_) { - throw std::runtime_error("Cannot create .dta file: " + path); - } -} - -DtaWriter::~DtaWriter() { - if (fp_) { - fclose(fp_); - fp_ = nullptr; - } -} - -// ─── WriteMetadata ────────────────────────────────────────────────────────── -// Writes everything from through (opening tag only). -// After this, caller writes raw row data, then calls Finalize(). - -void DtaWriter::WriteMetadata() { - // Format 118: K=2 bytes, varname=129, fmt=57, label_name=129, var_label=321, - // N=8 bytes, dataset_label_len=2 bytes, sortlist=2 bytes - - // - offsets_[0] = 0; - WriteTag(""); - - //
- WriteTag("
"); - WriteTag(""); - WriteTag("118"); - WriteTag(""); - WriteTag(""); - WriteTag("LSF"); - WriteTag(""); - WriteTag(""); - WriteU16(static_cast(columns_.size())); - WriteTag(""); - WriteTag(""); - WriteU64(0); // placeholder, rewritten in Finalize - WriteTag(""); - - // Dataset label - WriteTag(""); - - // Timestamp - WriteTag(""); - uint8_t ts_len = 0; - fwrite(&ts_len, 1, 1, fp_); - WriteTag(""); - WriteTag("
"); - - // — placeholder offsets, rewritten in Finalize - offsets_[1] = static_cast(ftell(fp_)); - WriteTag(""); - map_content_pos_ = static_cast(ftell(fp_)); - for (int i = 0; i < 14; i++) { - WriteU64(0); - } - WriteTag(""); - - // - offsets_[2] = static_cast(ftell(fp_)); - WriteTag(""); - for (auto &col : columns_) { - WriteU16(col.type_code); - } - WriteTag(""); - - // (129 bytes each) - offsets_[3] = static_cast(ftell(fp_)); - WriteTag(""); - for (auto &col : columns_) { - WriteFixed(col.name, 129); - } - WriteTag(""); - - // (K+1 entries, 2 bytes each) - offsets_[4] = static_cast(ftell(fp_)); - WriteTag(""); - for (size_t i = 0; i <= columns_.size(); i++) { - WriteU16(0); - } - WriteTag(""); - - // (57 bytes each) - offsets_[5] = static_cast(ftell(fp_)); - WriteTag(""); - for (auto &col : columns_) { - WriteFixed(col.format, 57); - } - WriteTag(""); - - // (129 bytes each) - offsets_[6] = static_cast(ftell(fp_)); - WriteTag(""); - for (auto &col : columns_) { - WriteFixed(col.value_label_name, 129); - } - WriteTag(""); - - // (321 bytes each) - offsets_[7] = static_cast(ftell(fp_)); - WriteTag(""); - for (auto &col : columns_) { - WriteFixed(col.label, 321); - } - WriteTag(""); - - // (empty) - offsets_[8] = static_cast(ftell(fp_)); - WriteTag(""); - WriteTag(""); - - // — opening tag only; caller writes row data - offsets_[9] = static_cast(ftell(fp_)); - WriteTag(""); -} - -// ─── WriteRowData ─────────────────────────────────────────────────────────── - -void DtaWriter::WriteRowData(const char *data, size_t n_bytes) { - if (n_bytes > 0) { - if (fwrite(data, 1, n_bytes, fp_) != n_bytes) { - throw std::runtime_error("Failed to write row data to .dta file"); - } - } -} - -// ─── StrL ─────────────────────────────────────────────────────────────────── - -void DtaWriter::AddStrL(uint32_t v, uint64_t o, const std::string &value) { - strl_entries_.push_back({v, o, value}); -} - -// ─── Value labels ─────────────────────────────────────────────────────────── - -void DtaWriter::AddValueLabel(const WriterValueLabel &vl) { - value_labels_.push_back(vl); -} - -// ─── Finalize ─────────────────────────────────────────────────────────────── - -void DtaWriter::Finalize(uint64_t total_obs) { - // Close - WriteTag(""); - - // - offsets_[10] = static_cast(ftell(fp_)); - WriteTag(""); - for (auto &entry : strl_entries_) { - WriteTag("GSO"); - WriteU32(entry.v); - WriteU64(entry.o); - uint8_t t = 130; // ASCII/UTF-8 - fwrite(&t, 1, 1, fp_); - uint32_t len = static_cast(entry.value.size() + 1); // include null terminator - WriteU32(len); - fwrite(entry.value.data(), 1, entry.value.size(), fp_); - fwrite("\0", 1, 1, fp_); // null terminator - } - WriteTag(""); - - // - offsets_[11] = static_cast(ftell(fp_)); - WriteTag(""); - for (auto &vl : value_labels_) { - // Sort entries by value - std::vector> sorted(vl.mappings.begin(), vl.mappings.end()); - std::sort(sorted.begin(), sorted.end()); - - // Build txt table and off table - std::vector vals; - std::vector offs; - std::string txt_concat; - for (auto &p : sorted) { - vals.push_back(p.first); - offs.push_back(static_cast(txt_concat.size())); - txt_concat += p.second; - txt_concat += '\0'; - } - - uint32_t n = static_cast(sorted.size()); - uint32_t txtlen = static_cast(txt_concat.size()); - // Total length of value_label_table: 4(n) + 4(txtlen) + 4*n(off) + 4*n(val) + txtlen(txt) - uint32_t table_len = 4 + 4 + 4 * n + 4 * n + txtlen; - - WriteTag(""); - // len: total bytes after labname+padding = table_len - // Actually: len covers labname(129) + padding(3) + table - uint32_t total_len = 129 + 3 + table_len; - WriteU32(total_len); - - // labname (129 bytes, null-terminated) - WriteFixed(vl.name, 129); - - // 3 bytes padding - char pad[3] = {0, 0, 0}; - fwrite(pad, 1, 3, fp_); - - // value_label_table - WriteU32(n); - WriteU32(txtlen); - for (auto &o : offs) { - WriteU32(o); - } - for (auto &v : vals) { - uint32_t u; - memcpy(&u, &v, 4); - WriteU32(u); - } - fwrite(txt_concat.data(), 1, txtlen, fp_); - - WriteTag(""); - } - WriteTag(""); - - // - offsets_[12] = static_cast(ftell(fp_)); - WriteTag(""); - offsets_[13] = static_cast(ftell(fp_)); - - // Rewrite in header - // is at a fixed position. Let's find it: - //
118LSFxx - // Tag lengths: =11,
=8, =9, 118=3, =10, - // =11, LSF=3, =12, =3, 2bytes, =4, =3 - // Total to start of N content = 11+8+9+3+10+11+3+12+3+2+4+3 = 79 - uint64_t n_pos = 79; - fseek(fp_, static_cast(n_pos), SEEK_SET); - WriteU64(total_obs); - - // Rewrite offsets - fseek(fp_, static_cast(map_content_pos_), SEEK_SET); - for (int i = 0; i < 14; i++) { - WriteU64(offsets_[i]); - } - - fclose(fp_); - fp_ = nullptr; -} - -} // namespace dta diff --git a/src/dta/dta_writer.hpp b/src/dta/dta_writer.hpp deleted file mode 100644 index af0ba9e..0000000 --- a/src/dta/dta_writer.hpp +++ /dev/null @@ -1,74 +0,0 @@ -#pragma once - -#include -#include -#include -#include -#include - -namespace dta { - -struct DtaWriteColumn { - std::string name; // max 32 UTF-8 chars (129 bytes) - uint16_t type_code; // 1-2045=str, 32768=strL, 65526-65530=numeric - uint16_t byte_width; // bytes per observation - std::string format; // display format (e.g. "%9.0g", "%td") - std::string value_label_name; - std::string label; // variable label -}; - -struct WriterValueLabel { - std::string name; - std::unordered_map mappings; -}; - -struct StrLEntry { - uint32_t v; // 1-based variable index - uint64_t o; // 1-based observation index - std::string value; -}; - -class DtaWriter { -public: - DtaWriter(const std::string &path, const std::vector &columns, - const std::string &dataset_label = ""); - ~DtaWriter(); - - // Call after construction, before writing data rows - void WriteMetadata(); - - // Write raw row data (row_width bytes per row, already packed) - void WriteRowData(const char *data, size_t n_bytes); - - // Add a strL entry (call during sink, before Finalize) - void AddStrL(uint32_t v, uint64_t o, const std::string &value); - - // Add a value label definition (call before Finalize) - void AddValueLabel(const WriterValueLabel &vl); - - // Close the data section and write strls, value labels, map, footer - void Finalize(uint64_t total_obs); - - uint32_t RowWidth() const { return row_width_; } - -private: - FILE *fp_; - std::vector columns_; - std::string dataset_label_; - uint32_t row_width_; - - // File positions for (14 entries) - uint64_t offsets_[14]; - uint64_t map_content_pos_; // file position where the 14 offsets start - - std::vector strl_entries_; - std::vector value_labels_; - - void WriteTag(const char *tag); - void WriteFixed(const std::string &s, uint32_t len); - void WriteU16(uint16_t v); - void WriteU32(uint32_t v); - void WriteU64(uint64_t v); -}; - -} // namespace dta diff --git a/src/dta/read_dta_function.cpp b/src/dta/read_dta_function.cpp deleted file mode 100644 index e93b060..0000000 --- a/src/dta/read_dta_function.cpp +++ /dev/null @@ -1,374 +0,0 @@ -#include "read_dta_function.hpp" -#include "dta_reader.hpp" - -#include "duckdb/common/types/date.hpp" -#include "duckdb/common/types/timestamp.hpp" -#include "duckdb/function/table_function.hpp" -#include "duckdb/main/client_context.hpp" - -#include - -namespace duckdb { - -// ─── Bind data ────────────────────────────────────────────────────────────── - -struct ReadDtaBindData : public TableFunctionData { - string file_path; - shared_ptr reader; - vector return_types; - vector return_names; - bool apply_value_labels; - - // Column index mapping: which reader column each return column maps to - vector reader_col_indices; - - // Value label lookups: for each return column, the value label map (if any) - vector col_value_labels; - // Pre-built: stata value -> enum index for each column with value labels - vector> col_enum_index; -}; - -// ─── Init state ───────────────────────────────────────────────────────────── - -struct ReadDtaGlobalState : public GlobalTableFunctionState { - idx_t current_row; - vector row_buffer; - bool strls_loaded; - - ReadDtaGlobalState() : current_row(0), strls_loaded(false) { - } - - idx_t MaxThreads() const override { - return 1; - } -}; - -// ─── Date/time conversion helpers ─────────────────────────────────────────── - -// Stata epoch: 1960-01-01, DuckDB epoch: 1970-01-01 -// Difference: 3653 days -static constexpr int32_t STATA_EPOCH_OFFSET = 3653; - -// Stata %tc: milliseconds since 1960-01-01 00:00:00 -// DuckDB: microseconds since 1970-01-01 00:00:00 -static constexpr int64_t STATA_TC_EPOCH_OFFSET_MS = 3653LL * 24 * 60 * 60 * 1000; - -static bool IsDateFormat(const string &fmt) { - // %td, %d, %-td, %-d - if (fmt.empty()) return false; - string lower = fmt; - for (auto &c : lower) c = tolower(c); - return lower.find("%td") != string::npos || lower.find("%d") != string::npos; -} - -static bool IsDatetimeFormat(const string &fmt) { - // %tc, %tC - if (fmt.empty()) return false; - return fmt.find("%tc") != string::npos || fmt.find("%tC") != string::npos; -} - -// ─── Type mapping ─────────────────────────────────────────────────────────── - -static LogicalType MapDtaType(const dta::DtaColumn &col, bool apply_value_labels, - const dta::DtaValueLabel *vl) { - uint16_t tc = col.type_code; - - // String types - if (tc >= 1 && tc <= 2045) { - return LogicalType::VARCHAR; - } - if (tc == 32768) { // strL - return LogicalType::VARCHAR; - } - - // Numeric types with value labels -> ENUM - if (apply_value_labels && vl && !vl->mappings.empty()) { - // Build enum from value label - // Sort values to get deterministic order - vector> sorted_pairs(vl->mappings.begin(), vl->mappings.end()); - std::sort(sorted_pairs.begin(), sorted_pairs.end()); - Vector enum_strings(LogicalType::VARCHAR, sorted_pairs.size()); - auto str_data = FlatVector::GetData(enum_strings); - for (idx_t i = 0; i < sorted_pairs.size(); i++) { - str_data[i] = StringVector::AddString(enum_strings, sorted_pairs[i].second); - } - return LogicalType::ENUM(enum_strings, sorted_pairs.size()); - } - - // Numeric types - switch (tc) { - case 65530: return LogicalType::TINYINT; // byte - case 65529: return LogicalType::SMALLINT; // int - case 65528: return LogicalType::INTEGER; // long - case 65527: return LogicalType::FLOAT; // float - case 65526: { // double - if (IsDateFormat(col.format)) { - return LogicalType::DATE; - } - if (IsDatetimeFormat(col.format)) { - return LogicalType::TIMESTAMP; - } - return LogicalType::DOUBLE; - } - default: - return LogicalType::VARCHAR; - } -} - -// ─── Bind ─────────────────────────────────────────────────────────────────── - -static unique_ptr ReadDtaBind(ClientContext &context, TableFunctionBindInput &input, - vector &return_types, vector &names) { - auto result = make_uniq(); - - result->file_path = input.inputs[0].GetValue(); - - // Parse named parameters - result->apply_value_labels = false; - for (auto &kv : input.named_parameters) { - if (kv.first == "value_labels") { - result->apply_value_labels = kv.second.GetValue(); - } - } - - // Open reader - result->reader = make_shared_ptr(result->file_path); - auto &reader = *result->reader; - - // Load value labels if needed - if (result->apply_value_labels) { - reader.LoadValueLabels(); - } - - // Load strLs (we need them during scanning) - // Check if there are any strL columns - bool has_strls = false; - for (auto &col : reader.Columns()) { - if (col.type_code == 32768) { - has_strls = true; - break; - } - } - if (has_strls) { - reader.LoadStrLs(); - } - - // Map columns - auto &cols = reader.Columns(); - for (idx_t i = 0; i < cols.size(); i++) { - // Find value label for this column - const dta::DtaValueLabel *vl = nullptr; - if (result->apply_value_labels && !cols[i].value_label_name.empty()) { - for (auto &label : reader.ValueLabels()) { - if (label.name == cols[i].value_label_name) { - vl = &label; - break; - } - } - } - - LogicalType type = MapDtaType(cols[i], result->apply_value_labels, vl); - return_types.push_back(type); - names.push_back(cols[i].name); - result->reader_col_indices.push_back(i); - result->col_value_labels.push_back(vl); - - // Build enum index: stata_value -> sorted enum position - unordered_map enum_idx; - if (vl && type.id() == LogicalTypeId::ENUM) { - vector> sorted_pairs(vl->mappings.begin(), vl->mappings.end()); - std::sort(sorted_pairs.begin(), sorted_pairs.end()); - for (uint32_t ei = 0; ei < sorted_pairs.size(); ei++) { - enum_idx[sorted_pairs[ei].first] = ei; - } - } - result->col_enum_index.push_back(std::move(enum_idx)); - } - - result->return_types = return_types; - result->return_names = names; - return std::move(result); -} - -// ─── Init ─────────────────────────────────────────────────────────────────── - -static unique_ptr ReadDtaInit(ClientContext &context, TableFunctionInitInput &input) { - return make_uniq(); -} - -// ─── Scan ─────────────────────────────────────────────────────────────────── - -static void ReadDtaScan(ClientContext &context, TableFunctionInput &data, DataChunk &output) { - auto &bind_data = data.bind_data->Cast(); - auto &gstate = data.global_state->Cast(); - auto &reader = *bind_data.reader; - - idx_t remaining = reader.NumObs() - gstate.current_row; - if (remaining == 0) { - output.SetCardinality(0); - return; - } - - uint32_t count = static_cast(std::min(remaining, static_cast(STANDARD_VECTOR_SIZE))); - size_t actual = reader.ReadRows(gstate.current_row, count, gstate.row_buffer); - if (actual == 0) { - output.SetCardinality(0); - return; - } - - auto &cols = reader.Columns(); - uint32_t row_width = reader.RowWidth(); - - // For each output column, extract data from the row buffer - for (idx_t out_col = 0; out_col < output.ColumnCount(); out_col++) { - idx_t reader_col = bind_data.reader_col_indices[out_col]; - auto &col_def = cols[reader_col]; - - // Compute byte offset of this column within a row - uint32_t col_offset = 0; - for (idx_t c = 0; c < reader_col; c++) { - col_offset += cols[c].byte_width; - } - - auto &vec = output.data[out_col]; - auto &type = bind_data.return_types[out_col]; - - for (idx_t row = 0; row < actual; row++) { - const char *row_ptr = gstate.row_buffer.data() + row * row_width + col_offset; - - // Helper: write an enum value from a Stata integer - auto write_enum = [&](int32_t stata_val) { - auto &eidx = bind_data.col_enum_index[out_col]; - auto eit = eidx.find(stata_val); - if (eit != eidx.end()) { - auto enum_size = EnumType::GetSize(type); - if (enum_size <= 256) { - FlatVector::GetData(vec)[row] = static_cast(eit->second); - } else if (enum_size <= 65536) { - FlatVector::GetData(vec)[row] = static_cast(eit->second); - } else { - FlatVector::GetData(vec)[row] = eit->second; - } - } else { - FlatVector::SetNull(vec, row, true); - } - }; - - bool is_enum = (type.id() == LogicalTypeId::ENUM); - - switch (col_def.type_code) { - case 65530: { // byte - int8_t val; - memcpy(&val, row_ptr, 1); - if (dta::DtaMissing::IsMissingByte(val)) { - FlatVector::SetNull(vec, row, true); - } else if (is_enum) { - write_enum(val); - } else { - FlatVector::GetData(vec)[row] = val; - } - break; - } - case 65529: { // int (2-byte) - int16_t val; - memcpy(&val, row_ptr, 2); - val = reader.Swap(val); - if (dta::DtaMissing::IsMissingInt(val)) { - FlatVector::SetNull(vec, row, true); - } else if (is_enum) { - write_enum(val); - } else { - FlatVector::GetData(vec)[row] = val; - } - break; - } - case 65528: { // long (4-byte) - int32_t val; - memcpy(&val, row_ptr, 4); - val = reader.Swap(val); - if (dta::DtaMissing::IsMissingLong(val)) { - FlatVector::SetNull(vec, row, true); - } else if (is_enum) { - write_enum(val); - } else { - FlatVector::GetData(vec)[row] = val; - } - break; - } - case 65527: { // float - float val; - memcpy(&val, row_ptr, 4); - val = reader.Swap(val); - if (dta::DtaMissing::IsMissingFloat(val)) { - FlatVector::SetNull(vec, row, true); - } else if (is_enum) { - write_enum(static_cast(val)); - } else { - FlatVector::GetData(vec)[row] = val; - } - break; - } - case 65526: { // double - double val; - memcpy(&val, row_ptr, 8); - val = reader.Swap(val); - if (dta::DtaMissing::IsMissingDouble(val)) { - FlatVector::SetNull(vec, row, true); - } else if (is_enum) { - write_enum(static_cast(val)); - } else if (type.id() == LogicalTypeId::DATE) { - int32_t stata_days = static_cast(val); - FlatVector::GetData(vec)[row] = date_t(stata_days - STATA_EPOCH_OFFSET); - } else if (type.id() == LogicalTypeId::TIMESTAMP) { - int64_t stata_ms = static_cast(val); - int64_t unix_ms = stata_ms - STATA_TC_EPOCH_OFFSET_MS; - FlatVector::GetData(vec)[row] = timestamp_t(unix_ms * 1000); - } else { - FlatVector::GetData(vec)[row] = val; - } - break; - } - case 32768: { // strL - // Read (v, o) reference from 8 bytes - uint32_t v_ref, o_ref_lo; - memcpy(&v_ref, row_ptr, 4); - memcpy(&o_ref_lo, row_ptr + 4, 4); - v_ref = reader.Swap(v_ref); - o_ref_lo = reader.Swap(o_ref_lo); - if (v_ref == 0 && o_ref_lo == 0) { - FlatVector::SetNull(vec, row, true); - } else { - auto &str = reader.ResolveStrL(v_ref, o_ref_lo); - FlatVector::GetData(vec)[row] = StringVector::AddString(vec, str); - } - break; - } - default: { - // str1-str2045 - if (col_def.type_code >= 1 && col_def.type_code <= 2045) { - // Find null terminator or use full width - size_t len = strnlen(row_ptr, col_def.byte_width); - FlatVector::GetData(vec)[row] = StringVector::AddString(vec, row_ptr, len); - } else { - FlatVector::SetNull(vec, row, true); - } - break; - } - } - } - } - - output.SetCardinality(actual); - gstate.current_row += actual; -} - -// ─── Register ─────────────────────────────────────────────────────────────── - -TableFunction GetReadDtaFunction() { - TableFunction func("read_dta", {LogicalType::VARCHAR}, ReadDtaScan, ReadDtaBind, ReadDtaInit); - func.named_parameters["value_labels"] = LogicalType::BOOLEAN; - func.projection_pushdown = false; // TODO: implement projection pushdown - return func; -} - -} // namespace duckdb diff --git a/src/dta/read_dta_function.hpp b/src/dta/read_dta_function.hpp deleted file mode 100644 index dc85970..0000000 --- a/src/dta/read_dta_function.hpp +++ /dev/null @@ -1,10 +0,0 @@ -#pragma once - -#include "duckdb.hpp" -#include "duckdb/function/table_function.hpp" - -namespace duckdb { - -TableFunction GetReadDtaFunction(); - -} // namespace duckdb diff --git a/src/dta/write_dta_function.cpp b/src/dta/write_dta_function.cpp deleted file mode 100644 index a24e301..0000000 --- a/src/dta/write_dta_function.cpp +++ /dev/null @@ -1,446 +0,0 @@ -#include "write_dta_function.hpp" -#include "dta_writer.hpp" - -#include "duckdb/common/types/date.hpp" -#include "duckdb/common/types/timestamp.hpp" -#include "duckdb/function/copy_function.hpp" -#include "duckdb/main/client_context.hpp" -#include "duckdb/catalog/catalog.hpp" -#include "duckdb/catalog/catalog_entry/table_catalog_entry.hpp" - -#include -#include - -namespace duckdb { - -// ─── Bind data ────────────────────────────────────────────────────────────── - -struct WriteDtaBindData : public FunctionData { - vector columns; - string dataset_label; - // ENUM columns: column index -> value label - vector value_labels; - - unique_ptr Copy() const override { - auto copy = make_uniq(); - copy->columns = columns; - copy->dataset_label = dataset_label; - copy->value_labels = value_labels; - return std::move(copy); - } - - bool Equals(const FunctionData &other_p) const override { - auto &other = other_p.Cast(); - return columns.size() == other.columns.size() && dataset_label == other.dataset_label; - } -}; - -// ─── Global state ─────────────────────────────────────────────────────────── - -struct WriteDtaGlobalState : public GlobalFunctionData { - unique_ptr writer; - mutex lock; - uint64_t rows_written; - - WriteDtaGlobalState() : rows_written(0) { - } -}; - -// ─── Local state ──────────────────────────────────────────────────────────── - -struct WriteDtaLocalState : public LocalFunctionData {}; - -// ─── Constants ────────────────────────────────────────────────────────────── - -static constexpr int32_t STATA_EPOCH_OFFSET = 3653; -static constexpr int64_t STATA_TC_EPOCH_OFFSET_MS = 3653LL * 24 * 60 * 60 * 1000; - -// Missing value sentinels (LSF byte order) -static constexpr int8_t MISSING_BYTE = 101; // 0x65 -static constexpr int16_t MISSING_INT = 32741; // 0x7fe5 -static constexpr int32_t MISSING_LONG = 2147483621; // 0x7fffffe5 - -static double MissingDouble() { - uint64_t bits = 0x7fe0000000000000ULL; - double val; - memcpy(&val, &bits, 8); - return val; -} - -static float MissingFloat() { - uint32_t bits = 0x7f000000U; - float val; - memcpy(&val, &bits, 4); - return val; -} - -// ─── Type mapping ─────────────────────────────────────────────────────────── - -static dta::DtaWriteColumn MapDuckDBType(const string &name, const LogicalType &type) { - dta::DtaWriteColumn col; - col.name = name; - - switch (type.id()) { - case LogicalTypeId::BOOLEAN: - case LogicalTypeId::TINYINT: - case LogicalTypeId::UTINYINT: - col.type_code = 65530; // byte - col.byte_width = 1; - col.format = "%8.0g"; - break; - case LogicalTypeId::SMALLINT: - case LogicalTypeId::USMALLINT: - col.type_code = 65529; // int (2-byte) - col.byte_width = 2; - col.format = "%8.0g"; - break; - case LogicalTypeId::INTEGER: - case LogicalTypeId::UINTEGER: - col.type_code = 65528; // long (4-byte) - col.byte_width = 4; - col.format = "%12.0g"; - break; - case LogicalTypeId::BIGINT: - case LogicalTypeId::UBIGINT: - case LogicalTypeId::HUGEINT: - col.type_code = 65526; // double - col.byte_width = 8; - col.format = "%10.0g"; - break; - case LogicalTypeId::FLOAT: - col.type_code = 65527; // float - col.byte_width = 4; - col.format = "%9.0g"; - break; - case LogicalTypeId::DOUBLE: - case LogicalTypeId::DECIMAL: - col.type_code = 65526; // double - col.byte_width = 8; - col.format = "%10.0g"; - break; - case LogicalTypeId::DATE: - col.type_code = 65526; // double - col.byte_width = 8; - col.format = "%td"; - break; - case LogicalTypeId::TIMESTAMP: - case LogicalTypeId::TIMESTAMP_TZ: - col.type_code = 65526; // double - col.byte_width = 8; - col.format = "%tc"; - break; - case LogicalTypeId::VARCHAR: - col.type_code = 32768; // strL - col.byte_width = 8; // (v, o) reference - col.format = "%9s"; - break; - case LogicalTypeId::ENUM: { - // Map ENUM to the smallest integer type that fits - auto enum_size = EnumType::GetSize(type); - if (enum_size <= 100) { - col.type_code = 65530; // byte - col.byte_width = 1; - } else if (enum_size <= 32740) { - col.type_code = 65529; // int - col.byte_width = 2; - } else { - col.type_code = 65528; // long - col.byte_width = 4; - } - col.format = "%8.0g"; - break; - } - default: - // Fall back to strL for everything else - col.type_code = 32768; - col.byte_width = 8; - col.format = "%9s"; - break; - } - - return col; -} - -// ─── Bind ─────────────────────────────────────────────────────────────────── - -// Try to find column comments from known tables -static unordered_map LookupColumnComments(ClientContext &context, const vector &names) { - unordered_map comments; - try { - // Try dodo._current first (most common case for save command) - auto &catalog = Catalog::GetCatalog(context, INVALID_CATALOG); - auto entry = catalog.GetEntry(context, CatalogType::TABLE_ENTRY, "dodo", "_current", - OnEntryNotFound::RETURN_NULL); - if (entry) { - auto &table_entry = entry->Cast(); - for (auto &col_name : names) { - string n = col_name; - if (table_entry.ColumnExists(n)) { - const auto &col_def = table_entry.GetColumn(table_entry.GetColumnIndex(n)); - if (!col_def.Comment().IsNull()) { - comments[col_name] = col_def.Comment().ToString(); - } - } - } - } - } catch (...) { - // Table doesn't exist or no comments — ignore - } - return comments; -} - -static unique_ptr WriteDtaBind(ClientContext &context, CopyFunctionBindInput &input, - const vector &names, const vector &sql_types) { - auto result = make_uniq(); - - // Look up column comments for variable labels - auto comments = LookupColumnComments(context, names); - - for (idx_t i = 0; i < names.size(); i++) { - auto col = MapDuckDBType(names[i], sql_types[i]); - - // Set variable label from column comment - auto it = comments.find(names[i]); - if (it != comments.end()) { - col.label = it->second; - } - - // Handle ENUM: create value label - if (sql_types[i].id() == LogicalTypeId::ENUM) { - string label_name = names[i]; - col.value_label_name = label_name; - - dta::WriterValueLabel vl; - vl.name = label_name; - auto enum_size = EnumType::GetSize(sql_types[i]); - for (idx_t ei = 0; ei < enum_size; ei++) { - auto str = EnumType::GetString(sql_types[i], ei); - vl.mappings[static_cast(ei)] = str.GetString(); - } - result->value_labels.push_back(std::move(vl)); - } - - result->columns.push_back(std::move(col)); - } - - return std::move(result); -} - -// ─── Init global ──────────────────────────────────────────────────────────── - -static unique_ptr WriteDtaInitGlobal(ClientContext &context, FunctionData &bind_data, - const string &file_path) { - auto &bdata = bind_data.Cast(); - auto gstate = make_uniq(); - - gstate->writer = make_uniq(file_path, bdata.columns, bdata.dataset_label); - gstate->writer->WriteMetadata(); - - // Register value labels - for (auto &vl : bdata.value_labels) { - gstate->writer->AddValueLabel(vl); - } - - return std::move(gstate); -} - -// ─── Init local ───────────────────────────────────────────────────────────── - -static unique_ptr WriteDtaInitLocal(ExecutionContext &context, FunctionData &bind_data) { - return make_uniq(); -} - -// ─── Sink ─────────────────────────────────────────────────────────────────── - -static void WriteDtaSink(ExecutionContext &context, FunctionData &bind_data, GlobalFunctionData &gstate_p, - LocalFunctionData &lstate, DataChunk &input) { - auto &bdata = bind_data.Cast(); - auto &gstate = gstate_p.Cast(); - lock_guard guard(gstate.lock); - - auto &writer = *gstate.writer; - uint32_t row_width = writer.RowWidth(); - idx_t count = input.size(); - - // Flatten all vectors for direct access - input.Flatten(); - - // Row buffer for transposition - vector row_buf(row_width, 0); - - for (idx_t row = 0; row < count; row++) { - memset(row_buf.data(), 0, row_width); - uint32_t offset = 0; - uint64_t obs_idx = gstate.rows_written + row + 1; // 1-based for strL - - for (idx_t col = 0; col < input.ColumnCount(); col++) { - auto &vec = input.data[col]; - auto &col_def = bdata.columns[col]; - char *dest = row_buf.data() + offset; - - if (FlatVector::IsNull(vec, row)) { - // Write missing value sentinel - switch (col_def.type_code) { - case 65530: { int8_t mv = MISSING_BYTE; memcpy(dest, &mv, 1); break; } - case 65529: { int16_t mv = MISSING_INT; memcpy(dest, &mv, 2); break; } - case 65528: { int32_t mv = MISSING_LONG; memcpy(dest, &mv, 4); break; } - case 65527: { float mv = MissingFloat(); memcpy(dest, &mv, 4); break; } - case 65526: { double mv = MissingDouble(); memcpy(dest, &mv, 8); break; } - case 32768: { memset(dest, 0, 8); break; } // strL: (0,0) = NULL - default: break; - } - } else { - auto val_type = vec.GetType().id(); - - switch (col_def.type_code) { - case 65530: { // byte - int8_t val; - if (val_type == LogicalTypeId::BOOLEAN) { - val = FlatVector::GetData(vec)[row] ? 1 : 0; - } else if (val_type == LogicalTypeId::ENUM) { - auto enum_size = EnumType::GetSize(vec.GetType()); - if (enum_size <= 256) { - val = static_cast(FlatVector::GetData(vec)[row]); - } else { - val = static_cast(FlatVector::GetData(vec)[row]); - } - } else if (val_type == LogicalTypeId::UTINYINT) { - val = static_cast(FlatVector::GetData(vec)[row]); - } else if (val_type == LogicalTypeId::TINYINT) { - val = FlatVector::GetData(vec)[row]; - } else { - val = static_cast(vec.GetValue(row).CastAs(context.client, LogicalType::TINYINT).GetValue()); - } - memcpy(dest, &val, 1); - break; - } - case 65529: { // int (2-byte) - int16_t val; - if (val_type == LogicalTypeId::ENUM) { - val = static_cast(FlatVector::GetData(vec)[row]); - } else if (val_type == LogicalTypeId::USMALLINT) { - val = static_cast(FlatVector::GetData(vec)[row]); - } else if (val_type == LogicalTypeId::SMALLINT) { - val = FlatVector::GetData(vec)[row]; - } else { - val = vec.GetValue(row).CastAs(context.client, LogicalType::SMALLINT).GetValue(); - } - memcpy(dest, &val, 2); - break; - } - case 65528: { // long (4-byte) - int32_t val; - if (val_type == LogicalTypeId::ENUM) { - val = static_cast(FlatVector::GetData(vec)[row]); - } else if (val_type == LogicalTypeId::INTEGER) { - val = FlatVector::GetData(vec)[row]; - } else if (val_type == LogicalTypeId::UINTEGER) { - val = static_cast(FlatVector::GetData(vec)[row]); - } else { - val = vec.GetValue(row).CastAs(context.client, LogicalType::INTEGER).GetValue(); - } - memcpy(dest, &val, 4); - break; - } - case 65527: { // float - float val; - if (val_type == LogicalTypeId::FLOAT) { - val = FlatVector::GetData(vec)[row]; - } else { - val = vec.GetValue(row).CastAs(context.client, LogicalType::FLOAT).GetValue(); - } - memcpy(dest, &val, 4); - break; - } - case 65526: { // double - double val; - if (val_type == LogicalTypeId::DATE) { - auto date = FlatVector::GetData(vec)[row]; - val = static_cast(date.days + STATA_EPOCH_OFFSET); - } else if (val_type == LogicalTypeId::TIMESTAMP || val_type == LogicalTypeId::TIMESTAMP_TZ) { - auto ts = FlatVector::GetData(vec)[row]; - int64_t unix_us = ts.value; - int64_t stata_ms = unix_us / 1000 + STATA_TC_EPOCH_OFFSET_MS; - val = static_cast(stata_ms); - } else if (val_type == LogicalTypeId::DOUBLE) { - val = FlatVector::GetData(vec)[row]; - } else if (val_type == LogicalTypeId::FLOAT) { - val = static_cast(FlatVector::GetData(vec)[row]); - } else if (val_type == LogicalTypeId::BIGINT) { - val = static_cast(FlatVector::GetData(vec)[row]); - } else if (val_type == LogicalTypeId::UBIGINT) { - val = static_cast(FlatVector::GetData(vec)[row]); - } else if (val_type == LogicalTypeId::INTEGER) { - val = static_cast(FlatVector::GetData(vec)[row]); - } else if (val_type == LogicalTypeId::SMALLINT) { - val = static_cast(FlatVector::GetData(vec)[row]); - } else if (val_type == LogicalTypeId::TINYINT) { - val = static_cast(FlatVector::GetData(vec)[row]); - } else { - // Fallback: use Value conversion - val = vec.GetValue(row).CastAs(context.client, LogicalType::DOUBLE).GetValue(); - } - memcpy(dest, &val, 8); - break; - } - case 32768: { // strL - string s; - if (val_type == LogicalTypeId::VARCHAR) { - s = FlatVector::GetData(vec)[row].GetString(); - } else { - s = vec.GetValue(row).ToString(); - } - if (s.empty()) { - memset(dest, 0, 8); - } else { - // v = 1-based column index, o = 1-based observation index - uint32_t v_ref = static_cast(col + 1); - uint32_t o_ref = static_cast(obs_idx); - memcpy(dest, &v_ref, 4); - memcpy(dest + 4, &o_ref, 4); - writer.AddStrL(v_ref, obs_idx, s); - } - break; - } - default: - break; - } - } - offset += col_def.byte_width; - } - - writer.WriteRowData(row_buf.data(), row_width); - } - - gstate.rows_written += count; -} - -// ─── Combine ──────────────────────────────────────────────────────────────── - -static void WriteDtaCombine(ExecutionContext &context, FunctionData &bind_data, GlobalFunctionData &gstate, - LocalFunctionData &lstate) { - // no-op -} - -// ─── Finalize ─────────────────────────────────────────────────────────────── - -static void WriteDtaFinalize(ClientContext &context, FunctionData &bind_data, GlobalFunctionData &gstate_p) { - auto &gstate = gstate_p.Cast(); - gstate.writer->Finalize(gstate.rows_written); -} - -// ─── Register ─────────────────────────────────────────────────────────────── - -CopyFunction GetDtaCopyFunction() { - CopyFunction func("dta"); - func.copy_to_bind = WriteDtaBind; - func.copy_to_initialize_global = WriteDtaInitGlobal; - func.copy_to_initialize_local = WriteDtaInitLocal; - func.copy_to_sink = WriteDtaSink; - func.copy_to_combine = WriteDtaCombine; - func.copy_to_finalize = WriteDtaFinalize; - func.extension = "dta"; - return func; -} - -} // namespace duckdb diff --git a/src/dta/write_dta_function.hpp b/src/dta/write_dta_function.hpp deleted file mode 100644 index 2ae6350..0000000 --- a/src/dta/write_dta_function.hpp +++ /dev/null @@ -1,10 +0,0 @@ -#pragma once - -#include "duckdb.hpp" -#include "duckdb/function/copy_function.hpp" - -namespace duckdb { - -CopyFunction GetDtaCopyFunction(); - -} // namespace duckdb diff --git a/src/extension/dodo_extension.cpp b/src/extension/dodo_extension.cpp index ed84753..51affde 100644 --- a/src/extension/dodo_extension.cpp +++ b/src/extension/dodo_extension.cpp @@ -1,8 +1,6 @@ #define DUCKDB_EXTENSION_MAIN #include "dodo_extension.hpp" -#include "read_dta_function.hpp" -#include "write_dta_function.hpp" #include "duckdb.hpp" #include "duckdb/common/exception.hpp" #include "duckdb/common/string_util.hpp" @@ -263,10 +261,6 @@ BoundStatement dodo_bind(ClientContext &context, Binder &binder, OperatorExtensi //===--------------------------------------------------------------------===// static void LoadInternal(ExtensionLoader &loader) { - // Register read_dta table function and dta copy function - loader.RegisterFunction(GetReadDtaFunction()); - loader.RegisterFunction(GetDtaCopyFunction()); - auto &instance = loader.GetDatabaseInstance(); auto &config = DBConfig::GetConfig(instance); diff --git a/test/sql/dodo_dta.test b/test/sql/dodo_dta.test new file mode 100644 index 0000000..e094a4a --- /dev/null +++ b/test/sql/dodo_dta.test @@ -0,0 +1,96 @@ +# name: test/sql/dodo_dta.test +# description: test dodo commands with .dta files +# group: [sql] + +require dodo + +require dta + +# ============================================================ +# Variable labels read from .dta +# ============================================================ + +# auto.dta has variable labels set by the Python generator +statement ok +use "test/data/auto.dta", clear + +# Verify labels were read +query IIII +codebook +---- +id INTEGER Firm identifier (empty) +name VARCHAR Firm name (empty) +revenue DOUBLE Annual revenue (empty) +year SMALLINT Fiscal year (empty) + +# label variable overrides +statement ok +label variable id "Unique firm ID" + +query IIII +codebook +---- +id INTEGER Unique firm ID (empty) +name VARCHAR Firm name (empty) +revenue DOUBLE Annual revenue (empty) +year SMALLINT Fiscal year (empty) + +statement ok +clear + +# ============================================================ +# save command with .dta extension +# ============================================================ + +statement ok +use "test/data/auto.dta", clear + +statement ok +keep if year == 2018 + +statement ok +save "__TEST_DIR__/save_test.dta" + +statement ok +clear + +query I +SELECT count(*) FROM read_dta('__TEST_DIR__/save_test.dta') +---- +2 + +query IIR +SELECT id, name, revenue FROM read_dta('__TEST_DIR__/save_test.dta') ORDER BY id +---- +2 Beta Corp 2500.5 +3 Gamma NULL + +# ============================================================ +# use .dta then transform +# ============================================================ + +statement ok +use "test/data/auto.dta", clear + +query I +count +---- +3 + +query IIRI +list +---- +1 Acme 1000.0 2017 +2 Beta Corp 2500.5 2018 +3 Gamma NULL 2018 + +statement ok +keep if year == 2018 + +query I +count +---- +2 + +statement ok +clear diff --git a/test/sql/read_dta.test b/test/sql/read_dta.test deleted file mode 100644 index ad8a61e..0000000 --- a/test/sql/read_dta.test +++ /dev/null @@ -1,122 +0,0 @@ -# name: test/sql/read_dta.test -# description: test read_dta table function -# group: [sql] - -require dodo - -# ============================================================ -# Basic read: column types, names, values -# ============================================================ - -query IIRI -SELECT * FROM read_dta('test/data/auto.dta') ORDER BY id ----- -1 Acme 1000.0 2017 -2 Beta Corp 2500.5 2018 -3 Gamma NULL 2018 - -# ============================================================ -# Column count and row count -# ============================================================ - -query I -SELECT count(*) FROM read_dta('test/data/auto.dta') ----- -3 - -# ============================================================ -# Type inference: int32 (long), varchar (str10), double, int16 (int) -# ============================================================ - -query IIII -SELECT typeof(id), typeof(name), typeof(revenue), typeof(year) -FROM read_dta('test/data/auto.dta') LIMIT 1 ----- -INTEGER VARCHAR DOUBLE SMALLINT - -# ============================================================ -# Missing value handling (double NULL sentinel) -# ============================================================ - -query I -SELECT count(*) FROM read_dta('test/data/auto.dta') WHERE revenue IS NULL ----- -1 - -query I -SELECT id FROM read_dta('test/data/auto.dta') WHERE revenue IS NULL ----- -3 - -# ============================================================ -# Filtering pushdown — WHERE works -# ============================================================ - -query IR -SELECT id, revenue FROM read_dta('test/data/auto.dta') WHERE year = 2018 ORDER BY id ----- -2 2500.5 -3 NULL - -# ============================================================ -# Aggregation over .dta file -# ============================================================ - -query R -SELECT avg(revenue) FROM read_dta('test/data/auto.dta') ----- -1750.25 - -# ============================================================ -# String trimming (str10 should not have trailing nulls) -# ============================================================ - -query I -SELECT length(name) FROM read_dta('test/data/auto.dta') WHERE id = 1 ----- -4 - -query I -SELECT length(name) FROM read_dta('test/data/auto.dta') WHERE id = 2 ----- -9 - -# ============================================================ -# Integration with dodo commands: use .dta file -# ============================================================ - -statement ok -use "test/data/auto.dta", clear - -query I -count ----- -3 - -query IIRI -list ----- -1 Acme 1000.0 2017 -2 Beta Corp 2500.5 2018 -3 Gamma NULL 2018 - -# Transformations on .dta data -statement ok -keep if year == 2018 - -query I -count ----- -2 - -statement ok -clear - -# ============================================================ -# Error handling: nonexistent file -# ============================================================ - -statement error -SELECT * FROM read_dta('nonexistent.dta') ----- -Cannot open .dta file diff --git a/test/sql/write_dta.test b/test/sql/write_dta.test deleted file mode 100644 index 6469879..0000000 --- a/test/sql/write_dta.test +++ /dev/null @@ -1,471 +0,0 @@ -# name: test/sql/write_dta.test -# description: test write_dta (COPY TO FORMAT dta) and round-trip -# group: [sql] - -require dodo - -# ============================================================ -# Basic round-trip: COPY TO then read_dta -# ============================================================ - -statement ok -COPY (SELECT 1::INTEGER AS id, 'hello' AS name, 42.5 AS value) TO '__TEST_DIR__/basic.dta' (FORMAT dta) - -query IIR -SELECT * FROM read_dta('__TEST_DIR__/basic.dta') ----- -1 hello 42.5 - -# ============================================================ -# Type mapping: INTEGER -> long (int32) -# ============================================================ - -statement ok -COPY (SELECT 1::INTEGER AS x, -100::INTEGER AS y, 2147483620::INTEGER AS z) TO '__TEST_DIR__/int.dta' (FORMAT dta) - -query III -SELECT typeof(x), typeof(y), typeof(z) FROM read_dta('__TEST_DIR__/int.dta') ----- -INTEGER INTEGER INTEGER - -query III -SELECT * FROM read_dta('__TEST_DIR__/int.dta') ----- -1 -100 2147483620 - -# ============================================================ -# Type mapping: TINYINT -> byte -# ============================================================ - -statement ok -COPY (SELECT 1::TINYINT AS a, 100::TINYINT AS b) TO '__TEST_DIR__/tinyint.dta' (FORMAT dta) - -query II -SELECT typeof(a), typeof(b) FROM read_dta('__TEST_DIR__/tinyint.dta') ----- -TINYINT TINYINT - -query II -SELECT * FROM read_dta('__TEST_DIR__/tinyint.dta') ----- -1 100 - -# ============================================================ -# Type mapping: SMALLINT -> int (int16) -# ============================================================ - -statement ok -COPY (SELECT 1::SMALLINT AS a, 32000::SMALLINT AS b) TO '__TEST_DIR__/smallint.dta' (FORMAT dta) - -query II -SELECT typeof(a), typeof(b) FROM read_dta('__TEST_DIR__/smallint.dta') ----- -SMALLINT SMALLINT - -query II -SELECT * FROM read_dta('__TEST_DIR__/smallint.dta') ----- -1 32000 - -# ============================================================ -# Type mapping: BIGINT -> double -# ============================================================ - -statement ok -COPY (SELECT 1000000000000::BIGINT AS big) TO '__TEST_DIR__/bigint.dta' (FORMAT dta) - -query I -SELECT typeof(big) FROM read_dta('__TEST_DIR__/bigint.dta') ----- -DOUBLE - -query R -SELECT * FROM read_dta('__TEST_DIR__/bigint.dta') ----- -1000000000000.0 - -# ============================================================ -# Type mapping: FLOAT -> float -# ============================================================ - -statement ok -COPY (SELECT 3.14::FLOAT AS pi) TO '__TEST_DIR__/float.dta' (FORMAT dta) - -query I -SELECT typeof(pi) FROM read_dta('__TEST_DIR__/float.dta') ----- -FLOAT - -# ============================================================ -# Type mapping: DOUBLE -> double -# ============================================================ - -statement ok -COPY (SELECT 2.718281828::DOUBLE AS e) TO '__TEST_DIR__/double.dta' (FORMAT dta) - -query R -SELECT * FROM read_dta('__TEST_DIR__/double.dta') ----- -2.718281828 - -# ============================================================ -# Type mapping: BOOLEAN -> byte -# ============================================================ - -statement ok -COPY (SELECT true AS a, false AS b) TO '__TEST_DIR__/bool.dta' (FORMAT dta) - -query II -SELECT * FROM read_dta('__TEST_DIR__/bool.dta') ----- -1 0 - -# ============================================================ -# Type mapping: DATE -> double with %td format -# ============================================================ - -statement ok -COPY (SELECT DATE '2024-01-15' AS dt) TO '__TEST_DIR__/date.dta' (FORMAT dta) - -query I -SELECT typeof(dt) FROM read_dta('__TEST_DIR__/date.dta') ----- -DATE - -query T -SELECT * FROM read_dta('__TEST_DIR__/date.dta') ----- -2024-01-15 - -# ============================================================ -# Type mapping: TIMESTAMP -> double with %tc format -# ============================================================ - -statement ok -COPY (SELECT TIMESTAMP '2024-01-15 10:30:00' AS ts) TO '__TEST_DIR__/timestamp.dta' (FORMAT dta) - -query I -SELECT typeof(ts) FROM read_dta('__TEST_DIR__/timestamp.dta') ----- -TIMESTAMP - -query T -SELECT * FROM read_dta('__TEST_DIR__/timestamp.dta') ----- -2024-01-15 10:30:00 - -# ============================================================ -# Type mapping: VARCHAR -> strL -# ============================================================ - -statement ok -COPY (SELECT 'short' AS a, 'a longer string value here' AS b) TO '__TEST_DIR__/varchar.dta' (FORMAT dta) - -query II -SELECT * FROM read_dta('__TEST_DIR__/varchar.dta') ----- -short a longer string value here - -# ============================================================ -# NULL handling for all types -# ============================================================ - -statement ok -COPY (SELECT NULL::INTEGER AS a, NULL::DOUBLE AS b, NULL::VARCHAR AS c, NULL::DATE AS d, NULL::TINYINT AS e) TO '__TEST_DIR__/nulls.dta' (FORMAT dta) - -query IIIII -SELECT a IS NULL, b IS NULL, c IS NULL, d IS NULL, e IS NULL FROM read_dta('__TEST_DIR__/nulls.dta') ----- -true true true true true - -# ============================================================ -# Multiple rows -# ============================================================ - -statement ok -COPY (SELECT * FROM (VALUES (1, 'a', 10.0), (2, 'b', 20.0), (3, 'c', 30.0)) AS t(id, name, value)) TO '__TEST_DIR__/multi.dta' (FORMAT dta) - -query I -SELECT count(*) FROM read_dta('__TEST_DIR__/multi.dta') ----- -3 - -query IIR -SELECT * FROM read_dta('__TEST_DIR__/multi.dta') ORDER BY id ----- -1 a 10.0 -2 b 20.0 -3 c 30.0 - -# ============================================================ -# Round-trip: read existing .dta, write, read back -# ============================================================ - -statement ok -COPY (SELECT * FROM read_dta('test/data/auto.dta')) TO '__TEST_DIR__/roundtrip.dta' (FORMAT dta) - -query IIRI -SELECT * FROM read_dta('__TEST_DIR__/roundtrip.dta') ORDER BY id ----- -1 Acme 1000.0 2017 -2 Beta Corp 2500.5 2018 -3 Gamma NULL 2018 - -# ============================================================ -# Boundary values: TINYINT (byte) range -127..100 -# ============================================================ - -statement ok -COPY (SELECT (-127)::TINYINT AS mn, 100::TINYINT AS mx) TO '__TEST_DIR__/byte_bounds.dta' (FORMAT dta) - -query II -SELECT * FROM read_dta('__TEST_DIR__/byte_bounds.dta') ----- --127 100 - -# Values above 100 are missing sentinels in Stata -statement ok -COPY (SELECT 101::TINYINT AS v) TO '__TEST_DIR__/byte_missing.dta' (FORMAT dta) - -query I -SELECT v IS NULL FROM read_dta('__TEST_DIR__/byte_missing.dta') ----- -true - -# ============================================================ -# Boundary values: SMALLINT (int) range -32767..32740 -# ============================================================ - -statement ok -COPY (SELECT (-32767)::SMALLINT AS mn, 32740::SMALLINT AS mx) TO '__TEST_DIR__/int_bounds.dta' (FORMAT dta) - -query II -SELECT * FROM read_dta('__TEST_DIR__/int_bounds.dta') ----- --32767 32740 - -# ============================================================ -# Boundary values: INTEGER (long) range -2147483647..2147483620 -# ============================================================ - -statement ok -COPY (SELECT (-2147483647)::INTEGER AS mn, 2147483620::INTEGER AS mx) TO '__TEST_DIR__/long_bounds.dta' (FORMAT dta) - -query II -SELECT * FROM read_dta('__TEST_DIR__/long_bounds.dta') ----- --2147483647 2147483620 - -# ============================================================ -# Unsigned integers: cast to larger signed type -# ============================================================ - -# UTINYINT -> byte (fits since max utinyint=255 > 100 means some become missing) -statement ok -COPY (SELECT 50::UTINYINT AS a) TO '__TEST_DIR__/utinyint.dta' (FORMAT dta) - -query I -SELECT * FROM read_dta('__TEST_DIR__/utinyint.dta') ----- -50 - -# USMALLINT -> int -statement ok -COPY (SELECT 30000::USMALLINT AS a) TO '__TEST_DIR__/usmallint.dta' (FORMAT dta) - -query I -SELECT * FROM read_dta('__TEST_DIR__/usmallint.dta') ----- -30000 - -# UINTEGER -> long -statement ok -COPY (SELECT 1000000::UINTEGER AS a) TO '__TEST_DIR__/uint.dta' (FORMAT dta) - -query I -SELECT * FROM read_dta('__TEST_DIR__/uint.dta') ----- -1000000 - -# ============================================================ -# Boolean: true=1, false=0 -# ============================================================ - -statement ok -COPY (SELECT true AS t, false AS f, NULL::BOOLEAN AS n) TO '__TEST_DIR__/bool2.dta' (FORMAT dta) - -query III -SELECT * FROM read_dta('__TEST_DIR__/bool2.dta') ----- -1 0 NULL - -# ============================================================ -# FLOAT precision round-trip -# ============================================================ - -statement ok -COPY (SELECT 1.5::FLOAT AS a, (-0.001)::FLOAT AS b, 0.0::FLOAT AS c) TO '__TEST_DIR__/float_prec.dta' (FORMAT dta) - -query RRR -SELECT * FROM read_dta('__TEST_DIR__/float_prec.dta') ----- -1.5 -0.001 0.0 - -# ============================================================ -# DOUBLE precision round-trip -# ============================================================ - -statement ok -COPY (SELECT 1.23456789012345::DOUBLE AS a, (-999.999)::DOUBLE AS b) TO '__TEST_DIR__/double_prec.dta' (FORMAT dta) - -query RR -SELECT * FROM read_dta('__TEST_DIR__/double_prec.dta') ----- -1.23456789012345 -999.999 - -# ============================================================ -# DATE boundary: epoch (1970-01-01), Stata epoch (1960-01-01) -# ============================================================ - -statement ok -COPY (SELECT DATE '1960-01-01' AS stata_epoch, DATE '1970-01-01' AS unix_epoch, DATE '2000-06-15' AS y2k) TO '__TEST_DIR__/date_bounds.dta' (FORMAT dta) - -query TTT -SELECT * FROM read_dta('__TEST_DIR__/date_bounds.dta') ----- -1960-01-01 1970-01-01 2000-06-15 - -# ============================================================ -# Mixed NULLs and values in multiple rows -# ============================================================ - -statement ok -COPY (SELECT * FROM (VALUES - (1, 'a', 10.0, true), - (NULL, 'b', NULL, false), - (3, NULL, 30.0, NULL) -) AS t(id, name, val, flag)) TO '__TEST_DIR__/mixed_nulls.dta' (FORMAT dta) - -query IIRI -SELECT * FROM read_dta('__TEST_DIR__/mixed_nulls.dta') ORDER BY COALESCE(id, 999) ----- -1 a 10.0 1 -3 NULL 30.0 NULL -NULL b NULL 0 - -# ============================================================ -# Empty string vs NULL string -# ============================================================ - -statement ok -COPY (SELECT '' AS empty, NULL::VARCHAR AS missing) TO '__TEST_DIR__/empty_str.dta' (FORMAT dta) - -query II -SELECT empty IS NULL, missing IS NULL FROM read_dta('__TEST_DIR__/empty_str.dta') ----- -true true - -# ============================================================ -# Variable labels round-trip: .dta -> dodo -> save .dta -> read -# ============================================================ - -# auto.dta has variable labels set by the Python generator -statement ok -use "test/data/auto.dta", clear - -# Verify labels were read -query IIII -codebook ----- -id INTEGER Firm identifier (empty) -name VARCHAR Firm name (empty) -revenue DOUBLE Annual revenue (empty) -year SMALLINT Fiscal year (empty) - -# Save and re-read — labels should survive -statement ok -save "__TEST_DIR__/labeled.dta" - -statement ok -clear - -statement ok -use "__TEST_DIR__/labeled.dta", clear - -query IIII -codebook ----- -id INTEGER Firm identifier (empty) -name VARCHAR Firm name (empty) -revenue DOUBLE Annual revenue (empty) -year SMALLINT Fiscal year (empty) - -statement ok -clear - -# ============================================================ -# Variable labels from label variable command -# ============================================================ - -statement ok -use "test/data/auto.dta", clear - -statement ok -label variable id "Unique firm ID" - -# id label should be overridden -query IIII -codebook ----- -id INTEGER Unique firm ID (empty) -name VARCHAR Firm name (empty) -revenue DOUBLE Annual revenue (empty) -year SMALLINT Fiscal year (empty) - -# Save and verify override persists -statement ok -save "__TEST_DIR__/relabeled.dta" - -statement ok -clear - -statement ok -use "__TEST_DIR__/relabeled.dta", clear - -query IIII -codebook ----- -id INTEGER Unique firm ID (empty) -name VARCHAR Firm name (empty) -revenue DOUBLE Annual revenue (empty) -year SMALLINT Fiscal year (empty) - -statement ok -clear - -# ============================================================ -# save command with .dta extension -# ============================================================ - -statement ok -use "test/data/auto.dta", clear - -statement ok -keep if year == 2018 - -statement ok -save "__TEST_DIR__/save_test.dta" - -statement ok -clear - -query I -SELECT count(*) FROM read_dta('__TEST_DIR__/save_test.dta') ----- -2 - -query IIR -SELECT id, name, revenue FROM read_dta('__TEST_DIR__/save_test.dta') ORDER BY id ----- -2 Beta Corp 2500.5 -3 Gamma NULL