duckdb_engine.py

# duckdb_engine.py
"""
Движок для больших данных на DuckDB.

Ключевые принципы для больших файлов (7+ ГБ, 100+ млн строк):

  1. НИКОГДА не читать весь файл в RAM.
     • проверка кодировки — только первые 64 КБ;
     • нормализация хвоста — только последние 8 КБ;
     • перекодировка — потоково, chunks по 4 МБ.

  2. НЕ создавать временные копии, если файл уже в UTF-8
     и не имеет лишних пустых строк в конце.

  3. DuckDB view — ленивый, реальное чтение начинается
     только при первом SELECT. Регистрация view — мгновенная.
"""

import os
import tempfile
import duckdb


DUCKDB_THRESHOLD = 250 * 1024 * 1024  # 250 МБ
MAX_PRINT_ROWS = 20
MAX_SHOW_ROWS = 1000

# Размеры «быстрых» проверок
ENCODING_SNIFF_BYTES = 64 * 1024       # 64 КБ — для определения кодировки
TRAILING_SNIFF_BYTES = 8 * 1024        # 8 КБ  — для проверки хвоста
TRANSCODE_CHUNK_BYTES = 4 * 1024 * 1024  # 4 МБ — для потоковой перекодировки


# ============================================================
# ГЛОБАЛЬНОЕ СОЕДИНЕНИЕ
# ============================================================
# Все DuckDBTable используют ОДНО соединение, чтобы JOIN, UNPIVOT
# и другие операции могли ссылаться на таблицы друг друга по имени.
_GLOBAL_CON = None


def get_global_connection():
    """Возвращает глобальное соединение DuckDB (создаёт при первом вызове)."""
    global _GLOBAL_CON
    if _GLOBAL_CON is None:
        _GLOBAL_CON = duckdb.connect()
    return _GLOBAL_CON


def should_use_duckdb(path):
    if not os.path.exists(path):
        return False
    return os.path.getsize(path) >= DUCKDB_THRESHOLD


# ============================================================
# ОПРЕДЕЛЕНИЕ КОДИРОВКИ (только по первым N КБ)
# ============================================================
def _detect_encoding(path):
    """
    Определяет кодировку CSV по ПЕРВЫМ 64 КБ.

    НЕ читает весь файл. Для 7,9 ГБ — это один seek + read 64 КБ.

    Возвращает имя кодировки или 'utf-8' как fallback.
    """
    candidates = ['utf-8', 'cp1251', 'utf-16', 'latin-1', 'koi8-r']

    try:
        with open(path, 'rb') as f:
            head = f.read(ENCODING_SNIFF_BYTES)
    except Exception:
        return 'utf-8'

    # BOM-проверки
    if head.startswith(b'\xef\xbb\xbf'):
        return 'utf-8-sig'
    if head.startswith(b'\xff\xfe') or head.startswith(b'\xfe\xff'):
        return 'utf-16'
    if head.startswith(b'\xff\xfe\x00\x00') or head.startswith(b'\x00\x00\xfe\xff'):
        return 'utf-32'

    for enc in candidates:
        try:
            head.decode(enc)
            return enc
        except (UnicodeDecodeError, LookupError):
            continue

    return 'utf-8'


def _ensure_utf8(path):
    """
    Проверяет, что файл в UTF-8. Если нет — создаёт временную копию
    ПОТОКОВО (chunks по 4 МБ), НЕ читая весь файл в RAM.

    Возвращает (utf8_path, is_temp).

    ВАЖНО:
        • UTF-8 и UTF-8 BOM → возвращаем исходный путь, is_temp=False.
          BOM DuckDB сам обрабатывает в read_csv_auto.
        • Другие кодировки → потоковая перекодировка во временный файл.
    """
    encoding = _detect_encoding(path)

    if encoding in ('utf-8', 'utf-8-sig'):
        return path, False

    tmp = tempfile.NamedTemporaryFile(
        mode='wb',
        suffix='.utf8.csv',
        delete=False,
    )

    try:
        with open(path, 'rb') as src:
            decoder = None
            try:
                import codecs
                decoder = codecs.getincrementaldecoder(encoding)(errors='replace')
            except Exception:
                decoder = None

            if decoder is None:
                # fallback: построчное чтение
                with open(path, 'r', encoding=encoding, errors='replace') as sf:
                    for line in sf:
                        tmp.write(line.encode('utf-8'))
            else:
                while True:
                    chunk = src.read(TRANSCODE_CHUNK_BYTES)
                    if not chunk:
                        tail = decoder.decode(b'', final=True)
                        if tail:
                            tmp.write(tail.encode('utf-8'))
                        break
                    text = decoder.decode(chunk, final=False)
                    if text:
                        tmp.write(text.encode('utf-8'))
    except Exception as e:
        tmp.close()
        try:
            os.remove(tmp.name)
        except Exception:
            pass
        raise RuntimeError(f"Не удалось перекодировать файл как {encoding}: {e}")

    tmp.close()
    return tmp.name, True


# ============================================================
# ПРОВЕРКА ХВОСТА ФАЙЛА (только последние 8 КБ)
# ============================================================
def _has_trailing_blank_lines(path):
    """
    Быстро проверяет, есть ли в КОНЦЕ файла «лишние» пустые строки
    (артефакт CSV-формата).

    Читает ТОЛЬКО последние TRAILING_SNIFF_BYTES байт.
    Для 7,9 ГБ — это один seek + read 8 КБ.

    Возвращает:
        • 0 — если в хвосте нет пустых строк (файл нормальный)
        • N — сколько байт «лишнего» в конце (можно обрезать)
    """
    try:
        size = os.path.getsize(path)
    except Exception:
        return 0

    if size == 0:
        return 0

    tail_size = min(TRAILING_SNIFF_BYTES, size)

    try:
        with open(path, 'rb') as f:
            f.seek(size - tail_size)
            tail = f.read()
    except Exception:
        return 0

    # Отрезаем завершающие \r и \n
    stripped_len = len(tail.rstrip(b'\r\n'))
    trailing_ws_len = len(tail) - stripped_len

    # Если в конце 0 или 1 перевод строки — это нормально,
    # лишних пустых строк нет.
    # Считаем «лишними» только 2+ подряд идущих \n (с учётом \r\n).
    # Например: "a,b\n\n" → лишняя одна пустая строка.
    # "a,b\n"     → норма.
    if trailing_ws_len <= 2:
        # ≤ 2 байта — это либо "\n", либо "\r\n" — норма
        return 0

    # Если в хвосте одни \n и \r — значит, там только пустые строки.
    # Вернём количество лишних байт, но не больше tail_size.
    # (точное число определим при обрезке)
    return trailing_ws_len


def _trim_trailing_blank_lines(src_path, dst_file):
    """
    Копирует src_path в dst_file, отрезая ЗАВЕРШАЮЩИЕ пустые строки.

    Потоково, chunks по 4 МБ. НЕ читает весь файл в RAM.

    Возвращает (bytes_written, bytes_trimmed).
    """
    try:
        size = os.path.getsize(src_path)
    except Exception:
        size = 0

    if size == 0:
        return 0, 0

    # Найдём, сколько байт надо отрезать с конца.
    # Идём с конца блоками по 64 КБ, пока не найдём последний непустой байт.
    trim_bytes = 0
    block = 64 * 1024

    with open(src_path, 'rb') as f:
        pos = size
        while pos > 0:
            step = min(block, pos)
            f.seek(pos - step)
            chunk = f.read(step)

            # Ищем с конца первый байт, не являющийся \r или \n
            idx = len(chunk)
            while idx > 0 and chunk[idx - 1] in (0x0A, 0x0D):
                idx -= 1

            if idx == len(chunk):
                # В этом блоке нет завершающих \r\n — обрезать нечего
                break

            # Сколько байт в этом блоке — завершающие \r\n
            ws_in_block = len(chunk) - idx
            trim_bytes += ws_in_block

            if idx > 0:
                # Нашли непустой байт в этом блоке — дальше не идём
                break

            pos -= step

    # Оставляем ровно один завершающий \n (если файл не пустой)
    # Если trim_bytes > 0, значит, в конце были \r\n — оставляем один \n.
    # Если trim_bytes == 0 — ничего не трогаем.

    if trim_bytes == 0:
        # Нечего обрезать — просто копируем как есть
        written = 0
        with open(src_path, 'rb') as src, open(dst_file, 'wb') as dst:
            while True:
                chunk = src.read(TRANSCODE_CHUNK_BYTES)
                if not chunk:
                    break
                dst.write(chunk)
                written += len(chunk)
        return written, 0

    # Копируем до (size - trim_bytes), затем добавляем один \n
    keep_until = size - trim_bytes
    written = 0

    with open(src_path, 'rb') as src, open(dst_file, 'wb') as dst:
        remaining = keep_until
        while remaining > 0:
            step = min(TRANSCODE_CHUNK_BYTES, remaining)
            chunk = src.read(step)
            if not chunk:
                break
            dst.write(chunk)
            written += len(chunk)
            remaining -= len(chunk)

        # Добавим ровно один \n в конец, чтобы файл не склеивался
        # со следующим при конкатенации.
        if keep_until > 0:
            dst.write(b'\n')
            written += 1

    return written, trim_bytes


# ============================================================
# DUCKDBTABLE
# ============================================================
class DuckDBTable:
    def __init__(self, path, table_name=None, connection=None,
                 preserve_insertion_order=True, sample_size=20000):
        self.path = path
        self.table_name = table_name or self._make_table_name(path)

        # ============================================================
        # ГЛОБАЛЬНОЕ соединение — чтобы все таблицы были видны
        # друг другу (JOIN, UNPIVOT).
        # ============================================================
        self.con = connection or get_global_connection()

        self._preserve_insertion_order = preserve_insertion_order
        self._sample_size = sample_size

        # ============================================================
        # ДИАГНОСТИКА — чтобы при отладке было видно, что произошло.
        # ============================================================
        self._encoding_detected = _detect_encoding(path)
        self._was_transcoded = False
        self._trailing_trimmed = False
        self._trailing_trimmed_bytes = 0

        # ============================================================
        # ШАГ 1: кодировка.
        # Если файл уже UTF-8 — НЕ трогаем его, utf8_path = path.
        # Если другая — потоково перекодируем во временный файл.
        # ============================================================
        self.utf8_path, self.is_temp = _ensure_utf8(path)
        self._was_transcoded = self.is_temp
        self._tmp_path = self.utf8_path if self.is_temp else None

        # ============================================================
        # ШАГ 2: нормализация хвоста.
        # Быстрая проверка последних 8 КБ. Если в хвосте нет
        # «лишних» пустых строк — НИЧЕГО не делаем.
        # Если есть — потоково копируем без них во временный файл.
        # ============================================================
        self._normalize_csv_trailing_newlines()

        # ============================================================
        # ШАГ 3: регистрация view (ленивая, DuckDB не читает данные).
        # ============================================================
        self._register()

    def _make_table_name(self, path):
        base = os.path.basename(path).replace('.', '_').replace(' ', '_')
        clean = ''.join(c if c.isalnum() or c == '_' else '_' for c in base)
        return f"t_{clean}"

    def _detect_delimiter(self):
        try:
            with open(self.utf8_path, 'r', encoding='utf-8', errors='replace') as f:
                first_line = f.readline()

            candidates = [',', ';', '\t', '|']
            counts = {d: first_line.count(d) for d in candidates}
            best = max(counts, key=counts.get)

            if counts[best] == 0:
                return ','
            return best
        except Exception:
            return ','

    # ============================================================
    # НОРМАЛИЗАЦИЯ ХВОСТА CSV (только последние 8 КБ)
    # ============================================================
    def _normalize_csv_trailing_newlines(self):
        """
        Отрезает ЗАВЕРШАЮЩИЕ пустые строки, НЕ читая весь файл в RAM.

        ЗАЧЕМ:
            read_csv_auto с header=true читает завершающую пустую строку
            как строку из None по всем столбцам. COUNT(*) её не считает
            (ignore_errors=false), а SELECT * возвращает. Из-за этого
            get_row_count() и get_data_for_show() дают разное число строк,
            и joinarray(horizontal) получает лишнюю None-строку.

        КАК:
            1. Быстрая проверка: читаем последние 8 КБ.
            2. Если лишних пустых строк нет — немедленный возврат.
            3. Если есть — потоково копируем без них во временный файл.
        """
        # ---- 1. Быстрая проверка хвоста ----
        trim_len = _has_trailing_blank_lines(self.utf8_path)

        if trim_len == 0:
            # Хвост нормальный — НИЧЕГО не делаем.
            return

        # ---- 2. Медленный путь: потоковое копирование с обрезкой ----
        tmp = tempfile.NamedTemporaryFile(
            mode='wb', suffix='.norm.csv', delete=False,
        )

        try:
            written, trimmed = _trim_trailing_blank_lines(
                self.utf8_path, tmp.name,
            )
        except Exception:
            # Не смогли — оставляем как есть, read_csv_auto сам разберётся
            try:
                os.remove(tmp.name)
            except Exception:
                pass
            return

        tmp.close()

        if trimmed == 0:
            # Ничего не отрезали — временный файл не нужен
            try:
                os.remove(tmp.name)
            except Exception:
                pass
            return

        # ---- 3. Подменяем путь на новый временный ----
        old_tmp = self._tmp_path
        self.utf8_path = tmp.name
        self._tmp_path = tmp.name
        self.is_temp = True
        self._trailing_trimmed = True
        self._trailing_trimmed_bytes = trimmed

        # Удаляем старый временный файл, если был
        if old_tmp and old_tmp != tmp.name and os.path.exists(old_tmp):
            try:
                os.remove(old_tmp)
            except Exception:
                pass

    # ============================================================
    # РЕГИСТРАЦИЯ VIEW
    # ============================================================
    def _register(self):
        """
        Регистрирует view на CSV-файл.

        ВАЖНО:
            sample_size    — сколько первых строк прочитать для детекции типов.
                             По умолчанию 20000, но можно уменьшить.
            ignore_errors  — false: падать при ошибке парсинга,
                             а не молча терять строки.
            preserve_insertion_order — сохранять порядок строк.
                             Для очень больших файлов можно отключить
                             (ускоряет загрузку, но порядок не гарантирован).

        View ленивый: DuckDB не читает данные при CREATE VIEW.
        Реальное чтение начинается при первом SELECT.
        """
        safe_path = self.utf8_path.replace("'", "''").replace("\\", "\\\\")
        delimiter = self._detect_delimiter()

        if delimiter == '\t':
            delim_sql = '\\t'
        else:
            delim_sql = delimiter

        # preserve_insertion_order — сессионная настройка DuckDB.
        # Если False — ускоряет загрузку очень больших файлов.
        if not self._preserve_insertion_order:
            try:
                self.con.execute("SET preserve_insertion_order = false")
            except Exception:
                pass

        self.con.execute(f"""
            CREATE OR REPLACE VIEW {self.table_name} AS
            SELECT * FROM read_csv_auto('{safe_path}', 
                                        sample_size={int(self._sample_size)},
                                        ignore_errors=false,
                                        header=true,
                                        delim='{delim_sql}')
        """)

    # ============================================================
    # ПОДСЧЁТ СТРОК
    # ============================================================
    def _count_data(self):
        """
        Количество ДАННЫХ (без заголовка).
        Для внутренних нужд движка (LIMIT, превью и т.п.).
        """
        return self.con.execute(
            f"SELECT COUNT(*) FROM {self.table_name}"
        ).fetchone()[0]

    def get_row_count(self):
        """
        Количество строк ВКЛЮЧАЯ ЗАГОЛОВОК.
        Согласовано с Matrix: lenrow(m) = 5 для 4 строк данных.

        В Matrix data[0] — заголовок, поэтому lenrow = 5.
        В DuckDB view содержит только данные, поэтому +1.
        """
        return self._count_data() + 1

    # ============================================================
    # ДОСТУП К ДАННЫМ
    # ============================================================
    def get_columns(self):
        result = self.con.execute(
            f"SELECT * FROM {self.table_name} LIMIT 0"
        ).description
        return [col[0] for col in result]

    def query(self, sql):
        return self.con.execute(sql).fetchall()

    def query_to_matrix(self, sql):
        from runtime.matrix import MatrExMatrix
        rows = self.con.execute(sql).fetchall()
        data = [list(row) for row in rows]
        return MatrExMatrix(data, True)

    def get_data_for_show(self, limit=MAX_SHOW_ROWS):
        try:
            data = self.query(
                f"SELECT * FROM {self.table_name} LIMIT {limit}"
            )
            return [list(row) for row in data]
        except Exception as e:
            print(f"DuckDB get_data_for_show error: {e}")
            return [[""]]

    def get_headers(self):
        return self.get_columns()

    def get_column_type(self, col_name):
        try:
            safe_col = '"' + str(col_name).replace('"', '""') + '"'
            result = self.con.execute(
                f'SELECT typeof({safe_col}) FROM {self.table_name} LIMIT 1'
            ).fetchone()
            return result[0] if result else 'VARCHAR'
        except Exception:
            return 'VARCHAR'

    def close(self):
        """
        Закрывает таблицу.

        ВАЖНО:
            • Глобальное соединение НЕ закрываем — оно общее.
            • Удаляем ТОЛЬКО временные файлы, созданные нами.
            • Исходный файл пользователя НИКОГДА не удаляем.
        """
        try:
            if self.con is not None and self.con is not get_global_connection():
                self.con.close()
        except Exception:
            pass
        self.con = None

        # Удаляем только временный файл, и только если он не совпадает
        # с исходным путём пользователя.
        if (self._tmp_path
                and os.path.exists(self._tmp_path)
                and os.path.abspath(self._tmp_path) != os.path.abspath(self.path)):
            try:
                os.remove(self._tmp_path)
            except Exception:
                pass
            self._tmp_path = None

    def __str__(self):
        try:
            data_rows = self._count_data()
            columns = self.get_columns()
            cols = len(columns)

            preview_rows = min(MAX_PRINT_ROWS, data_rows)
            data = self.query(
                f"SELECT * FROM {self.table_name} LIMIT {preview_rows}"
            )

            lines = self._format_table(columns, data)

            if data_rows > MAX_PRINT_ROWS:
                lines.append(f"... (всего {data_rows} строк × {cols} столбцов)")
            else:
                lines.append(f"({data_rows} строк × {cols} столбцов)")

            return "\n".join(lines)

        except Exception as e:
            import traceback
            return f"DuckDBTable({self.table_name}, ошибка: {e})\n{traceback.format_exc()}"

    def _format_table(self, columns, rows):
        if not rows:
            return ["(пусто)"]

        widths = []
        for i, col in enumerate(columns):
            max_width = len(str(col))
            for row in rows:
                if i < len(row):
                    max_width = max(max_width, len(str(row[i])))
            widths.append(max_width)

        lines = []
        header = "  ".join(
            str(col).ljust(widths[i])
            for i, col in enumerate(columns)
        )
        lines.append(header)
        lines.append("-" * len(header))

        for row in rows:
            line = "  ".join(
                str(row[i] if i < len(row) else "").ljust(widths[i])
                for i in range(len(columns))
            )
            lines.append(line)

        return lines

    def __repr__(self):
        return self.__str__()


def escape_identifier(name):
    return '"' + str(name).replace('"', '""') + '"'


def escape_value(value):
    if value is None:
        return "NULL"
    if isinstance(value, bool):
        return "TRUE" if value else "FALSE"
    if isinstance(value, (int, float)):
        return str(value)
    return "'" + str(value).replace("'", "''") + "'"